通过梯度引导剪枝,让大模型学会聚焦关键信息,减少干扰。
Learning to Focus: Causal Attention Distillation via Gradient-Guided Token Pruning
- 基于梯度对比识别训练数据中的干扰词,定位混淆因素。
- 在蒸馏中剪除干扰词,使学生模型注意力更聚焦于关键上下文。
- 提升数学推理与代码生成准确率,结果更可解释可信。
大型语言模型在上下文理解方面已取得显著进展,但在长上下文推理与生成中关注真正关键信息的能力仍不足。初步实验发现,某些干扰模式会误导模型注意力,移除这些模式能显著提升推理准确率与生成质量。我们归因于训练数据中的伪相关性,阻碍了模型对真实因果关系的推断。这可能导致冗余推理过程,带来显著计算开销,并引发错误或次优输出。为此,我们提出两阶段框架Learning to Focus(LeaF),采用干预式推理分离混杂因素。第一阶段,利用梯度对比与先进教师模型,自动识别训练语料中基于因果关系的混淆标记;第二阶段,在蒸馏过程中剪除这些标记以实现干预,使学生模型注意力分布与教师对关键上下文标记的关注一致。实验表明,LeaF不仅在数学推理、代码生成和多跳问答等基准上取得绝对性能提升,还能有效抑制推理中对混淆标记的关注,获得更可解释、可靠的推理模型。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。