Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Manifold-Guided Attention Steering

About

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 propose Manifold-Guided Attention Steering (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.

Ian Li, Kapilesh Guruprasad, Raunak Sengupta, Ninad Satish, Loris D'Antoni, Rose Yu• 2026

Related benchmarks

TaskDatasetResultRank
Code GenerationHumanEval (test)--
612
Code GenerationMBPP (test)--
405
Mathematical ReasoningMATH-500 (test)
Accuracy64.8
46
Code GenerationHumanEval
Accuracy60.4
5
Code GenerationMBPP
Accuracy57.4
5
Mathematical ReasoningMATH 500
Accuracy53
5
Mathematical ReasoningGSM8K
Accuracy86.7
5
Mathematical ReasoningGSM8K (test)
Accuracy91.3
5
Molecular GenerationMolecular Generation SMILES
Validity54.8
5
Showing 9 of 9 rows

Other info

Follow for update