TLDR
Can we apply reasoning models of today to problems in niche domains such as biology and chemistry?
We answer this question for drug repositioning, a problem in biology where known drugs are repurposed for uncommon side-effects.
We steer reasoning models to guess the correct drug by intervening at middle layers during the thinking stage.
Our analysis shows that steering helps in revealing geometric patterns of reasoning tokens and also guides models towards faithful traces.
Particularly, crosscoder steering stabilizes activations of reasoning models when we sample our steering vetors from a prior distribution.
Why Drug Repositioning?
Reasoning models have shown a lot of progress in domains such as mathematics and physics. Most models reason towards a single hypothesis, i.e.- the problem has a single solution. However, things often become problematic in niche domains such as biology, chemistry and pharmacy. These domains lack structured input-output pairings and problems may consist of multiple plausible solutions. This provides a unique challenge in assessing capabilities of reasoning models.
We pick the drug repositioning task which asks the question “What other existing drugs can be used to treat a disease?” Intuitively, this requires biological reasoning over multiple possible chemical combinations of molecules. Since most drugs share chemical and biological properties with their targets, drug repositioning is an effective strategy for commercial drug discovery. Our setup studies drug repositioning as a multi-answer reasoning task where reasoning models are provided with a problem definition and a side-effect (such as Anxiety Disorder or PTSD) and yield top-10 existing drugs as solutions. These solutions are then verified by an external verifier for correctness. Additionally, we also verify the reasoning trace for correctness.
To study reasoning capabilities of models, we additively apply interventions and steer models. These interventions, called steering vectors, perturb the model at the middle layers during its thinking and generation stages. But before we dive into our steering strategy and its consequences, we take a brief detour and review steering vectors in reasoning models.
What is Crosscoder Steering?
A Brief Review
Steering broadly refers to intervening upon the behaviors of a model during the generation stage. Recent works also call this representation engineering where steering is used to refer to additively applying vectors as linear perturbations to the representations of a model. Let’s formalize this definition. Suppose we have the hidden state of a layer i as hi and we want to steer the layer using a vector v with a strength, then we perturb hi in the direction of v.
While the vector v can have any form, we would like it to produce correct solutions in reasoning. One way to do this is to use v as the hidden state of another reasoning model that got the solution correct, i.e.- hcorrect. Another way to do this is to use the opposite direction of an incorrect solution, hincorrect. A third way to do this is to perturb the model along the gradient change between the correct and incorrect direction, hcorrect- hincorrect. We use this third strategy as it is more intuitive and in line with gradient updates. This is called diff-based steering.
Okay, we have now formalized steering and we have a strategy for steering reasoning models. But it still remains unclear on how to acquire hcorrect and hincorrect. This is where SAEs and crosscoders come in. We start with SAEs which are short for Sparse Auto-Encoders. SAEs are similar to VAEs, except that they enforce a sparsity constraint on their latent embeddings. That is, an SAE consists of an encoder and a decoder to reconstruct the inputs. Unlike VAEs, SAEs are single-layer encoders and decoders and they are trained to make their latents z sparse via a sparsity penalty. Suppose we want to reconstruct the hidden states of our reasoning model hi, we would apply an SAE to get the reconstruction and train it using the below loss function.
Similarly, SAEs can be trained on diverse input sources wherein hi could be the hidden states of other reasoning models. This is called a crosscoder, which encodes and decodes hidden states across models. So intuitively, a crosscoder is a very large SAE trained on hidden states of different models. This provides us a way to acquire hidden states for correct and incorrect solutions, hcorrect and hincorrect. We first harvest hidden states from different models. We then pretrain a crosscoder across the harvested hidden states. We then pass a correct hidden state from the test set and decode it to hcorrect. Similarly, we acquire hincorrect and compute their difference as the steering vector. This process is called SAE or crosscoder steering.
Crosscoders and Activations
Does crosscoder steering help? What benefits do crosscoders have? We answer these questions by steering DeepSeek R1 8B at layer 14 using a pretrained crosscoder. The side-effect is Familial Biliary Degeneration and the model is asked to output 10 drugs that can be repurposed for this disease. We compare the base model (without steering) with naive diff-based steering and a distributional variant of crosscoder steering, we call this crosscoder-guided steering (we introduce this in the next section). We see that the projection of token hidden states reveals geometric patterns. Reasoning tokens form 10 clusters where each cluster corresponds to one of the 10 solutions. These clusters also drift along the steering direction in a curved arc. While no such pattern is observed in the base model, crosscoder steering encodes geometry in its hidden states.
Could steering also help in minimizing uncertainty in reasoning? We carry along this line of thought and look at the reasoning trace trajectory. Upon monitoring KL divergences between hidden states of adjacent tokens, we observe that crosscoder-guided steering better minimizes uncertainty in earlier tokens. Base model and naive steering possess high uncertainty in earlier tokens. Crosscoder-guided steering, on the other hand, minimizes divergences early on and forces a consistent token distribution along the trace.
What else can crosscoder-guided steering do? So far we have only looked at hidden states of one layer, what about the whole model. We monitor layer-wise statistics and observe that activity grows exponentially in the base model as we progress along the layers. Later layers have significantly larger activity in their hidden states. Extrapolating towards a deeper model would indicate a blowup in model activity and instability in features. Upon applying steering we observe that activity grows sharply following the intervention and then starts saturating in the later layers. While stability gains are modest, we are still able to prevent exponential blowups using steering.
Okay, we now look at the crosscoder-guided steering variant which we utilized for our experiment.
Distributional Steering
Sampling from the Gumbel Distribution
We now describe our distributional steering strategy that we briefly introduced in the previous section. The main idea, similar to VAEs, is to sample latent embeddings of the crosscoder from an expressive distribution. Intuitively, sampling latent embeddings has two benefits. Firstly, the distribution acts as an expressive prior over hidden states by making sparse entries of embeddings into continuous values. Secondly, sampling steering vectors provides additional diversity into our steering vectors when compared to vectors sampled from a discrete space. Thus, we fit a distribution over latent embeddings of the crosscoder during inference (one can study sampling during training as well). We then sample latent embeddings from this distribution, pass it through the decoder and obtain our steering vector v = hcorrect hincorrect as before.
We are still left with one thing though. Which distribution should we sample the steering vector from? A natural choice is the normal distribution N(0, 1). We experimented with this choice and found it to perform suboptimally. Instead, we use the Gumbel distribution. This is a desirable choice for two reasons. Firstly, the Gumbel distribution has a long tail which is suitable for capturing the distribution of crosscoder embeddings. These embeddings also have a heavy-tail behavior as we only select the top-k entries from embeddings. Secondly, the distribution takes the form of a double-exponential over its inputs (see below equation) which provides more expressivity when compared to the quadratic-exponential form of N(0, 1).
Does It Work?
We would like to verify that our sampling strategy is able to influence hidden states of the reasoning model. We monitor the cosine similarity between hidden states of thinking and generated response tokens across model layers. Intuitively, a high similarity would indicate that our activations are in alignment and converge to the same solution hypothesis. A lower similarity, on the other hand, would indicate deviation from consistent behavior and hence, lower influence of the steering vector. We compare 4 reasoning models above and observe variable results. Following the addition of sampled steering vector at layer 14, most reasoning models see an increase in cosine similarities. The smaller Phi 4 Mini 4B model is not influenced by the intervention. Overall, majority of the models, especially larger models, are found to be steerable and sampled vectors are able to influence hidden states towards correct solutions.
Why Distributions Matter
Diversity and Spread
Sampling steering vectors from a Gumbel distribution provides us with a richer prior and greater diversity. Compared to a normal distribution over steering vectors, the Gumbel distribution presents more expressivity around the mean (at 57% probability). We observe a sawed-off peak that prevents overlapping of latent embeddings and boosts expressivity via greater entropy. Also, as opposed to the normal distribution, the Gumbel distribution provides a longer and heavier tail. This results in greater structural diversity. Intuitively, sampled entries balance between sparsity and diversity, allowing the steering vector to pack more information.
Reasoning Performance
Another reason for using distributional steering is the performance benefit obtained via customization. Upon lowering the scale of the Gumbel distribution to smaller values, steering vectors improve reasoning performance and trace correctness. This results in a sharper prior which is found to be more amenable for steering. Similarly, biasing the distribution mean with the Euler-Mascheroni constant (0.5772) improves solution hit rates and trace quality. We believe that additional customization can further improve reasoning capabilities.
Future Work
Several avenues exist for future work, we list three of them here.
Firstly, our setup can be extended to additional multi-answer domains such as chemistry, programming and philosophy. Such domains challenge reasoning models by assessing procedural reasoning while maintaining multi-answer formats.
Secondly, we isolate the effect of steering during reasoning. Future work may analyze reasoning models by combining steering with finetuning and few-shot in-context learning. Such a comparison will clarify the need and relative contributions of steering techniques within the broader language modeling pipeline.
Lastly, our setting considers reasoning models with moderate parameters (4B to 8B). Future work may consider assessing reasoning models with larger sizes as these would present diverse and unconventional rationales during the thinking stage.