AtManRL: Towards Faithful Reasoning via Differentiable Attention Saliency
TL;DR
The results suggest that training explicitly for saliency is a promising route towards chain-of-thought reasoning that is more faithful to the model's actual decision process.
Abstract
Large language models (LLMs) increasingly rely on chain-of-thought (CoT) reasoning to solve complex tasks, yet a reasoning trace may merely accompany the final answer rather than causally contribute to it or faithfully reflect how it was reached. We introduce differentiable attention manipulation to train salient reasoning with reinforcement learning -- method which we call AtManRL. By optimizing an additive attention mask that identifies the CoT tokens needed to recover the model's answer, we derive a saliency reward and combine it with an outcome reward in GRPO with decoupled advantages. Because both saliency and faithfulness depend strongly on trace length, we evaluate against a sweep of length-penalty baselines that traces the trade-off a pure length penalty can reach, rather than against a single baseline. Across three models (Qwen2.5-3B, Granite-3.3-2B, and Phi-4-mini) and four benchmarks (GSM8K, MMLU, ProofWriter, and QuestBench), AtManRL sits above this trade-off curve by 3.9-4.0pp in saliency, measured as answer changes under mistake injection, and by 5.7-7.7pp in faithfulness, measured as hint verbalization under hint injection, while matching the accuracy of baselines of similar length. These results suggest that training explicitly for saliency is a promising route towards chain-of-thought reasoning that is more faithful to the model's actual decision process.
