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

TL;DR

This paper introduces LeaF, a two-stage framework that improves large language models' focus on critical tokens during reasoning by identifying and pruning confounding tokens, leading to better accuracy and interpretability.

Contribution

LeaF is a novel intervention-based distillation method that automatically identifies and prunes confounding tokens to enhance model focus and reasoning accuracy.

Findings

Improves reasoning accuracy on multiple benchmarks

Reduces attention to confounding tokens during inference

Enhances interpretability and reliability of LLMs

Abstract

Large language models (LLMs) have demonstrated significant improvements in contextual understanding. However, their ability to attend to truly critical information during long-context reasoning and generation still falls behind the pace. Specifically, our preliminary experiments reveal that certain distracting patterns can misdirect the model's attention during inference, and removing these patterns substantially improves reasoning accuracy and generation quality. We attribute this phenomenon to spurious correlations in the training data, which obstruct the model's capacity to infer authentic causal instruction-response relationships. This phenomenon may induce redundant reasoning processes, potentially resulting in significant inference overhead and, more critically, the generation of erroneous or suboptimal responses. To mitigate this, we introduce a two-stage framework called Learning to Focus (LeaF) leveraging intervention-based inference to disentangle confounding factors. In the first stage, LeaF employs gradient-based comparisons with an advanced teacher to automatically identify confounding tokens based on causal relationships in the training corpus. Then, in the second stage, it prunes these tokens during distillation to enact intervention, aligning the student's attention with the teacher's focus distribution on truly critical context tokens. Experimental results demonstrate that LeaF not only achieves an absolute improvement in various mathematical reasoning, code generation and multi-hop question answering benchmarks but also effectively suppresses attention to confounding tokens during inference, yielding a more interpretable and reliable reasoning model.