Predicting Future Behaviors in Reasoning Models Enables Better Steering

Paper: 2606.11172 Authors: Evgenii Kortukov, Piotr Komorowski, Florian Klein, Paula Engl, Gabriele Sarti, Seong Joon Oh, Sebastian Lapuschkin, Wojciech Samek Categories: cs.LG

The Gap

Existing work on steering large reasoning models (LRMs) intervenes directly on hidden representations at test time. These methods assume that internal features which detect already‑generated behavior (e.g., “this sentence is toxic”) are the right intervention targets. But detection features are intrinsically retrospective—they identify behavior that has already occurred, not behavior that will happen after the next few tokens.

The authors show that detection features are poor predictors of future behavioral outcomes. Steering based on them often degrades output quality because the intervention is applied to the wrong signal. The paper’s core contribution is to separate detection features from prediction features—representations that can forecast future behavior from intermediate reasoning steps. By training linear probes to predict future behavior likelihoods (64–91% accuracy on held‑out steps), they uncover a distinct type of internal feature. They then build a text‑level steering method that uses these prediction probes to select among candidate generations, eliminating the need for direct hidden‑state manipulation and preserving output quality.

[Problem] LRM steering degrades output quality when using detection features
     |
     v
[Prior assumption] detection features are natural intervention targets
     |
     v (contradiction)
[New insight] prediction features (forecasting future behavior) exist and are separable
     |
     v
[Method] train linear probes on intermediate hidden states to predict future behavior likelihood
     |
     v
[Evidence] probes achieve 64-91% accuracy; FPCG steers with minimal quality loss
     |
     v
[Conclusion] distinguishing detection vs prediction features enables nuanced, quality‑preserving steering

The Increment

One sentence: Before this paper, steering relied on detection features that degrade output quality; after, prediction features enable text‑level steering that preserves output quality and works where activation steering fails.

Core Mechanism

The method has two phases: probe training and probe‑guided generation.

Probe training: During a forward pass of the LRM on a task, the authors extract hidden states at a specific intermediate layer (e.g., after N reasoning steps). For each step, they define a future behavior label—for example, whether the eventual answer will be correct, or whether the final response will be polite. They train linear classifiers (probes) on these hidden states to predict the future label. The probes are regularized to prevent overfitting and are validated on held‑out steps. Critically, the probes are trained *without any intervention on the model; they only read internal representations.

Probe‑guided generation (FPCG): During inference, given a partial input, the model first generates several candidate continuations (e.g., 5–10 samples). Each candidate is fed back into the model, and at the same intermediate layer, the trained prediction probe produces a score indicating the likelihood that the candidate will lead to the desired future behavior. The candidate with the highest score is selected as the final output. No hidden states are modified during generation—the probe only informs a top‑k or beam‑style selection process.

   [Input prompt] ----+----> [LRM forward pass] ----+----> activation at layer L
                      |                             |
                      v                             v
            [sample K candidates]       [probe: hidden state -> future behavior score]
                      |                             |
                      +-----> [score each candidate] -----> [select best candidate]
                                                                 |
                                                                 v
                                                           [final output]

Structural metaphor: Think of a chess grandmaster who doesn’t just look at the pieces on the board (detection features—“this is a sharp position”) but calculates several candidate moves forward and evaluates which one leads to a winning endgame (prediction features). The intermediate board states are the hidden representations; the evaluation function is the probe; each candidate continuation is a possible move; picking the highest‑scoring move is the selection mechanism. The grandmaster doesn’t touch the pieces during calculation—they only decide after forecasting.

Key Concepts

• Detection features vs Prediction features: Detection features are internal representations that correlate with properties of tokens that have *already been generated. For example, in a toxic‑content detector, the activation pattern might fire after seeing “I hate you” but not help predict whether the next sentence will continue being toxic. Prediction features, in contrast, encode information about what the model will do in the immediate future. The paper trains probes to explicitly convert activation patterns into probability estimates over future behavior. A concrete example: at reasoning step 5, a detection feature might correctly flag “I’m confused” as a sign of uncertainty, while a prediction feature could forecast that the model’s final answer will be wrong with 80% probability, even if the current tokens appear confident.

• Future Probe Controlled Generation (FPCG): A non‑intrusive steering method. Instead of modifying model weights or hidden states, FPCG samples multiple candidate completions (e.g., using different random seeds or sampling temperatures) and uses a trained prediction probe to select the most promising one. The probe is a lightweight linear classifier; training it only requires a few hundred forward passes on a validation set. FPCG is especially useful when the desired behavior is rare in the unconditional distribution, because it allows the model to generate several attempts and then pick the one that the probe rates highest.

Framework Shift

Before (mainstream approach):
                  +-----------+
  input --> | LRM with | --> output (single path)
                  | steering  |        ^
                  | intervention|        |
                  +-----------+        (hidden states modified directly)
                        |
                        v
                  degraded quality due to rough intervention

After (this paper):
                  +-------------------+
  input --> | sample K candidates | --> scores from probe --> pick best
                  +-------------------+
                        |
                        v
                  output (preserves quality, no hidden state modification)

From intervention in representational space to selection among generation possibilities based on predicted future behavior. The core shift is moving the steering locus from the hidden state level to the token level: instead of pushing the model’s internal dynamics, you let it run freely many times and then choose the run that is predicted to lead to the desired outcome.

Expert Assessment

Problem choice: Real and well‑targeted. Many practitioners have observed that activation steering often degrades fluency or coherence. The paper pinpoints a specific reason (use of detection features) and provides a clean alternative. It sits at the intersection of interpretability and alignment, where small conceptual distinctions can yield large practical improvements.

Method maturity: Clever insight + simple execution. The key novelty is conceptual (separation of detection vs prediction features), not algorithmic. The training of linear probes is standard, and the selection approach is reminiscent of rejection sampling with a learned critic. The authors do not claim algorithmic breakthroughs, and that honesty is a strength. However, one could ask: why not train a single critic to score candidates directly, without intermediate hidden states? The answer (that probes on hidden states generalize better across tasks) is reasonable but could be argued.

Experimental integrity: Baselines appear fair—they compare FPCG against activation steering and a no‑steering baseline across several tasks (correctness, helpfulness, safety). The accuracy numbers (64–91%) are reported with hold‑out sets. The paper claims FPCG causes “almost no output quality degradation” and backs it with perplexity and human evaluation. I would like to see a finer‑grained analysis of where FPCG fails (e.g., when the probe is wrong) to fully trust the numbers. No obvious red flags, but the sample sizes in human evaluation are not reported in the abstract.

Writing quality: The abstract and introduction are crisp. The main weakness is the middle sections: the explanation of how probes are trained on “future behavior labels” could be more explicit (what exactly constitutes a future behavior for different tasks?). Rewriting Section 3.1 with a concrete example (e.g., step‑by‑step for a math word problem) would make the paper much more accessible.

Verdict: weak accept — clean conceptual contribution with practical benefits, but the method itself does not introduce new machinery, and the experimental scope is modest.

Takeaways

• Diagnostic tool: When steering degrades quality, investigate whether your steering signal comes from detection or prediction features. Extract hidden states from an intermediate layer and train a linear probe to predict the desired future behavior—if the probe performs well, using it for selection (rather than intervention) may fix the quality issue.
• Selection over manipulation: For any controlled generation task (style transfer, factuality, safety), consider a simple “sample‑and‑score” loop where a lightweight probe trained on the model’s own hidden states serves as the scorer. This avoids the brittleness of activation steering and requires only a few hundred forward passes of training data.
• Transferable technique: The idea of separating detection and prediction features is not limited to reasoning models; it likely applies to any autoregressive model where you care about future properties (toxicity, bias, coherence). Practitioners can adapt the probe training pipeline to their own architecture with minimal effort.