AtManRL: Training More Honest Reasoning Models with Differentiable Attention Saliency

This article introduces the AtManRL method, which aims to address the "dishonesty" problem in Chain-of-Thought (CoT) reasoning of Large Language Models (LLMs) — i.e., the reasoning process may be irrelevant to answer generation. The method identifies key tokens in reasoning chains via differentiable attention masks, combines saliency rewards with outcome rewards, and jointly optimizes the correctness and interpretability of reasoning under the GRPO framework, providing a new path for building trustworthy AI.

Although LLMs have strong CoT reasoning capabilities, there is a fundamental question: do the reasoning steps actually affect answer generation? Researchers define "faithful reasoning" to meet three criteria: 1. Causal relevance (reasoning steps participate in answer generation); 2. Interpretability (humans can understand the reasoning logic); 3. Consistency (the same reasoning leads to the same conclusion). Existing models often have "reasoning shortcuts", such as generating irrelevant steps or retroactively constructing explanations.

The core of AtManRL (Attention Manipulation Reinforcement Learning) includes: 1. Additive Attention Mask: Identifies key tokens in CoT that affect the answer, supports end-to-end training, and has sparsity constraints; 2. Saliency Reward: Evaluates the actual impact of key tokens on predictions based on the mask, giving positive rewards only when reasoning tokens truly affect the answer, directly optimizing reasoning faithfulness.

AtManRL combines two types of rewards under the GRPO (Group Relative Policy Optimization) framework: 1. Outcome Reward: Based on answer correctness (positive for correct, negative for incorrect); 2. Saliency Reward: Based on the actual impact of reasoning on the answer (positive for relevant, negative for irrelevant). Joint optimization balances correctness and interpretability, avoiding the limitations of a single objective (e.g., optimizing only correctness leads to shortcut reasoning).

The research team validated AtManRL using Llama-3.2-3B-Instruct as the base model on mathematical reasoning (GSM8K) and general knowledge reasoning (MMLU) tasks: 1. Successfully identified key tokens in CoT (e.g., intermediate calculation results, logical transitions); 2. Generated CoT with more coherent logic, fewer irrelevant steps, and stronger interpretability; 3. Maintained accuracy comparable to using only outcome rewards, while significantly improving faithfulness.

The significance of AtManRL lies in: 1. Causal Modeling: Going beyond the correlation of traditional attention visualization to explicitly model the causal impact of tokens on predictions; 2. Training-Time Intervention: Guiding the model to generate faithful reasoning from the source, which is more efficient than post-hoc explanations; 3. Scalability: Compatible with Transformer LLMs and existing RLHF frameworks (e.g., GRPO, PPO), with controllable computational overhead.

AtManRL has limitations: 1. Quantitative evaluation of faithfulness remains an open problem; 2. Only validated on reasoning-intensive tasks, and the effect on open-ended tasks needs to be tested; 3. Reward balancing requires fine parameter tuning. Future directions: Develop more refined faithfulness metrics, explore multimodal scenarios, and expand to larger model scales.

AtManRL transforms "faithful reasoning" into an optimizable training objective, improves the transparency of LLM reasoning, and lays the foundation for building trustworthy AI systems. As LLMs are increasingly applied in high-risk decision-making scenarios, ensuring the honesty of reasoning becomes more important, and AtManRL provides a promising technical direction.