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

要約

大規模な言語モデル（LLM）は、コンテキストの理解に大幅な改善を実証しています。
ただし、長い文脈の推論と世代の間に真に重要な情報に出席する能力は、まだペースの遅れになります。
具体的には、予備的な実験では、特定の注意散漫パターンが推論中のモデルの注意を誤って指示し、これらのパターンを削除すると、推論の正確性と生成品質が大幅に向上することが明らかになります。
この現象は、トレーニングデータの偽の相関関係に起因するものであり、これにより、本物の因果的指導反応関係を推測するモデルの能力を妨害します。
この現象は、冗長な推論プロセスを誘導する可能性があり、潜在的に有意な推論オーバーヘッドをもたらし、より重要なことに、誤ったまたは最適ではない応答の生成をもたらします。
これを緩和するために、解き間交絡因子に対する介入ベースの推論を活用するLearning to Focus（Leaf）という2段階のフレームワークを紹介します。
最初の段階では、Leafは上級教師と勾配ベースの比較を使用して、トレーニングコーパスの因果関係に基づいて交絡トークンを自動的に識別します。
次に、第2段階では、蒸留中にこれらのトークンを剪定して介入を制定し、生徒の注意を真の重要なコンテキストトークンに教師の焦点分布に合わせます。
実験結果は、葉がさまざまな数学的推論とコード生成ベンチマークの絶対的な改善を達成するだけでなく、推論中の交絡トークンへの注意を効果的に抑制し、より解釈可能で信頼性の高い推論モデルをもたらすことを示しています。

要約(オリジナル)

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 and code generation benchmarks but also effectively suppresses attention to confounding tokens during inference, yielding a more interpretable and reliable reasoning model.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google