Research map / Large language models / Reasoning and chain of thought

Learning to Focus: Causal Attention Distillation via Gradient‐Guided Token Pruning

NeurIPS 2025 · Large language models › Reasoning and chain of thought · approach: Test-time compute scaling · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes LeaF, which uses teacher-student gradient sensitivity differences to find confounding tokens in instructions and responses, prunes them into counterfactual samples, and distills the student on both original and pruned inputs to improve reasoning, code, and multi-hop QA.

Problem

LLMs are distracted by spurious confounding tokens in instructions and generated responses, which misdirects attention and degrades long-context reasoning accuracy and generation quality in small student models.

Approach

Identify confounding tokens by comparing gradient sensitivities of a high-capacity teacher and a student on mispredicted examples, create counterfactual samples by span pruning, and apply a combined KL distillation loss on original and counterfactual inputs so the student aligns with the teacher's focus.

Key results

• LeaF (Instr & Resp Mask) raises LLaMA3.2-1B-Instruct MathBench average from 33.03 (KD) to 34.71 and LLaMA3.2-3B from 50.29 to 53.21.
• On CodeBench, LeaF lifts LLaMA3.2-3B average from 31.89 (KD) to 34.53.
• On multi-hop QA with LLaMA3.2-3B, LeaF (Instr Mask) improves average from 39.03 to 42.27.

Measured effects reported in the paper

Surprising findings

• Collective pruning of all confounding tokens at once degrades LLaMA3.2-3B performance on MATH-500 relative to KD, while span pruning helps.
• Removing distracting patterns from instructions without training already yields over 20% gain on MATH training corpus for small models (preliminary experiment).

Limitations

• Depends on an advanced teacher model (70B-class) to identify confounding tokens.
• Only validated on math and code tasks; long-form generation and other domains untested.
• Gains are modest on some benchmarks and evaluated on relatively small student models.

Experimental setup

Datasets: NuminaMath-CoT AceCode-87K KILT HotpotQA NQ PopQA

Benchmarks: GSM8K MATH-500 OlympiadBench HumanEval+ LeetCode LiveCodeBench 2WikiMultiHopQA Musique HotpotQA

Metrics: accuracy pass@1 pass@10 EM F1 SubEM

Baselines: KD w/o Mask Instruct model (Pre-KD) Random Masking PPL-based Masking Collective Pruning

Builds on: chain-of-thought knowledge distillation Step-DPO RHO-1 TokenSkip LLMLingua

Closest papers by meaning

• Don't Take Things Out of Context: Attention Intervention for Enhancing Chain-of-Thought Reasoning in Large Language Models (ICLR 2025 · similarity 0.88)
• Eliciting Causal Abilities in Large Language Models for Reasoning Tasks (AAAI 2025 · similarity 0.88)
• Can Large Language Models Infer Causation from Correlation? (ICLR 2024 · similarity 0.87)
• Longer Context, Deeper Thinking: Uncovering the Role of Long-Context Ability in Reasoning (NeurIPS 2025 · similarity 0.87)
• Attention Illuminates LLM Reasoning: The Uncovered Preplan-and-Anchor Rhythm Enables Fine-Grained Policy Optimization (ICML 2026 · similarity 0.87)
• Enhancing Elusive Clues in Knowledge Learning by Contrasting Attention of Language Models (AAAI 2025 · similarity 0.86)