Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs

要約

専門家は、トランスフォーマーベースの大規模言語モデル (LLM) で 3 つの不可解な現象、つまりアテンション シンク、値状態のドレイン、および残差状態のピークを一貫して観察してきました。これらは総称してエクストリーム トークン現象と呼ばれます。
これらの現象は、特定のいわゆる「シンク トークン」が不釣り合いに高い注意の重みを受け取り、著しく小さい値の状態を示し、他のトークンよりもはるかに大きな残差状態規範を持つことを特徴としています。
これらの極端なトークンは、LLM 推論、量子化、解釈可能性においてさまざまな課題を引き起こします。
エクストリームトークン現象のメカニズムを解明します。
まず、これらの現象が、おもちゃのモデルであるバイグラム バックコピー (BB) タスクで訓練された非常に単純なアーキテクチャ (1 ～ 3 層のトランスフォーマー) で発生することを示します。
この設定では、アテンション ヘッドが特定の入力ドメインに対してシンクとなり、他の入力ドメインに対しては非シンクのままである、アクティブ-ドーマント メカニズムを特定します。
トレーニングダイナミクスの理論的分析により、これらの現象は相互強化メカニズムによって引き起こされていることが明らかになりました。
これらの洞察に基づいて、softmax を ReLU に置き換え、Adam を SGD に置き換えるなど、事前トレーニング中の極端なトークン現象を軽減する戦略を提案します。
次に、解析を Llama や OLMo などの事前学習済み LLM に拡張し、多くのアテンション ヘッドが BB タスクと同様のアクティブ – ドーマント メカニズムを示し、相互強化メカニズムも LLM 中の極端なトークン現象の出現を支配していることを示します。
事前トレーニング。
私たちの結果は、BB タスクによって予測されたエクストリーム トークン現象の静的および動的特性の多くが、事前学習済み LLM での観察と一致することを明らかにしました。

要約(オリジナル)

Practitioners have consistently observed three puzzling phenomena in transformer-based large language models (LLMs): attention sinks, value-state drains, and residual-state peaks, collectively referred to as extreme-token phenomena. These phenomena are characterized by certain so-called ‘sink tokens’ receiving disproportionately high attention weights, exhibiting significantly smaller value states, and having much larger residual-state norms than those of other tokens. These extreme tokens give rise to various challenges in LLM inference, quantization, and interpretability. We elucidate the mechanisms behind extreme-token phenomena. First, we show that these phenomena arise in very simple architectures — transformers with one to three layers — trained on a toy model, the Bigram-Backcopy (BB) task. In this setting, we identify an active-dormant mechanism, where attention heads become sinks for specific input domains while remaining non-sinks for others. Our theoretical analysis of the training dynamics reveals that these phenomena are driven by a mutual reinforcement mechanism. Building on these insights, we propose strategies to mitigate extreme-token phenomena during pretraining, including replacing softmax with ReLU and Adam with SGD. Next, we extend our analysis to pretrained LLMs, including Llama and OLMo, showing that many attention heads exhibit a similar active-dormant mechanism as in the BB task, and that the mutual reinforcement mechanism also governs the emergence of extreme-token phenomena during LLM pretraining. Our results reveal that many of the static and dynamic properties of extreme-token phenomena predicted by the BB task align with observations in pretrained LLMs.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google