Steering Vectors as a Training Signal in LLM Post-Training

AI summary

Overview

Research area: Machine learning / LLM post-training (supervised fine-tuning, distillation, and reinforcement learning with verifiable rewards), with roots in mechanistic interpretability.

Technical level: Advanced. The paper assumes familiarity with residual-stream activations, LoRA, knowledge distillation, and GRPO/PPO-style policy optimization.

Scope: A systematic study of whether steering vectors — directions in the residual stream previously used only to control frozen models at inference time — can instead be injected during training to shape parameter updates in both supervised and reinforcement-learning post-training regimes, plus the conditions under which this fails.

What This Paper Is About

Steering vectors are residual-stream directions that shift an LLM's behavior when added to hidden states at generation time, and prior work treats them purely as post-hoc controls on a frozen model. This paper asks whether the same directions can act as a training signal: injected into hidden states during teacher-forced supervised training (SFT and distillation), or used to repair the advantage on wrong rollouts in GRPO. The goal is to determine whether such injection reliably improves post-training, which design choices matter, and where the approach breaks down.

Key Contributions

1. Framing SV injection as a training signal. The authors define two intervention regimes — a hidden-state hook during teacher-forced supervised training and an advantage-repair mechanism on wrong GRPO rollouts — and describe each along three axes: the steering direction, its placement (layer and token positions), and the operation that combines it with existing computation. They state this is the first systematic study of SVs as a persistent training signal rather than an inference-time control.

2. Configurations that reliably help in both regimes. Across multiple students, datasets, and teachers, SV injection during supervised training improves over Plain SFT and distillation, with per-prompt directions plus spherical interpolation (Per-sample Slerp) as the strongest configuration tested. In RL, a correctness-contrast direction applied as advantage repair on wrong rollouts (SV-Rescue) recovers a learning signal on otherwise zero-gradient all-wrong groups.

3. A composable RL remedy. Because SV-Rescue only modifies the advantage estimator on wrong rollouts, it stacks with orthogonal rollout-side mechanisms such as entropy-aware advantage shaping and a temperature schedule (SV-Explore).

4. A map of three failure boundaries. Code generation breaks the direction (the reference direction conflicts with the test-passing direction), cross-family layer transfer breaks the placement, and applying the SV at inference time on an SV-trained checkpoint breaks the regime, removing rather than amplifying the gain.

Main Findings

• SV injection helps supervised training, mostly. Across four models (Qwen2.5-1.5B-Instruct, Qwen2.5-7B-Instruct, Qwen3-4B in no-thinking mode, llama3.2-3B) and four datasets (GSM8K, MATH, MBPP, StrategyQA), SV injection is positive over Plain on the three Qwen models in 11 of 12 cells; the single negative Qwen cell is Qwen2.5-1.5B on MBPP. Example: Qwen2.5-1.5B on GSM8K goes from 59.14 (Plain) to 64.14 (Per-sample Slerp).

• Per-prompt directions generally beat the global mean. Per-sample Slerp wins in 8 of 12 Qwen cells over Mean-SV, despite using a 5× smaller injection coefficient.

• Direction and geometry both matter. On Qwen2.5-1.5B / GSM8K, Per-sample Slerp reaches 64.14 (+5.00), Per-sample additive reaches 62.93 (+3.79), multi-layer Mean-SV with a tuned per-layer schedule reaches 62.62 (+3.48), PCA SV reaches 59.29 (+0.15), and a random per-sample Gaussian control drops to 54.82 (−4.32). The random control hurting the model is the evidence that the specific ground-truth-versus-prediction direction, not arbitrary perturbation, drives the gain.

• Slerp is sharper; additive is safer. On Qwen2.5-1.5B, Per-sample Slerp collapses to 0% at α = 1.0 because the hidden state is fully rotated away from the prediction direction, while Mean-SV additive has a broad plateau.

• Layer choice is model-specific, not universal. Llama-3.2-3B's default layer ℓ⋆ = 15 performs poorly, but ℓ⋆ ∈ {10, 13, 19} all surpass Plain SFT, with 10 and 19 tied at the peak. The authors conclude this is a layer-transfer failure rather than a failure of the SV mechanism; a 0.25-epoch layer sweep validated on the last 500 training examples recovers the best layer on both Qwen2.5-1.5B and Llama-3.2-3B without test data.

• Teacher mode dominates teacher family in distillation. Both thinking-mode teachers (Qwen3-8B, Gemma-4-26B thinking) give the largest gains over Plain, up to 12%, while both no-thinking teachers are essentially flat. The geometric explanation offered is that same-mode teachers from different families remain well-aligned in the SV direction, while different-mode pairs align poorly even within the same family.

• Distillation mitigation is measurable without parameter updates. Applying the training-time SV hook to 500 training samples raises label log-probability (ΔNLL_sum) and label probability mass (ΔP_sum) through teacher forcing alone, with large effects on thinking-teacher cells and weak effects on non-thinking-teacher cells.

• Gains survive multiple seeds. Reruns with seeds {42, 1, 2} show every reasoning and distillation cell improving over Plain on all three seeds — e.g. Qwen2.5-1.5B GSM8K +3.94, Qwen3-4B MATH +4.60, 7B MATH +3.73, and distillation from the Qwen3-8B thinking teacher on 1.5B MATH +15.34 — while MBPP stays neutral within noise (+0.34). The 1.5B GSM8K gain is described as about eight times the larger seed standard deviation.

• Inference-time SV on an SV-trained checkpoint hurts. On three checkpoints, adding the SV at inference does not help, even on the checkpoint trained with the SV on every answer token. The authors attribute this to the SV offloading the style and format component of the reference shift: the change training makes to the layer-ℓ⋆ hidden state is nearly orthogonal to the SV (cosine ≈ −0.05), whereas SV-free training aligns positively with it.

• RL: SV-Explore is strongest on math; SV-Rescue gives a smaller positive lift. In-domain final-step validation improves on GSM8K and DeepScaleR across all three students — e.g. Llama-3.2-3B GSM8K 81.83 → 84.49 (SV-Rescue) → 86.71 (SV-Explore), and DeepScaleR 36.69 → 38.89 → 38.92. The largest gains appear on Llama-3.2-3B, the weakest base model and therefore the one producing the most wrong rollouts.

• RL gains transfer out of distribution. MATH-500 for GSM8K-trained checkpoints: Qwen2.5-1.5B 34.60 → 35.80 → 37.00; Llama-3.2-3B 26.60 → 32.00 → 30.80. AIME-24 for DeepScaleR-trained checkpoints: Qwen3-4B 22.00 → 23.33 → 24.66. AIME accuracy is averaged over 10 independent samples per problem.

• Code is the consistent weak spot in both regimes. No method wins consistently across the three students on DeepCoder, mirroring the MBPP behavior in supervised training and supporting the claim that the direction that works for math is the wrong target for code.

Methodology in Plain English

The authors build a direction, then reuse it in two ways.

Supervised regime. Before training, they take each training prompt and (a) let the current model produce its own greedy answer, and (b) take the reference answer from the dataset or teacher. They run teacher forcing on both and average the hidden states over the answer tokens at a fixed middle layer. The difference between the reference average and the model's own average is the "off-policy gap" vector for that prompt. During SFT or distillation, they add or rotate the hidden state toward this direction at the chosen layer, on answer tokens only, so the injected signal flows into the gradient. They compare a dataset-level mean direction with additive injection (Mean-SV, α = 0.5) against a per-prompt direction with spherical interpolation that preserves hidden-state norm (Per-sample Slerp, α = 0.1). All runs use LoRA with rank 8, learning rate 5×10⁻⁵, batch size 8, and one epoch; the layer is fixed in advance at L/2+1, and the operation and α are selected once on Qwen2.5-1.5B with GSM8K (6,000 training examples) then reused unchanged everywhere.

RL regime. Before GRPO, they collect rollouts labeled correct or wrong by the verifier and compute a correctness-contrast direction from mean hidden states. During training, each rollout gets a scalar score by projecting its late-answer-token hidden states onto this direction; the score is standardized within the group, and added to the advantage only on rollouts the verifier marked wrong. Correct rollouts are untouched. On an all-wrong group, where the GRPO advantage is exactly zero, the advantage reduces to this ranking, turning a wasted prompt into a gradient. Runs use the verl GRPO implementation with 8 rollouts per prompt, KL coefficient 5×10⁻³, base temperature 0.7, learning rate 5×10⁻⁵, λ = 0.1, and a direction computed once from 500 samples.

The empirical work spans sensitivity sweeps over α and layer, ablation variants (random Gaussian, PCA, multi-layer, per-sample additive), a multi-seed reliability check, an inference-time control, distillation with four teachers of differing mode and family, and in-domain plus out-of-distribution RL evaluation.

Why This Matters

Impact on research. The paper reframes steering vectors from an inference-time knob into a training-time signal, and in doing so connects interpretability work on linear behavior directions to the practical problem of the off-policy gap in teacher forcing and zero-variance groups in GRPO. The three failure boundaries are as useful as the gains: they specify when the intervention is valid rather than presenting it as universally beneficial.

Real-world applications:

• Cheap behavioral correction during fine-tuning. Data teams can compute a direction from model-versus-reference answers and use it to get more out of the same reference data, without new labels or a larger teacher.
• Distillation from reasoning models. The largest reported effects come from thinking-mode teachers, relevant to teams distilling long chain-of-thought traces into smaller deployable students.
• Rescuing wasted RL compute. All-wrong groups consume rollout budget while producing zero gradient; the advantage repair converts them into signal, which matters for hard math and reasoning curricula where most prompts are failing.
• Composability with existing training stacks. Because the RL method only touches the advantage estimator, it can be layered onto other zero-variance remedies rather than replacing them.

Industry relevance. The intervention is lightweight — a fixed precomputed vector, a hook or a modified advantage, LoRA-scale fine-tuning, and experiments run on 2 Nvidia L40-S GPUs — which places it within reach of standard post-training pipelines rather than requiring new infrastructure.

Future Directions

• A principled way to pick the injection layer. The paper shows the mid-stack region transfers but the exact best layer does not, and offers a 0.25-epoch sweep on 500 held-out examples as a workaround; a predictive rule remains open.
• A correctness direction that works for code. Code generation breaks the direction in both the SFT and RL regimes, so finding a code-appropriate contrast direction (or a way to select among candidate directions) is an explicit open problem.
• Removing the need to choose Slerp versus additive per setting. Slerp peaks higher but collapses at α = 1.0 while additive is forgiving, which suggests a schedule or adaptive α rather than a fixed choice.
• Understanding the offloading effect. The near-orthogonal (cosine ≈ −0.05) relationship between the trained update and the SV suggests the vector absorbs style and format while the gradient learns correctness; whether that division of labor can be exploited deliberately is unresolved.

Target Audience

Researchers and practitioners working on LLM post-training — especially those running SFT, distillation, or GRPO pipelines — and interpretability researchers interested in whether linear behavior directions have causal force during optimization rather than only at inference. Readers need comfort with residual-stream activations, LoRA, and policy-gradient methods; the qualitative takeaways (where it helps, where it breaks) are accessible to a broader ML audience, but the reported configuration details assume an intermediate-to-advanced background.

Authors’ abstract

Steering vectors (SVs) are residual-stream directions that shift the behavior of LLMs when added to hidden states at inference time. While prior work has mainly used SVs as post-hoc controls, we ask whether the same directions can also guide parameter updates during post-training. We study this question across supervised learning, including SFT and distillation, and reinforcement learning with verifiable rewards (RLVR). Our framework considers two intervention regimes: a hidden-state hook during teacher-forced supervised training, and an advantage-repair mechanism on wrong rollouts in GRPO. For each regime, we analyze three design axes, which are the steering direction, its placement, and the operation used to combine it with training. In supervised settings, SV injection improves over plain SFT and distillation, with gains scaling with the off-policy gap and reaching up to a $12\%$ improvement when distilling from thinking-mode teachers. These gains hold across three random seeds. In RL settings, applying a correctness-contrast direction as advantage repair on wrong rollouts, including otherwise zero-gradient all-wrong groups, recovers a learning signal that transfers to held-out math benchmarks and acts as a composable auxiliary signal that adds to existing zero-variance remedies such as entropy-aware advantage shaping. We further identify three failure boundaries. In detail, code generation can break the direction, cross-family layer transfer can break the placement, and inference-time application on an SV-trained checkpoint can break the intervention regime. Together, these results establish steering vectors as a conditional but effective training signal for LLM post-training. Our code is available here.