Manifold-Guided Attention Steering

# Manifold-Guided Attention Steering
Source: [https://arxiv.org/html/2605.21770](https://arxiv.org/html/2605.21770)
Ian LiKapilesh GuruprasadRaunak Senguptai6li@ucsd\.edukguruprasad@ucsd\.edur2sengupta@ucsd\.eduNinad SatishLoris D’AntoniRose Yunsatish@ucsd\.eduldantoni@ucsd\.eduroseyu@ucsd\.edu University of California, San Diego

###### Abstract

Large language models frequently produce errors in reasoning tasks despite possessing the underlying knowledge required for correct reasoning\. One possible approach to improve reasoning consistency is through activation steering\. However, existing activation steering approaches apply fixed, pre\-computed correction vectors, ignoring where the model currently sits along its generation trajectory; the result is indiscriminate perturbation that disrupts already\-correct steps as freely as erroneous ones\. We proposeManifold\-Guided AttentionSteering \(MAGS\), a trajectory\-aware inference\-time intervention grounded in a geometric observation: the output activations of specific attention heads diverge from a low\-dimensional*correctness manifold*at the point of error, and this deviation compounds through subsequent steps\. For each identified attention head, we learn a low\-dimensional subspace from contrastive pairs of correct and incorrect traces that capture the directions along which error behavior deviates from correct behavior\. During inference, we monitor each head’s proximity to this manifold and apply a targeted projection correction when deviation exceeds a learned threshold, steering the attention output back toward the correct subspace before the error propagates\. MAGS consistently outperforms both unsteered baselines and static steering approaches across benchmarks spanning mathematical reasoning \(MATH\-500, GSM8K\), code generation \(HumanEval, MBPP\), and molecular generation \(SMILES\), suggesting that correctness manifolds are a general feature of LLM attention geometry\.

## 1Introduction

Large language models \(LLMs\) frequently produce reasoning errors in multi\-step reasoning tasks despite possessing the underlying capability to solve them\. Under repeated sampling, the same model and prompt that produce an incorrect solution will often produce a correct one\[[3](https://arxiv.org/html/2605.21770#bib.bib17),[20](https://arxiv.org/html/2605.21770#bib.bib21),[4](https://arxiv.org/html/2605.21770#bib.bib19)\], suggesting that the capability is present, but its reliable expression is not\. Process\-level annotations further confirm that most errors arise at intermediate reasoning steps rather than from a terminal absence of knowledge\[[11](https://arxiv.org/html/2605.21770#bib.bib6),[17](https://arxiv.org/html/2605.21770#bib.bib22)\]\. Since errors arise in the generation process rather than from missing knowledge, correcting them at inference time is a natural and practical target\.

Existing activation steering methods\[[16](https://arxiv.org/html/2605.21770#bib.bib1),[24](https://arxiv.org/html/2605.21770#bib.bib2),[14](https://arxiv.org/html/2605.21770#bib.bib11),[18](https://arxiv.org/html/2605.21770#bib.bib13)\]apply a fixed correction vector to the residual stream at generation steps\. These approaches are well\-suited for persistent, global behaviors \(tone, style, sentiment\) but are structurally mismatched to reasoning\. A reasoning trajectory may proceed correctly for many steps before committing a localized error at stept∗t^\{\\ast\}; applying constant corrections may corrupt the correct intermediate steps while offering no guarantee of intercepting the error\.

We hypothesize that reasoning errors manifest as a drift in a low\-dimensional subspace of individual attention heads’ output space: correct and incorrect trajectories occupy geometrically separable regions, and the transition from correct to incorrect behavior follows a structured, low\-rank direction\. We perform diagnostic experiments to confirm this hypothesis and in fact discover that correct and incorrect trajectories are highly separable by a low\-dimensional subspace of attention\-head activations\. This is consistent with mechanistic interpretability findings that individual attention heads are functionally specialized\[[5](https://arxiv.org/html/2605.21770#bib.bib4),[19](https://arxiv.org/html/2605.21770#bib.bib5)\], and with the linear representation hypothesis\[[13](https://arxiv.org/html/2605.21770#bib.bib3)\], which posits that semantically meaningful distinctions are encoded along low\-dimensional linear directions\.

Therefore, one should steer only when their attention outputs have drifted into the error subspace\.

We proposeManifold\-Guided Attention Steering \(MAGS\): an adaptive intervention that dynamically steers the attention head outputs when reasoning errors are detected\. MAGS outperforms static steering baselines across reasoning benchmarks and molecular generation on three model families, including Llama, Gemma, and GPT\-OSS\.

In summary, our contributions are as follows:

1. 1\.We hypothesize that reasoning errors manifest as structured drift in a low\-dimensional subspace of individual attention heads’ output space, and confirm this hypothesis with diagnostic experiments showing that correct and incorrect trajectories are highly separable \(Section[3](https://arxiv.org/html/2605.21770#S3)\)\.
2. 2\.We propose Manifold\-Guided Attention Steering \(MAGS\), an adaptive inference\-time mechanism that monitors attention heads for reasoning drift and applies dynamic correction only when needed \(Section[4](https://arxiv.org/html/2605.21770#S4)\)\.
3. 3\.Empirically, MAGS consistently outperforms static steering baselines on benchmarks across three model families, by up to10\.8%10\.8\\%while incurring negligible inference overhead \(Section[5](https://arxiv.org/html/2605.21770#S5)\)\.

## 2Related Work

We discuss existing inference\-time steering methods and geometric interpretability work\. Existing steering methods apply fixed corrections without error\-detection mechanisms; existing interpretability work establishes the geometric structure that we exploit and extend for adaptive intervention\. To our knowledge, MAGS is the first method to combine per\-step detection with geometry\-aware, conditional correction at the attention\-head level\.

![Refer to caption](https://arxiv.org/html/2605.21770v1/figures/method_figure.png)Figure 1:Comparison of static and Manifold\-Guided Attention Steering \(MAGS\) on an example problem\.Step\-by\-step reasoning traces for a static baseline and MAGS\. Blue boxes denote correct reasoning steps; red boxes denote erroneous ones\.### 2\.1Activation Steering and Inference\-time Intervention

Activation steering methods modify internal representations at inference time without updating model parameters\.*Activation Addition*\[[16](https://arxiv.org/html/2605.21770#bib.bib1)\]adds a fixed difference vector to the residual stream throughout generation\.*Contrastive Activation Addition*\(CAA;Rimskyet al\.[14](https://arxiv.org/html/2605.21770#bib.bib11)\) improves reliability by averaging difference vectors across many contrastive prompt pairs\.*Representation Engineering*\(RepE;Zouet al\.[24](https://arxiv.org/html/2605.21770#bib.bib2)\) extracts principal steering directions from contrastive activations via PCA\.*Angular Steering*\[[18](https://arxiv.org/html/2605.21770#bib.bib13)\]replaces additive correction with a 2D rotation applied uniformly across all layers\.*Inference\-Time Intervention*\(ITI;Liet al\.[9](https://arxiv.org/html/2605.21770#bib.bib12)\) shifts attention head outputs along a probing direction to improve truthfulness\.*CREST*\[[23](https://arxiv.org/html/2605.21770#bib.bib7)\]identifies reasoning\-relevant attention heads and applies fixed steering vectors to them, but does not adapt the correction to the model’s current trajectory state\.

All of these methods share a common limitation: they steer along a fixed direction regardless of the current activation state\. MAGS addresses this by introducing a dynamic proximity trigger that fires only when a head drifts toward the error subspace and applying a step\-dependent correction whose direction is determined by the current activation’s projection onto the error subspace rather than a fixed vector \(as illustrated in Figure[1](https://arxiv.org/html/2605.21770#S2.F1)\)\.

### 2\.2Geometric Structure of Transformer Representations

A growing body of work establishes that transformer representations have rich geometric structure that can be leveraged for analysis and intervention\.

*Mechanistic interpretability*studies decompose transformer computation into interpretable circuits\.Elhageet al\.\[[5](https://arxiv.org/html/2605.21770#bib.bib4)\]show that attention heads implement primitive operations \(copying, retrieval, inhibition\) whose outputs compose additively in the residual stream\.Wanget al\.\[[19](https://arxiv.org/html/2605.21770#bib.bib5)\]shows that multi\-step tasks are implemented by sparse circuits across a small number of heads\. Together, these results suggest that reasoning failures are likely attributable to specific heads in failure modes, which motivates our head\-level intervention\.

The*linear representation hypothesis*\[[13](https://arxiv.org/html/2605.21770#bib.bib3)\]posits that semantically meaningful distinctions are encoded along low\-dimensional linear directions in transformer representations\.Burnset al\.\[[2](https://arxiv.org/html/2605.21770#bib.bib20)\]show that truth has a linear representation findable by contrastive probing, establishing a precedent for our contrastive PCA construction\.Zouet al\.\[[24](https://arxiv.org/html/2605.21770#bib.bib2)\]confirm that high\-level concepts, including reasoning quality, are linearly decodable from residual stream activations\. MAGS extends this line of work to the per\-head level, showing that the correct\-to\-error transition in reasoning trajectories is also linearly structured within individual head\-output spaces\.

## 3Detecting Error Drift in Attention Heads

We hypothesize that incorrect reasoning traces induce drift in the output of a subset of attention heads toward a low\-dimensional error subspace, geometrically separable from the subspace occupied by correct traces\. We empirically validate this hypothesis by constructing a contrastive error manifold per head and demonstrating that a proximity\-based score achieves strong trajectory\-level error detection across layers and heads\.

### 3\.1Setup and Notation

We consider a set of reasoning problems𝒫=\{p1,…,pN\}\\mathcal\{P\}=\\\{p\_\{1\},\\ldots,p\_\{N\}\\\}\. For each problempip\_\{i\}, we generateSSindependent reasoning traces𝒯i=\{τi,1,…,τi,S\}\\mathcal\{T\}\_\{i\}=\\\{\\tau\_\{i,1\},\\ldots,\\tau\_\{i,S\}\\\}\. Each traceτ\\tauis a token sequence of lengthLτL\_\{\\tau\}, and is assigned a binary labelyτ∈\{0,1\}y\_\{\\tau\}\\in\\\{0,1\\\}\(1 = correct final answer\)\. Assume\|𝒯i\+\|≥1\|\\mathcal\{T\}\_\{i\}^\{\+\}\|\\geq 1and\|𝒯i−\|≥1\|\\mathcal\{T\}\_\{i\}^\{\-\}\|\\geq 1for allpip\_\{i\}\. We write𝒯i\+=\{τ∈𝒯i:yτ=1\}\\mathcal\{T\}\_\{i\}^\{\+\}=\\\{\\tau\\in\\mathcal\{T\}\_\{i\}:y\_\{\\tau\}=1\\\}and𝒯i−=\{τ∈𝒯i:yτ=0\}\\mathcal\{T\}\_\{i\}^\{\-\}=\\\{\\tau\\in\\mathcal\{T\}\_\{i\}:y\_\{\\tau\}=0\\\}for the correct\-trace and error\-trace sets of problempip\_\{i\}\.

Given a transformer\-based language model withLLlayers andHHattention heads per layer, where each head operates on adhd\_\{h\}\-dimensional output space\. For a sampleτ\\tauand head\(l,h\)\(l,h\), the sequence of attention head outputs is:

𝐀τ\(l,h\)=\[𝒂1\(l,h,τ\),𝒂2\(l,h,τ\),…,𝒂Lτ\(l,h,τ\)\]∈ℝLτ×dh\.\\mathbf\{A\}\_\{\\tau\}^\{\(l,h\)\}\\;=\\;\\bigl\[\\bm\{a\}\_\{1\}^\{\(l,h,\\tau\)\},\\;\\bm\{a\}\_\{2\}^\{\(l,h,\\tau\)\},\\;\\ldots,\\;\\bm\{a\}\_\{L\_\{\\tau\}\}^\{\(l,h,\\tau\)\}\\bigr\]\\;\\in\\;\\mathbb\{R\}^\{L\_\{\\tau\}\\times d\_\{h\}\}\.\(1\)

### 3\.2Contrastive Error Manifold Construction

Given correct and incorrect reasoning traces, we construct a per\-head error subspace by identifying the low\-dimensional directions along which correct and incorrect activations diverge\.

#### Per\-problem difference vectors\.

For each problem and each head\(l,h\)\(l,h\), define the per\-class mean:

𝝁c,i\(l,h\)=1∑τ∈𝒯i\+Lτ​∑τ∈𝒯i\+∑t=1Lτ𝒂t\(l,h,τ\),𝝁e,i\(l,h\)=1∑τ∈𝒯i−Lτ​∑τ∈𝒯i−∑t=1Lτ𝒂t\(l,h,τ\)\.\\bm\{\\mu\}\_\{c,i\}^\{\(l,h\)\}=\\frac\{1\}\{\\sum\_\{\\tau\\in\\mathcal\{T\}\_\{i\}^\{\+\}\}L\_\{\\tau\}\}\\sum\_\{\\tau\\in\\mathcal\{T\}\_\{i\}^\{\+\}\}\\sum\_\{t=1\}^\{L\_\{\\tau\}\}\\bm\{a\}\_\{t\}^\{\(l,h,\\tau\)\},\\qquad\\bm\{\\mu\}\_\{e,i\}^\{\(l,h\)\}=\\frac\{1\}\{\\sum\_\{\\tau\\in\\mathcal\{T\}\_\{i\}^\{\-\}\}L\_\{\\tau\}\}\\sum\_\{\\tau\\in\\mathcal\{T\}\_\{i\}^\{\-\}\}\\sum\_\{t=1\}^\{L\_\{\\tau\}\}\\bm\{a\}\_\{t\}^\{\(l,h,\\tau\)\}\.\(2\)The*contrastive difference vector*for problempip\_\{i\}and head\(l,h\)\(l,h\)is:

𝜹i\(l,h\)=𝝁e,i\(l,h\)−𝝁c,i\(l,h\)∈ℝdh\.\\bm\{\\delta\}\_\{i\}^\{\(l,h\)\}\\;=\\;\\bm\{\\mu\}\_\{e,i\}^\{\(l,h\)\}\\;\-\\;\\bm\{\\mu\}\_\{c,i\}^\{\(l,h\)\}\\;\\in\\;\\mathbb\{R\}^\{d\_\{h\}\}\.\(3\)By construction,𝜹i\(l,h\)\\bm\{\\delta\}\_\{i\}^\{\(l,h\)\}cancels all directions uniformly activated by problempip\_\{i\}regardless of correctness, isolating the*directional shift*attributable to the error\.

#### Difference matrix\.

![Refer to caption](https://arxiv.org/html/2605.21770v1/figures/manifold_schema.png)Figure 2:Schematic of the contrastive error manifold\. Correct and error activation spaces are separated along the learned subspacespan​\(𝐁\)\\mathrm\{span\}\(\\mathbf\{B\}\)\. Given an error activationata\_\{t\}, the projection𝐁⊤​𝐁\\mathbf\{B\}^\{\\top\}\\\!\\mathbf\{B\}gives the direction to map back towardμc\(l,h\)\\mu\_\{c\}^\{\(l,h\)\}\.For a set ofNNproblems, we stack the difference vectors row\-wise:

\(𝐃\(l,h\)\)T=\[𝜹1\(l,h\)…𝜹P\(l,h\)\]∈ℝdh×N\.\(\\mathbf\{D\}^\{\(l,h\)\}\)^\{T\}\\;=\\;\\begin\{bmatrix\}\\bm\{\\delta\}\_\{1\}^\{\(l,h\)\}&\\dots&\\bm\{\\delta\}\_\{P\}^\{\(l,h\)\}\\end\{bmatrix\}\\;\\in\\;\\mathbb\{R\}^\{d\_\{h\}\\times N\}\.\(4\)

#### Error subspace via PCA\.

Compute the compact singular value decomposition𝑫\(l,h\)=𝑼​𝚺​𝑽⊤\\bm\{D\}^\{\(l,h\)\}=\\bm\{U\}\\bm\{\\Sigma\}\\bm\{V\}^\{\\top\}\. Define the*error subspace basis*as the top\-kkright singular vectors:

𝑩\(l,h\)=𝑽:,1:k⊤∈ℝk×dh,\\bm\{B\}^\{\(l,h\)\}\\;=\\;\\bm\{V\}\_\{:,1:k\}^\{\\top\}\\;\\in\\;\\mathbb\{R\}^\{k\\times d\_\{h\}\},\(5\)where the rows of𝑩\(l,h\)\\bm\{B\}^\{\(l,h\)\}are orthonormal\. The error subspace captures thekkdirections in head\-output space along which correct\-to\-incorrect deviation has the greatest variance across problems, as illustrated in Figure[2](https://arxiv.org/html/2605.21770#S3.F2)\.

#### Global correct\-state centroid\.

Compute a global reference point from all correct traces:

𝝁c\(l,h\)=∑i∑τ∈𝒯i\+∑t=1Lτ𝒂t\(l,h,τ\)∑i∑τ∈𝒯i\+Lτ\.\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\;=\\;\\frac\{\\displaystyle\\sum\_\{i\}\\sum\_\{\\tau\\in\\mathcal\{T\}\_\{i\}^\{\+\}\}\\sum\_\{t=1\}^\{L\_\{\\tau\}\}\\bm\{a\}\_\{t\}^\{\(l,h,\\tau\)\}\}\{\\displaystyle\\sum\_\{i\}\\sum\_\{\\tau\\in\\mathcal\{T\}\_\{i\}^\{\+\}\}L\_\{\\tau\}\}\.\(6\)
This serves as the centering reference at inference time, since per\-problem means are not available during generation\.

### 3\.3Proximity\-Based Error Detection

To detect when a head has drifted into the error subspace, we measure how much its current output projects onto the learned error subspace\. A large projection indicates the head is behaving similarly to how it behaves in erroneous traces\. At each decode stepttduring inference, for each monitored head\(l,h\)\(l,h\), we compute the*proximity score*:

dt\(l,h\)=‖𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\)‖2=\(𝒂t\(l,h\)−𝝁c\(l,h\)\)⊤​𝑩\(l,h\)⊤​𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\)\.d\_\{t\}^\{\(l,h\)\}\\;=\\;\\left\\\|\\bm\{B\}^\{\(l,h\)\}\\bigl\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\bigr\)\\right\\\|^\{2\}\\;=\\;\\bigl\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\bigr\)^\{\\top\}\\bm\{B\}^\{\(l,h\)\\top\}\\bm\{B\}^\{\(l,h\)\}\\bigl\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\bigr\)\.\(7\)This is the squared norm of the projection of the centered head output onto the error subspace\. A large value indicates that the current head output has a substantial component along the learned error directions\.

To avoid flagging normal generation steps, we calibrate a per\-head threshold on correct traces and trigger only when the proximity score exceeds it\. We fire a correction only when the proximity score exceeds a per\-head threshold calibrated on correct traces, ensuring that already\-correct steps are left undisturbed\. A*trigger*fires at stepttfor head\(l,h\)\(l,h\)when:

dt\(l,h\)\>τ\(l,h\),d\_\{t\}^\{\(l,h\)\}\\;\>\\;\\tau^\{\(l,h\)\},\(8\)whereτ\(l,h\)\\tau^\{\(l,h\)\}is set to theqq\-th percentile of\{dt\(l,h\)\}\\\{d\_\{t\}^\{\(l,h\)\}\\\}computed over all token steps from correct trajectories in the training set\.

![Refer to caption](https://arxiv.org/html/2605.21770v1/figures/v1_auroc_heatmap.png)Figure 3:Per\-head error detection AUROC across four monitored layers for the Math\-Instruct dataset with Llama\-3\.1\-8b\-Instruct\. Dots mark heads with AUROC\>0\.65\>0\.65\. The signal is sparse and concentrated in specific heads\.
### 3\.4Empirical Validation of the Drift Hypothesis

We validate whether the learned error subspace carries a detectable signal for distinguishing correct from incorrect reasoning trajectories\. Using Math\-Instruct traces from Llama\-3\.1\-8B\-Instruct collected at layers\{8,16,24,31\}\\\{8,16,24,31\\\}, we perform a problem\-level 70/30 train/test split, ensuring that all traces of a given problem land in the same split and the manifold never observes test problems during construction\. We build the contrastive error manifold on the training split and evaluate on the held\-out test split\.

For each trace, we aggregate the per\-step proximity scores\{dt\(l,h\)\}\\\{d\_\{t\}^\{\(l,h\)\}\\\}into a single scalar using max aggregation, and classify the trace as incorrect if the score exceeds a thresholdτ\(l,h\)\\tau^\{\(l,h\)\}calibrated on the training set by maximizing balanced accuracy\. We report trajectory\-level AUROC for each\(l,h\)\(l,h\)pair independently\. Figure[3](https://arxiv.org/html/2605.21770#S3.F3)shows the per\-head AUROC across all monitored layers\. The signal is sparse and concentrated in specific heads \(marked dots for AUROC\>0\.65\>0\.65\), confirming that error drift is a structured, localized phenomenon rather than a diffuse property of all heads\.

## 4Manifold\-Guided Attention Steering

Having established that proximity scores reliably detect drift toward the error subspace, we now describe how MAGS exploits this signal to apply a targeted correction to the attention head output at inference time\.

#### Head selection\.

Rather than monitoring allL×HL\\times Hheads, we pre\-select the top\-KKheads by*held\-out AUROC*: for each head, AUROC is computed between the trajectory\-level error label and the mean proximity score over the trajectory, on a held\-out problem split\. Monitoring only the top\-KKheads reduces per\-step overhead fromO​\(L⋅H⋅k⋅dh\)O\(L\\cdot H\\cdot k\\cdot d\_\{h\}\)toO​\(K⋅k⋅dh\)O\(K\\cdot k\\cdot d\_\{h\}\)\.

#### Steering by error\-component correction\.

When head\(l,h\)\(l,h\)triggers at steptt, we apply the following in\-place correction to its output*before*it is passed to the output projection of layerll,WO\(l\)W\_\{O\}^\{\(l\)\}:

𝒂~t\(l,h\)=𝒂t\(l,h\)−α​𝑩\(l,h\)⊤​𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\),\\tilde\{\\bm\{a\}\}\_\{t\}^\{\(l,h\)\}\\;=\\;\\bm\{a\}\_\{t\}^\{\(l,h\)\}\\;\-\\;\\alpha\\,\\bm\{B\}^\{\(l,h\)\\top\}\\bm\{B\}^\{\(l,h\)\}\\bigl\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\bigr\),\(9\)whereα∈\(0,1\]\\alpha\\in\(0,1\]is a steering strength hyperparameter that we can control empirically\.

Supposeα=1\\alpha=1, and let𝑷⟂\(l,h\)=𝑰dh−𝑩\(l,h\)⊤​𝑩\(l,h\)\\bm\{P\}\_\{\\perp\}^\{\(l,h\)\}=\\bm\{I\}\_\{d\_\{h\}\}\-\\bm\{B\}^\{\(l,h\)\\top\}\\bm\{B\}^\{\(l,h\)\}denote the orthogonal projector onto the*complement*of the error subspace\. Then \([9](https://arxiv.org/html/2605.21770#S4.E9)\) can be written equivalently as:

𝒂~t\(l,h\)=𝝁c\(l,h\)\+𝑷⟂\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\)\.\\tilde\{\\bm\{a\}\}\_\{t\}^\{\(l,h\)\}\\;=\\;\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\;\+\\;\\bm\{P\}\_\{\\perp\}^\{\(l,h\)\}\\bigl\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\\bigr\)\.\(10\)This form makes the semantics transparent: we decompose the deviation of the head output from the correct\-state mean into an error\-subspace component and a complement component, then discard only the former\. Algorithm[1](https://arxiv.org/html/2605.21770#alg1)summarizes the complete inference procedure, combining the proximity check \(Section[3\.3](https://arxiv.org/html/2605.21770#S3.SS3)\) and the error\-component correction into a single decode loop\.

We formalize the key advantage over full residual stream correction: thedh−kd\_\{h\}\-kdirections unrelated to the error manifold are completely untouched\.

###### Proposition 1\(Information Preservation\)\.

The correction \([9](https://arxiv.org/html/2605.21770#S4.E9)\) preserves all information in𝐚t\(l,h\)\\bm\{a\}\_\{t\}^\{\(l,h\)\}that lies in the\(dh−k\)\(d\_\{h\}\-k\)\-dimensional complement of the error subspace\. Specifically, for any vector𝐯∈ℝdh\\bm\{v\}\\in\\mathbb\{R\}^\{d\_\{h\}\}with𝐁\(l,h\)​𝐯=𝟎\\bm\{B\}^\{\(l,h\)\}\\bm\{v\}=\\bm\{0\}:

⟨𝒂~t\(l,h\),𝒗⟩=⟨𝒂t\(l,h\),𝒗⟩\.\\left\\langle\\tilde\{\\bm\{a\}\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle\\;=\\;\\left\\langle\\bm\{a\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle\.\(11\)

Algorithm 1Manifold\-Guided Head Steering \(MAGS\) — Inference1:Pre\-computed manifolds

\{𝑩\(l,h\),𝝁c\(l,h\),τ\(l,h\)\}\\\{\\bm\{B\}^\{\(l,h\)\},\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\},\\tau^\{\(l,h\)\}\\\}for top\-

KKheads; prompt

x1:promptx\_\{1:\\mathrm\{prompt\}\}
2:

t←0t\\leftarrow 0
3:whilegeneration not completedo

4:

t←t\+1t\\leftarrow t\+1
5:Run forward pass for token

tt; collect

𝒂t\(l,h\)\\bm\{a\}\_\{t\}^\{\(l,h\)\}for all monitored

\(l,h\)\(l,h\)via pre\-hooks on

WO\(l\)W\_\{O\}^\{\(l\)\}
6:foreach monitored head

\(l,h\)\(l,h\)in layer orderdo

7:Compute

dt\(l,h\)←‖𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\)‖2d\_\{t\}^\{\(l,h\)\}\\leftarrow\\left\\\|\\bm\{B\}^\{\(l,h\)\}\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\)\\right\\\|^\{2\}// Eq\. \([7](https://arxiv.org/html/2605.21770#S3.E7)\)

8:if

dt\(l,h\)\>τ\(l,h\)d\_\{t\}^\{\(l,h\)\}\>\\tau^\{\(l,h\)\}then

9:

𝒂~t\(l,h\)←𝒂t\(l,h\)−α​𝑩\(l,h\)⊤​𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\)\\tilde\{\\bm\{a\}\}\_\{t\}^\{\(l,h\)\}\\leftarrow\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\alpha\\bm\{B\}^\{\(l,h\)\\top\}\\bm\{B\}^\{\(l,h\)\}\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\)// Eq\. \([9](https://arxiv.org/html/2605.21770#S4.E9)\)

10:Replace

𝒂t\(l,h\)\\bm\{a\}\_\{t\}^\{\(l,h\)\}with

𝒂~t\(l,h\)\\tilde\{\\bm\{a\}\}\_\{t\}^\{\(l,h\)\}in the input to

WO\(l\)W\_\{O\}^\{\(l\)\}
11:endif

12:endfor

13:Complete the forward pass; sample next token

xt\+1x\_\{t\+1\}
14:endwhile

The overhead of MAGS per decode step is dominated by theKKproximity score computations\. Each requires a matrix\-vector product𝑩\(l,h\)​𝒗\\bm\{B\}^\{\(l,h\)\}\\bm\{v\}of costO​\(k⋅dh\)O\(k\\cdot d\_\{h\}\), followed by a norm computation of costO​\(k\)O\(k\)\. The conditional correction, when triggered, requires one additional matrix\-vector product𝑩\(l,h\)⊤​\(𝑩\(l,h\)​𝒗\)\\bm\{B\}^\{\(l,h\)\\top\}\(\\bm\{B\}^\{\(l,h\)\}\\bm\{v\}\)of costO​\(k⋅dh\)O\(k\\cdot d\_\{h\}\)\. The total per\-step overhead isO​\(K⋅k⋅dh\)O\(K\\cdot k\\cdot d\_\{h\}\)\.

#### Compositional Steering for Multiple Objectives\.

A fundamental challenge in multi\-objective steering is that naively combining steering vectors for different constraints can produce conflicting corrections\. We sidestep this problem by introducing MAGSu\{\}^\{\\text\{u\}\}, which treats each objective independently from the start\. For each objectiveckc\_\{k\}, MAGSu\{\}^\{\\text\{u\}\}extracts a dedicated error manifold, which naturally yields disjoint head setsℋ1,…,ℋK\\mathcal\{H\}\_\{1\},\\dots,\\mathcal\{H\}\_\{K\}during head selection\. Then, at inference time, MAGSu\{\}^\{\\text\{u\}\}steers the*union*of all selected heads, applying each head’s correction through its own independently learned manifold\.

## 5Experiments

### 5\.1Reasoning Benchmarks

#### Setup and metrics\.

We evaluate across two reasoning domains chosen to span difficulty and output structure: mathematical reasoning onMATH\-500\[[11](https://arxiv.org/html/2605.21770#bib.bib6)\]andGSM8K\[[4](https://arxiv.org/html/2605.21770#bib.bib19)\], and code generation onHumanEval\[[3](https://arxiv.org/html/2605.21770#bib.bib17)\]andMBPP\[[1](https://arxiv.org/html/2605.21770#bib.bib16)\], where correctness is verified by executing test cases rather than by string matching\. To assess generality across model families, we run all reasoning and code experiments on two instruction\-tuned language models:Llama\-3\.1\-8B\-Instruct\[[7](https://arxiv.org/html/2605.21770#bib.bib28)\]andGemma\-4\-E4b\-it\[[6](https://arxiv.org/html/2605.21770#bib.bib29)\]\.

#### Baselines\.

We compare against two families of inference\-time steering methods\.

*Representation\-level methods*\.Inference\-Time Intervention\(ITI;Liet al\.[9](https://arxiv.org/html/2605.21770#bib.bib12)\), which steers individual attention head outputs along a fixed linear probe direction;Angular Steering\(AS;Vu and Nguyen[18](https://arxiv.org/html/2605.21770#bib.bib13)\), which applies a fixed 2D rotation in the mean\-difference span across all layers\. Since ITI and AS were originally designed for behavioral alignment rather than reasoning, we adapt them to our setting by treating correct and incorrect solution traces as the desired and undesired contrast sets, respectively, replacing their original prompt\-pair construction\.

*Decoding\-level methods*\.Contrastive Decoding\(CD;Liet al\.[10](https://arxiv.org/html/2605.21770#bib.bib23)\) contrasts the token distributions of a large expert model and a smaller amateur model at each decoding step, providing a baseline that operates at the output distribution level and a more computationally expensive method than all steering methods\.

#### Manifold construction\.

For each benchmark, we collect contrastive trace pairs \(i\.e\., same problem, one correct solution and one incorrect\) from the corresponding training split:Math\-Instruct\[[22](https://arxiv.org/html/2605.21770#bib.bib25)\]for MATH\-500, theGSM8K\[[4](https://arxiv.org/html/2605.21770#bib.bib19)\]training set for GSM8K, andAPPS\[[8](https://arxiv.org/html/2605.21770#bib.bib26)\]for both HumanEval\[[3](https://arxiv.org/html/2605.21770#bib.bib17)\]and MBPP\[[1](https://arxiv.org/html/2605.21770#bib.bib16)\]\. Traces are generated by sampling the base model; problems for which both a correct and an incorrect trace cannot be obtained within 8 samples are discarded\.

Table 1:Performance and output fluency \(perplexity,↓\\downarrow\) of steering methods on Llama\-3\.1\-8B\-Instruct\. Best accuracy per benchmark inbold; best perplexity initalics\.Table 2:Performance and output fluency \(perplexity,↓\\downarrow\) of steering methods on Gemma\-4\-E4b\-it\. Best accuracy per benchmark inbold; best perplexity initalics
#### Result\.

As shown in Table[1](https://arxiv.org/html/2605.21770#S5.T1)and Table[2](https://arxiv.org/html/2605.21770#S5.T2), MAGS consistently outperforms the unsteered baseline and all three steering methods across both models and all four benchmarks\. On MATH\-500, where multi\-step derivations provide the most opportunities for error compounding, MAGS achieves the largest gains: \+5\.2 points over the unsteered Llama baseline and \+3\.4 points over the unsteered Gemma baseline\. On GSM8K, where problems are shorter and errors more localized, the margin over baselines narrows, consistent with the intuition that proximity\-triggered correction is most valuable when correct steps substantially outnumber erroneous ones\.

Code generation reveals a sharper contrast\. MAGS improves HumanEval pass@1 by approximately 4 points over the unsteered baseline on both models\. All three baselines*degrade*HumanEval performance relative to the unsteered Gemma model, suggesting that unconditional interventions disrupt syntactic coherence even when the model’s reasoning is already correct\. MAGS avoids this failure mode since its proximity trigger suppresses corrections at steps where no drift is detected\.

Beyond accuracy, we evaluate output fluency via perplexity to assess whether steering distorts the model’s generation distribution\. MAGS consistently matches or approaches unsteered perplexity across benchmarks, confirming that the proximity threshold suppresses corrections on already\-correct steps and leaves the output distribution largely intact\. ITI and Angular Steering generally raise perplexity relative to the baseline, indicating that static interventions can potentially distort the generation distribution\.

### 5\.2Molecular Generation

#### Task\.

Molecular generation requires producing syntactically and chemically valid SMILES strings\[[21](https://arxiv.org/html/2605.21770#bib.bib18)\]\. Unlike natural language, SMILES has rigid grammar rules: a single misplaced token renders the entire molecule invalid\. Beyond validity, we additionally steer toward improved binding affinity against a target protein, measured via the AutoDock\-GPU docking score\[[15](https://arxiv.org/html/2605.21770#bib.bib31)\]\. Since an ideal molecule would be both syntactically valid and chemically strong binding, this task induces a natural multi\-objective structure\.

#### Setup and Metrics\.

For molecular generation, we evaluate usingGPT\-OSS 20B\[[12](https://arxiv.org/html/2605.21770#bib.bib30)\]\. We reportValidity\(fraction of generated SMILES parseable; higher is better\) andBinding Affinity\(measured as AutoDock\-GPU scores; lower is better\)\. We generate 500 molecules per method\. Contrastive Decoding is excluded as it requires a smaller companion model; no publicly available version of GPT\-OSS below 20B parameters exists at the time of writing\.

Table 3:Molecular generation by steering GPT\-OSS\-20B\. Validity: higher is better \(%\)\. Binding Affinity: lower score is better\(kcal/mol\)\.
#### Result\.

As shown in Table[3](https://arxiv.org/html/2605.21770#S5.T3), ITI illustrates the consequence of the entanglement of objectives: it substantially improves validity but incurs a significant drop in binding affinity, a sign of overcorrection toward common valid scaffolds that are grammatically safe but chemically generic\. MAGSu\{\}^\{\\text\{u\}\}avoids this interference by learning two independent manifolds: the affinity manifold \(contrasting high and low binding affinity\) and the validity manifold \(contrasting valid and invalid molecules\)\. At inference time, the union of both subspaces applies each correction in its own independent direction, allowing binding affinity to improve without the validity–affinity interference observed in joint methods\.

## 6Discussion

#### Ablation on hyperparameter sensitivity\.

While we detail the experiment hyperparameters in Appendix[B](https://arxiv.org/html/2605.21770#A2), Table[4](https://arxiv.org/html/2605.21770#A2.T4)reveals a clear contrast in robustness across methods\. Steering Llama\-3\.1\-8B on MATH\-500, MAGS accuracy varies narrowly \(0\.492–0\.530\) across all configurations tested, indicating that the proximity threshold already governs when corrections are applied and reduces sensitivity to other hyperparameters\. ITI spans a much wider range \(0\.214–0\.498\), with performance collapsing at high steering strength, confirming that static interventions are brittle to miscalibration\. Angular Steering shows the greatest variance \(0\.206–0\.506\), with accuracy swinging dramatically across rotation angles\.

![Refer to caption](https://arxiv.org/html/2605.21770v1/figures/latent_trajectory.png)Figure 4:Latent\-space trajectories of attention\-head activations projected onto the top\-4 principal components of the contrastive error subspace\.Filled contours: kernel\-density estimate of the correct\-output activation distribution\.Purple: unsteered mean trajectory \(light→\\todark = early→\\tolate generation steps\)\.Orange: MAGS\-steered mean trajectory\.Teal: ITI\-steered mean trajectory\. Faint lines show individual problem traces\. Dashed connectors mark a correction step for MAGS and ITI\.
#### Latent trajectory analysis\.

To visualize how MAGS affects the model’s internal representations during generation, we project the attention\-head output activations at a steered layer onto the top\-4 principal components of the contrastive error subspace\. We compare MAGS against ITI as ITI also operates directly on attention\-head representations\. We run parallel steered and unsteered decodes on the same set of problems, collect per\-step activations, and plot the resulting trajectories in the PC1\-PC2 and PC3\-PC4 planes\.

Figure[4](https://arxiv.org/html/2605.21770#S6.F4)reveals a clear divergence between methods\. The unsteered trajectory \(purple\) drifts steadily away from the correct\-output distribution and fails to recover, consistent with the hypothesis that reasoning errors manifest as directional activation drift\. MAGS \(orange\) intervenes at the correction step \(orange dashed connector\) and immediately redirects the trajectory back toward the high\-density correct region, where it remains for the rest of the generation\. ITI \(teal\), even after intervention, continues to diverge and settles in a region far from the correct\-output distribution, which geometrically explains why static interventions can degrade performance: without a subspace constraint, the correction vector pushes activations in an imprecise direction that does not align with the correct\-output manifold\.

## 7Conclusion

We presented Manifold\-Guided Attention Steering \(MAGS\), an inference\-time method that corrects attention\-head activations by projecting them back onto a low\-rank manifold learned from correct\-output traces, intervening only when a proximity thresholdτ\\tausignals that the current activation has drifted clearly off\-manifold and leaving already\-correct steps undisturbed\. Across mathematical reasoning, code generation, and molecular generation on three model families, MAGS consistently outperforms ITI, Angular Steering, and Contrastive Decoding\. The manifold perspective further reframes steering as a*drift detection*problem: the low\-rank basisBBidentifies the subspace most predictive of error onset, providing interpretable geometric structure that scalar\-magnitude methods lack, as visualized through latent\-space projections of steered trajectories being pulled back into the high\-density region of correct\-output activations\.

#### Limitations and future work\.

MAGS requires a set of contrastive\-output traces for manifold fitting, which may be difficult to curate in low\-resource domains\. The current method also operates on a fixed set of target heads identified offline; an adaptive scheme that selects heads dynamically based on generation context could reduce reliance on this pre\-selection step\. More broadly, we treat each head independently; jointly modeling interactions between heads may yield stronger corrections at lower computational cost\. Finally, while our experiments cover three domains, the degree to which the learned manifold transfers across tasks within a domain \(e\.g\., arithmetic to symbolic reasoning\) remains an open question\. We leave these directions to future work\.

## Acknowledgements

This work was supported, in part, by the National Science Foundation under grants CCF\-2546822, CCF\-2506134, CCF\-2446711, CCF\-2422214, \#2205093, \#2146343, \#2134274, \#2441832; the Schmidt Foundation; and a Microsoft Faculty Fellowship\. This work was also supported in part by the U\.S\. Army Research Office under Army\-ECASE award W911NF\-07\-R\-0003\-03; the U\.S\. Department of Energy, Office of Science; the ARPA\-H SOL\-24\-101 program; the IARPA HAYSTAC Program; DARPA YFA; and CDC\-RFA\-FT\-23\-0069\.

## References

- \[1\]J\. Austin, A\. Odena, M\. Nye, M\. Bosma, H\. Michalewski, D\. Dohan, E\. Jiang, C\. Cai, M\. Terry, Q\. Le, and C\. Sutton\(2021\)Program synthesis with large language models\.arXiv preprint arXiv:2108\.07732\.Cited by:[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px3.p1.1)\.
- \[2\]\(2022\)Discovering latent knowledge in language models without supervision\.arXiv preprint arXiv:2212\.03827\.Cited by:[§2\.2](https://arxiv.org/html/2605.21770#S2.SS2.p3.1)\.
- \[3\]M\. Chen, J\. Tworek, H\. Jun, Q\. Yuan, H\. Pondé, J\. Kaplan, H\. Edwards, Y\. Burda, N\. Joseph, G\. Brockman, A\. Ray, R\. Puri, G\. Krueger, M\. Petrov, H\. Khlaaf, G\. Sastry, P\. Mishkin, B\. Chan, S\. Gray, N\. Ryder, M\. Pavlov, A\. Power, L\. Kaiser, M\. Bavarian, C\. Winter, P\. Tillet, F\. P\. Such, D\. W\. Cummings, M\. Plappert, F\. Chantzis, E\. Barnes, A\. Herbert\-Voss, W\. H\. Guss, A\. Nichol, I\. Babuschkin, S\. Balaji, S\. Jain, A\. Carr, J\. Leike, J\. Achiam, V\. Misra, E\. Morikawa, A\. Radford, M\. M\. Knight, M\. Brundage, M\. Murati, K\. Mayer, P\. Welinder, B\. McGrew, D\. Amodei, S\. McCandlish, I\. Sutskever, and W\. Zaremba\(2021\)Evaluating large language models trained on code\.arXiv preprint arXiv: 2107:03374\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px3.p1.1)\.
- \[4\]K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano, C\. Hesse, and J\. Schulman\(2021\)Training verifiers to solve math word problems\.arXiv preprint arXiv:2110\.14168\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px3.p1.1)\.
- \[5\]N\. Elhage, N\. Nanda, C\. Olsson, T\. Henighan, N\. Joseph, B\. Mann, A\. Askell, Y\. Bai, A\. Chen, T\. Conerly, N\. DasSarma, D\. Drain, D\. Ganguli, Z\. Hatfield\-Dodds, D\. Hernandez, A\. Jones, J\. Kernion, L\. Lovitt, K\. Ndousse, D\. Amodei, T\. Brown, J\. Clark, J\. Kaplan, S\. McCandlish, and C\. Olah\(2021\)A mathematical framework for transformer circuits\.Transformer Circuits Thread\.Note:https://transformer\-circuits\.pub/2021/framework/index\.htmlCited by:[§1](https://arxiv.org/html/2605.21770#S1.p3.1),[§2\.2](https://arxiv.org/html/2605.21770#S2.SS2.p2.1)\.
- \[6\]Gemma Team, Google DeepMind\(2025\)Gemma 4 technical report\.Note:[https://ai\.google\.dev/gemma/docs/core/model\_card\_4](https://ai.google.dev/gemma/docs/core/model_card_4)Cited by:[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px1.p1.1)\.
- \[7\]A\. Grattafioriet al\.\(2024\)The llama 3 herd of models\.arXiv preprint arXiv:2407\.21783\.Cited by:[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px1.p1.1)\.
- \[8\]D\. Hendrycks, S\. Basart, S\. Kadavath, M\. Mazeika, A\. Arora, E\. Guo, C\. Burns, S\. Puranik, H\. He, D\. Song, and J\. Steinhardt\(2021\)Measuring coding challenge competence with apps\.InThe Thirty\-fifth Annual Conference on Neural Information Processing Systems,Cited by:[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px3.p1.1)\.
- \[9\]K\. Li, O\. Patel, F\. Viégas, H\. Pfister, and M\. Wattenberg\(2023\)Inference\-time intervention: eliciting truthful answers from a language model\.InThirty\-seventh Conference on Neural Information Processing Systems,Cited by:[§B\.1](https://arxiv.org/html/2605.21770#A2.SS1.p1.4),[§2\.1](https://arxiv.org/html/2605.21770#S2.SS1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px2.p2.1)\.
- \[10\]X\. L\. Li, A\. Holtzman, D\. Fried, P\. Liang, J\. Eisner, T\. Hashimoto, L\. Zettlemoyer, and M\. Lewis\(2023\-07\)Contrastive decoding: open\-ended text generation as optimization\.InProceedings of the 61st Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),A\. Rogers, J\. Boyd\-Graber, and N\. Okazaki \(Eds\.\),Toronto, Canada,pp\. 12286–12312\.External Links:[Link](https://aclanthology.org/2023.acl-long.687/),[Document](https://dx.doi.org/10.18653/v1/2023.acl-long.687)Cited by:[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px2.p3.1)\.
- \[11\]H\. Lightman, V\. Kosaraju, Y\. Burda, H\. Edwards, B\. Baker, T\. Lee, J\. Leike, J\. Schulman, I\. Sutskever, and K\. Cobbe\(2023\)Let’s verify step by step\.arXiv preprint arXiv:2305\.20050\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px1.p1.1)\.
- \[12\]OpenAI\(2025\)Gpt\-oss\-120b & gpt\-oss\-20b model card\.arXiv preprint arXiv:2508\.10925\.Cited by:[§5\.2](https://arxiv.org/html/2605.21770#S5.SS2.SSS0.Px2.p1.1)\.
- \[13\]K\. Park, Y\. J\. Choe, and V\. Veitch\(2023\)The linear representation hypothesis and the geometry of large language models\.arXiv preprint arXiv:2311\.03658\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p3.1),[§2\.2](https://arxiv.org/html/2605.21770#S2.SS2.p3.1)\.
- \[14\]N\. Rimsky, N\. Gabrieli, J\. Schulz, M\. Tong, E\. Hubinger, and A\. Turner\(2024\-08\)Steering llama 2 via contrastive activation addition\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\),L\. Ku, A\. Martins, and V\. Srikumar \(Eds\.\),Bangkok, Thailand,pp\. 15504–15522\.External Links:[Link](https://aclanthology.org/2024.acl-long.828/),[Document](https://dx.doi.org/10.18653/v1/2024.acl-long.828)Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p2.1),[§2\.1](https://arxiv.org/html/2605.21770#S2.SS1.p1.1)\.
- \[15\]D\. Santos\-Martins, L\. Solis\-Vasquez, A\. F\. Tillack, M\. F\. Sanner, A\. Koch, and S\. Forli\(2021\-02\-09\)Accelerating autodock4 with gpus and gradient\-based local search\.Journal of Chemical Theory and Computation17\(2\),pp\. 1060–1073\.External Links:ISSN 1549\-9618,[Document](https://dx.doi.org/10.1021/acs.jctc.0c01006),[Link](https://doi.org/10.1021/acs.jctc.0c01006)Cited by:[§5\.2](https://arxiv.org/html/2605.21770#S5.SS2.SSS0.Px1.p1.1)\.
- \[16\]A\. Turner, L\. Thiergart, D\. Udell, J\. Leike, U\. Mini, and M\. MacDiarmid\(2023\)Steering language models with activation engineering\.arXiv preprint arXiv:2308\.10248\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p2.1),[§2\.1](https://arxiv.org/html/2605.21770#S2.SS1.p1.1)\.
- \[17\]J\. Uesato, N\. Kushman, R\. Kumar, F\. Song, N\. Siegel, L\. Wang, A\. Creswell, G\. Irving, and I\. Higgins\(2022\)Solving math word problems with process\- and outcome\-based feedback\.arXiv preprint arXiv:2211\.14275\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p1.1)\.
- \[18\]H\. M\. Vu and T\. M\. Nguyen\(2025\)Angular steering: behavior control via rotation in activation space\.InThe Thirty\-ninth Annual Conference on Neural Information Processing Systems,Cited by:[§B\.1](https://arxiv.org/html/2605.21770#A2.SS1.p1.4),[§1](https://arxiv.org/html/2605.21770#S1.p2.1),[§2\.1](https://arxiv.org/html/2605.21770#S2.SS1.p1.1),[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px2.p2.1)\.
- \[19\]K\. Wang, A\. Variengien, A\. Conmy, B\. Shlegeris, and J\. Steinhardt\(2022\)Interpretability in the wild: a circuit for indirect object identification in GPT\-2 small\.arXiv preprint arXiv:2211\.00593\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p3.1),[§2\.2](https://arxiv.org/html/2605.21770#S2.SS2.p2.1)\.
- \[20\]X\. Wang, J\. Wei, D\. Schuurmans, Q\. V\. Le, E\. H\. Chi, S\. Narang, A\. Chowdhery, and D\. Zhou\(2023\)Self\-consistency improves chain of thought reasoning in language models\.InThe Eleventh International Conference on Learning Representations,Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p1.1)\.
- \[21\]D\. Weininger\(1988\)SMILES, a chemical language and information system\. 1\. introduction to methodology and encoding rules\.Journal of Chemical Information and Computer Sciences28\(1\),pp\. 31–36\.External Links:[Document](https://dx.doi.org/10.1021/ci00057a005),[Link](https://doi.org/10.1021/ci00057a005),https://doi\.org/10\.1021/ci00057a005Cited by:[§5\.2](https://arxiv.org/html/2605.21770#S5.SS2.SSS0.Px1.p1.1)\.
- \[22\]X\. Yue, X\. Qu, G\. Zhang, Y\. Fu, W\. Huang, H\. Sun, Y\. Su, and W\. Chen\(2023\)MAmmoTH: building math generalist models through hybrid instruction tuning\.arXiv preprint arXiv:2309\.05653\.Cited by:[§5\.1](https://arxiv.org/html/2605.21770#S5.SS1.SSS0.Px3.p1.1)\.
- \[23\]Z\. Zhang, X\. Wu, Z\. Zhou, Q\. Wu, Y\. Zhang, P\. Ponnusamy, H\. Subbaraj, J\. Wang, S\. L\. Song, and B\. Athiwaratkun\(2025\)Understanding and steering the cognitive behaviors of reasoning models at test\-time\.arXiv preprint arXiv:2512\.24574\.Cited by:[§2\.1](https://arxiv.org/html/2605.21770#S2.SS1.p1.1)\.
- \[24\]A\. Zou, L\. Phan, S\. Chen, J\. Campbell, P\. Guo, R\. Ren, A\. Pan, X\. Yin, M\. Mazeika, A\. Dombrowski,et al\.\(2023\)Representation engineering: a top\-down approach to ai transparency\.arXiv preprint arXiv:2310\.01405\.Cited by:[§1](https://arxiv.org/html/2605.21770#S1.p2.1),[§2\.1](https://arxiv.org/html/2605.21770#S2.SS1.p1.1),[§2\.2](https://arxiv.org/html/2605.21770#S2.SS2.p3.1)\.

## Appendix AProof of Proposition 1

###### Proof\.

Expanding using \([9](https://arxiv.org/html/2605.21770#S4.E9)\):

⟨𝒂~t\(l,h\),𝒗⟩\\displaystyle\\left\\langle\\tilde\{\\bm\{a\}\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle=⟨𝒂t\(l,h\),𝒗⟩−⟨𝑩\(l,h\)⊤​𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\),𝒗⟩\\displaystyle=\\left\\langle\\bm\{a\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle\-\\left\\langle\\bm\{B\}^\{\(l,h\)\\top\}\\bm\{B\}^\{\(l,h\)\}\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\),\\,\\bm\{v\}\\right\\rangle\(12\)=⟨𝒂t\(l,h\),𝒗⟩−⟨𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\),𝑩\(l,h\)​𝒗⟩\\displaystyle=\\left\\langle\\bm\{a\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle\-\\left\\langle\\bm\{B\}^\{\(l,h\)\}\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\),\\,\\bm\{B\}^\{\(l,h\)\}\\bm\{v\}\\right\\rangle\(13\)=⟨𝒂t\(l,h\),𝒗⟩−⟨𝑩\(l,h\)​\(𝒂t\(l,h\)−𝝁c\(l,h\)\),0⟩=⟨𝒂t\(l,h\),𝒗⟩\.∎\\displaystyle=\\left\\langle\\bm\{a\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle\-\\left\\langle\\bm\{B\}^\{\(l,h\)\}\(\\bm\{a\}\_\{t\}^\{\(l,h\)\}\-\\bm\{\\mu\}\_\{c\}^\{\(l,h\)\}\),\\,\\bm\{0\}\\right\\rangle\\;=\\;\\left\\langle\\bm\{a\}\_\{t\}^\{\(l,h\)\},\\,\\bm\{v\}\\right\\rangle\.\\qed\(14\)

## Appendix BAdditional Experimental Details

### B\.1Ablation on Hyperparameter Sensitivity

For ITI, we follow the hyperparameter ranges reported in the original paper\[[9](https://arxiv.org/html/2605.21770#bib.bib12)\], sweeping over the number of steered heads in\{24,48,96\}\\\{24,48,96\\\}and intervention strengthα∈\{0\.5,1\.0,5\.0\}\\alpha\\in\\\{0\.5,1\.0,5\.0\\\}\. For Angular Steering, rotation angles are sampled uniformly across the full360∘360^\{\\circ\}range as evaluated in the original work\[[18](https://arxiv.org/html/2605.21770#bib.bib13)\]\. For MAGS, we select hyperparameters to minimize unnecessary interference: we restrict the number of steered heads to a small set \(top\-1 or top\-3 by manifold signal strength\) and sweep projection strengthα∈\{0\.3,0\.5,0\.7,1\.0\}\\alpha\\in\\\{0\.3,0\.5,0\.7,1\.0\\\}, reflecting the principle that corrections should be both targeted and conservative\. Corrections are applied only to heads where drift is detected and only to the extent required to return activations to the correct\-output manifold\.

Table 4:Hyperparameter ablations on MATH\-500 \(Llama\-3\.1\-8B\-Instruct\)\. Best result per method inbold\.\(a\)MAGS: top\-kkheads and steering strengthα\\alpha\.
\(b\)ITI: number of steered headsKKand strengthα\\alpha\.
\(c\)Angular Steering: rotation angle\.

![Refer to caption](https://arxiv.org/html/2605.21770v1/x1.png)Figure 5:Relative attention shiftΔ​W/Wunsteered\\Delta W/W\_\{\\text\{unsteered\}\}at layerℓbip\\ell\_\{\\text\{bip\}\}for two representative problems from MATH\-500, steered with Gemma\-4\-E4b\-it\. Each cell shows the head\-averaged attention change between the steered and unsteered model on the same forced token sequence\. The dashed horizontal line marks the first correction step\.tfiret\_\{\\text\{fire\}\}

## Appendix CVisualization of the effect of steering on attention graph\.

To examine how MAGS reshapes information routing, we visualize the*relative attention shift*at a layer after the correction layer, denotedℓbip\\ell\_\{\\mathrm\{bip\}\}\. Since MAGS modifies the head*output*𝐚h=Ah​Vh\\mathbf\{a\}\_\{h\}=A\_\{h\}V\_\{h\}rather than the routing matrixAhA\_\{h\}directly, the corrected residual is written into the key\-value cache at positiontfiret\_\{\\mathrm\{fire\}\}\. Then, subsequent query steps can attend to this corrected entry\. To isolate the effect of accumulated KV\-cache differences, both steered and unsteered streams are run on the same forced token sequence \(the steered model’s greedy outputs\), so the two caches diverge solely through the corrections\. For each problem we average attention weights over all heads and compute

Δ​WW=Wsteered−WunsteeredWunsteered\+ε\.\\frac\{\\Delta W\}\{W\}\\;=\\;\\frac\{W\_\{\\mathrm\{steered\}\}\-W\_\{\\mathrm\{unsteered\}\}\}\{W\_\{\\mathrm\{unsteered\}\}\+\\varepsilon\}\.Figure[5](https://arxiv.org/html/2605.21770#A2.F5)shows two representative problems from the MATH\-500 dataset for Gemma\-4\-E4b\-it; red \(blue\) indicates positions attended to more \(less\) by the steered model relative to the baseline\.

The strongest signal in both panels is a set of*vertical stripes*: certain key tokens receive consistently higher attention across all subsequent query steps\. This reveals that MAGS does not merely fix a single prediction in isolation, but makes the corrected reasoning step a more*retrievable*context anchor: the entire generation following the correction keeps routing information through those tokens, reinforcing their influence persistently\. The effect is consistent across both problems despite their different mathematical content, suggesting that this “attention anchoring” behavior may be a general mechanism by which MAGS steers the model back onto a correct reasoning trajectory\.

### C\.1Computational infrastructure\.

Experiments with Llama\-3\.1\-8B\-Instruct and Gemma\-4\-E4b\-it were conducted on NVIDIA RTX 4090 GPUs \(24 GB VRAM\) with model weights loaded infloat16precision\. GPT\-OSS 20B experiments were run on NVIDIA H200 GPUs \(141 GB HBM3\) due to the larger memory footprint of the model\.

## Appendix DStatistical Significance

All confidence intervals are 95% percentile bootstrap CIs withB=10000B\{=\}10000resamples \(seed = 42\)\. Point estimates are original trial accuracies\.

Table 5:Bootstrap 95% CIs on Gemma\-4\-E4b\-it across all benchmarks \(B=10,000B\{=\}10\{,\}000resamples, percentile method\)\.Table 6:Bootstrap 95% CIs on Llama\-3\.1\-8B\-Instruct across all benchmarks \(B=10,000B\{=\}10\{,\}000resamples, percentile method\)\.