AtManRL: Towards Faithful Reasoning via Differentiable Attention Saliency
| Authors | Max Henning Höth et al. |
| Year | 2026 |
| Field | NLP |
| arXiv | 2604.16158 |
| Download | |
| Categories | cs.CL, cs.AI, cs.LG |
Abstract
Large language models (LLMs) increasingly rely on chain-of-thought (CoT) reasoning to solve complex tasks. Yet ensuring that the reasoning trace both contributes to and faithfully reflects the processes underlying the model's final answer, rather than merely accompanying it, remains challenging. We introduce AtManRL, a method that leverages differentiable attention manipulation to learn more faithful reasoning through reinforcement learning. By training an additive attention mask that identifies tokens in the CoT crucial for producing correct answers, we derive a saliency reward signal that encourages the model to generate reasoning traces that genuinely influence its final predictions. We integrate this saliency reward with outcome-based rewards within the GRPO framework to jointly optimize for correctness and interpretability. Experiments on GSM8K and MMLU with Llama-3.2-3B-Instruct demonstrate that our approach can identify influential reasoning tokens and enable training more transparent reasoning models.
Engineering Breakdown
Plain English
AtManRL addresses a fundamental problem with chain-of-thought reasoning in large language models: the reasoning traces LLMs generate often don't actually explain how the model arrives at its answer—they just accompany it. The paper introduces a method that uses differentiable attention manipulation combined with reinforcement learning to make reasoning traces more faithful and interpretable. The core idea is to train an additive attention mask that identifies which tokens in the reasoning chain are genuinely crucial for producing correct answers, then use this as a reward signal during training. By integrating saliency rewards with outcome-based rewards in the GRPO framework, the authors encourage models to generate reasoning that truly influences predictions rather than post-hoc rationalizations.
Core Technical Contribution
The key innovation is a differentiable attention saliency mechanism that can be optimized end-to-end through reinforcement learning to promote faithful reasoning. Rather than treating reasoning traces as a side effect or decoration on model outputs, the authors introduce an explicit saliency reward signal derived from attention manipulation—measuring which tokens in the CoT actually contribute to the final answer through gradient analysis. This is different from prior work that either treats interpretability separately from performance optimization, or uses fixed post-hoc explanation methods. The integration with reinforcement learning (specifically GRPO) creates a unified training objective that simultaneously optimizes both answer correctness and reasoning faithfulness, making them complementary rather than competing goals.
How It Works
The method operates by training an additive attention mask alongside the language model's standard attention computations. During forward passes, this learnable mask weights different tokens in the chain-of-thought sequence, and differentiable gradient analysis through this mask identifies which tokens have the highest attribution to the final logits (saliency). The model generates a reasoning trace and produces an answer; the attention mask determines which tokens in that trace are most consequential. A saliency reward signal is computed by measuring the relationship between attention mask weights and gradient attribution—tokens with high attention that also have high gradients receive higher rewards. This saliency reward is then combined with standard outcome rewards (whether the answer was correct) within the GRPO reinforcement learning framework, creating a composite reward that encourages both correctness and faithful reasoning. The training process updates both the language model parameters and the attention mask parameters to maximize this joint objective, converging toward policies that generate reasoning traces that demonstrably influence the final predictions.
Production Impact
In production systems, this addresses a critical trust and debuggability issue: when models fail or make unexpected decisions, engineers need to understand whether the reasoning trace is actually explanatory or just confabulation. Adopting AtManRL would mean incorporating attention saliency computation into your training pipeline, adding modest computational overhead during training but providing immediate benefits at inference time—you can now audit whether a model's reasoning actually influenced its output. For high-stakes applications (legal analysis, medical diagnosis, scientific discovery), this moves from "the model said this reasoning" to "the model demonstrably used this reasoning in its decision." The trade-off is increased training complexity: you need to maintain the attention mask parameters and compute gradient attributions during RLHF, roughly 15-25% additional GPU memory and compute overhead compared to standard RLHF. Integration requires modifying your RL training loop to extract and score saliency rewards, which is straightforward but non-trivial if you're using existing RLHF infrastructure.
Limitations and When Not to Use This
The method assumes that gradient-based saliency through attention is a faithful proxy for true causal contribution to model decisions—but attention weights and gradients can be misleading, and this approach doesn't prove counterfactual causation in the formal sense. The approach is computationally expensive during training and requires differentiable access to the full model (incompatible with APIs or frozen model weights), limiting applicability to open-weight models you can fine-tune. The paper's evaluation is presumably on standard benchmarks, but it's unclear how well the saliency rewards transfer to out-of-distribution prompts or whether the method avoids just learning to produce superficially coherent reasoning that passes the saliency test without genuine faithfulness. Finally, the technique requires explicit chain-of-thought generation, which doesn't help with models that reason implicitly or models where you want to extract faithfulness from standard autoregressive inference without step-by-step reasoning.
Research Context
This work builds on the chain-of-thought reasoning literature (Wei et al., 2022) and recent work on interpreting reasoning in LLMs, extending beyond earlier post-hoc explanation methods (LIME, SHAP, attention visualization) into the training loop itself. It's positioned as an improvement over standard RLHF and GRPO methods that optimize for answer correctness alone without explicitly enforcing that intermediate reasoning steps are causal. The saliency-based reward signal draws inspiration from mechanistic interpretability research that uses gradient analysis to identify important features and tokens. This opens up a research direction toward training objectives that jointly optimize performance and interpretability, potentially enabling future work on self-explaining models and AI systems that provide genuine explanations rather than plausible-sounding narratives.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::