AtManRL: Towards Faithful Reasoning via Differentiable Attention Saliency

Abstract: 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.

• The paper proposes a novel RL approach using a differentiable attention mask to optimize reasoning trace saliency and align intermediate steps with final answers.
• It achieves a 44%-46% reduction in chain-of-thought length on GSM8K and MMLU benchmarks, significantly improving inference efficiency and token density.
• The method maintains solution accuracy while promoting concise, information-dense reasoning, offering practical benefits for model transparency and auditing.

AtManRL: Differentiable Attention Saliency for Faithful LLM Reasoning

Introduction

The proliferation of chain-of-thought (CoT) prompting and RL-based fine-tuning has advanced the reasoning capabilities of LLMs, yet the critical issue of reasoning faithfulness remains unresolved. In "AtManRL: Towards Faithful Reasoning via Differentiable Attention Saliency" (2604.16158), the authors propose AtManRL—a reinforcement learning approach that utilizes a differentiable attention manipulation mask to explicitly optimize reasoning trace saliency. By targeting the causal dependency between generated reasoning steps and the model's final answer, AtManRL provides a nuanced mechanism for enforcing interpretability that surpasses post-hoc rationalization.

Methodology

Differentiable Attention Manipulation

AtManRL is grounded in the AtMan framework, which manipulates transformer attention by additively adjusting the pre-softmax attention scores through a learned mask:

Figure 1: Initialization of the additive attention mask $H^{\text{AtMan}}$, which suppresses attention over chain-of-thought tokens prior to differentiable optimization.

This mask is trained end-to-end to reveal which tokens in the CoT trace are truly salient for the model's prediction by measuring their impact on answer log-probabilities. The mask is optimized for each rollout as the only trainable object in the system, using gradient descent to recover the original answer probabilities when attention to reasoning tokens is suppressed.

Figure 2: Optimization of the attention mask per rollout, with the objective of restoring label log-probabilities and identifying non-salient tokens.

Saliency is then quantified for each token as the extent to which attention to it must be re-enabled to preserve correct answer likelihood, providing a granular, data-driven measure of influence within the reasoning trace.

Saliency-Guided RL

The computed token-wise saliency is integrated into the RL training loop via a bespoke reward. The overall reward is the sum of a traditional outcome-based component (for answer correctness) and a saliency-based component. The policy is trained using the clipped GRPO objective, with group-normalized, advantage-weighted updates. Training is conducted on Llama-3.2-3B-Instruct using GRPO, and the entire process is designed to encourage models to generate concise, information-dense, and causally-relevant reasoning steps that remain essential for downstream answer correctness.

Experimental Evaluation

Quantitative Performance

AtManRL is evaluated on GSM8K (math word problems) and MMLU (massive multitask language understanding), with ablation against a strong outcome-only RL baseline. Key results include:

• Average CoT length reduced by 44% (GSM8K) and 46% (MMLU), significantly lowering inference cost and trace verbosity.
• Content shifts toward higher numbers/symbols and reduced stop-word usage: on GSM8K, stop words drop by 25%, while numbers/symbols increase by 48%/61%, indicating higher information density.
• Negligible loss (or small improvement) in solution accuracy: Pass@4 drops only 0.4% on GSM8K and slightly improves on MMLU (+0.1%).

Training dynamics indicate that accuracy is preserved while mean response length consistently drops over time, further corroborating the efficiency of the approach.

Figure 3: Accuracy and mean response length of GSM8K evaluation set across training epochs. Both metrics remain stable as AtManRL reduces response length.

Figure 4: Accuracy and mean response length of MMLU evaluation set during training, demonstrating maintained accuracy with minimized reasoning traces.

Qualitative Analysis

Fine-grained comparisons of reasoning traces reveal a marked decrease in superfluous verbiage and a greater focus on computationally relevant tokens under AtManRL, without omitting necessary intermediate reasoning.

Figure 5: Reasoning trace comparison between baseline and AtManRL on GSM8K; AtManRL outputs are notably more concise and contain fewer filler tokens.

Figure 6: Additional GSM8K examples illustrating AtManRL’s preference for information-dense, critical reasoning steps.

Figure 7: On MMLU, AtManRL-generated traces similarly exhibit reduced length with no observed accuracy degradation.

Figure 8: Another MMLU example: AtManRL maintains key inferential content while omitting irrelevant or weakly causal tokens.

Theoretical and Practical Implications

By reframing the faithfulness problem as a measurable causal influence of reasoning trace tokens, AtManRL introduces a direct mechanism for aligning LLM explanations with their actual prediction pathways. This provides a more rigorous standard than plausibility- or outcome-based explanations and stands to improve the reliability of LLM-generated rationales in sensitive domains such as scientific reasoning, education, and safety-critical applications.

The marked reduction in non-salient reasoning content without compromising performance points to strong practical implications for both compute efficiency and transparency. As token-level saliency tools mature, they offer potential integration into model auditing, training pipeline monitoring, and dynamic reasoning compression during deployment.

Future Directions

The work highlights several open questions:

• Robustness across scales and domains: Further evaluation on larger models and diverse datasets is necessary to ensure generalization beyond small-scale experiments.
• Faithfulness beyond saliency: While token saliency is necessary for faithfulness, it is not sufficient. Future research should directly measure internal-external alignment of reasoning.
• Self-correction behavior: The current saliency reward discourages verbosity, which could inadvertently suppress beneficial self-correction or exploration behaviors. Reward shaping to explicitly preserve such phenomena warrants investigation.
• Integration with critical token discovery and long-context modeling: The approach may synergize with recent research on critical tokens (Lin et al., 2024), selective attention (Wang et al., 20 Feb 2025), and context compression (Akhauri et al., 10 Mar 2025).

Conclusion

AtManRL advances faithful LLM reasoning by leveraging differentiable attention manipulation to anchor chain-of-thought traces in true model causality, rather than mere output plausibility. The method jointly optimizes for outcome accuracy and interpretable saliency, substantiated by substantial reductions in spurious reasoning content without loss of accuracy. This strategy offers a principled path toward more transparent and trustworthy LLMs and sets a new paradigm for aligning explanations with underlying model computation.