世界を動かす技術を、日本語で。

注意メカニズムが言語モデルの安定性を保つ方法

概要

  • Transformer系モデル の長文会話処理における致命的な失敗原因の発見
  • Attention Sink (注意の逃げ場)という現象とその重要性の解明
  • StreamingLLM による簡単かつ効果的な解決策の提案
  • OpenAI など大手AI企業による実装と設計の違い
  • Attention Sinkの理論的背景と今後の設計方針への影響

長文会話におけるTransformerの失敗とAttention Sinkの発見

  • 言語モデル は長い会話で古いトークンを削除すると、出力が 完全なナンセンス になる現象
  • 最初の数トークン にモデルが大量の注意を割り当てる「 Attention Sink」現象の発見
  • Softmaxの制約 (重みの総和が1)により、使いきれない注意を「逃がす」場所が必要
  • StreamingLLM では最初の4トークンを常時保持し、それ以外はスライドウィンドウで管理する手法を提案
  • この手法により、 数百万トークン 規模の安定した長文処理が可能に

Attention Sinkの理論的背景と既存研究

  • Softmax正規化 が「必ずどこかに注意を割り当てる」という根本的な制約
  • 訓練データ で常に存在する最初のトークンが「注意の逃げ場」として機能
  • BERTVision Transformer でも区切り記号や背景パッチに同様の現象が観測
  • 過去の研究 でも初期トークンが捨てられない設計が見られる

StreamingLLMの実装と効果

  • KVキャッシュ管理 において最初の数トークンだけは絶対に削除しない実装
  • 実装例:最初の4トークン+直近ウィンドウサイズ分のトークンでキャッシュを構成
  • 従来法 では数千トークンで破綻していたモデルが、StreamingLLMで 400万トークン超 の安定生成を実現
  • パープレキシティ(困惑度) も安定し、長文生成が現実的に

Attention Sinkの最適数と訓練時設計

  • なぜ「 4つのSink」が必要なのか、訓練時の設計で検証
  • 専用Sinkトークン を訓練時から導入したモデルは、推論時に1つのSinkで十分安定
  • 通常訓練モデルは4つの内容トークンをSinkとして流用しないと安定しない
  • Sinkトークンの導入 で学習効率や下流性能も向上

OpenAIの実装とStreamingLLMの違い

  • OpenAI は各Attention HeadのSoftmaxに スカラー値1つ を加える設計
  • StreamingLLMは 専用Sinkトークン を用意し、より柔軟な注意分配が可能
  • OpenAI方式は パラメータ効率 が高く、既存モデルへの導入が容易
  • 両方式とも「注意の逃げ場」を確保することで安定性を実現

理論的意義と今後の設計指針

  • Attention Sink は「情報の過混合(over-mixing)」を防ぐ 圧力弁 として機能
  • Sinkがあることで、情報やノイズの拡散を抑制し、 表現の安定性 が向上
  • 大規模モデル ほどAttention Sinkの利用率が高い傾向
  • Softmax正規化 という根本制約から生じる現象であり、今後のアーキテクチャ設計にも重要な示唆

まとめ

  • Attention Sink はTransformerの長文安定化に不可欠な機構
  • StreamingLLM の発見と実装が、業界標準となりつつある現状
  • 今後は Sinkトークンの訓練時導入効率的な逃げ場設計 が重要課題
  • 理論的・実用的意義 の両面でAIモデル設計を大きく前進させる発見

Hackerたちの意見

最初のいくつかのトークンは、しばしば最小限の意味情報しか持っていなかった。時には、シーケンスの開始マーカーや「the」や「a」といった一般的な単語だけだったりする。最初の単語をタイトルみたいに使った方が、文法的に正しい文にすぐ入るよりもいいのかなって思う。

「磁石。どうやって動くの?」

なんか、習慣的な礼儀から「こんにちは」とか「お願いします」とかでプロンプトを始める人がいるみたい。もしそのプロンプトが本当に効果的だったら、モデルがその言葉をアテンションシンクとして使えるから面白いよね。

研究者たちは、BERTでも似たようなパターンを観察していて、「驚くほど多くの注意が区切りトークン [SEP] やピリオドに集中している」と主張していた。これはモデルが一種のノーオプとして使っているということだ。同じ夏、Metaではビジョントランスフォーマーを研究している研究者たちが似たような挙動を見つけて、モデルが情報のない背景パッチを計算用のスクラッチパッドとして再利用する様子を観察した。これはトランスフォーマーだけに留まらないようだ。例えば、少し前に読んだ論文で、GAN/U-Netアーキテクチャを使った画像から画像へのモデルでも似たような効果が示されていたのを思い出す。[1] https://arxiv.org/abs/1712.02950

GANが懐かしいな。トランスフォーマーと同じパフォーマンスを出すのに、データ量やパラメータが多くてもトレーニングがずっと難しいのは分かるけど、そこから出てきた最適化の研究やトリックは本当に良かったよね。識別器のキャパシティに関する研究も超クールだった。

これはトランスフォーマーを超えているように見えるし、それ以上のことだね。時々、AI研究って他のところにある車輪を再発明してる気がするんだ。車輪だけかもしれないけど、それでもね。

これは素晴らしくて役立つね。新しいGPT-OSSモデルがこの技術を使ってるから。元の著者たちに拍手を送りたい!

そして、いつものように、FOSSエコシステムは素早く動いてるね。llama.cppはもう完全にサポートしてる! https://github.com/ggml-org/llama.cpp/pull/15157

Barberoらは、アテンションシンクが「圧力弁」として機能し、研究者たちが「オーバーミキシング」と呼ぶ病的な状態を防ぐことを示している。これは、長いシーケンスを処理する深層モデルがトークン間の重要な区別をぼやけさせる状態だ。シンクの存在は他のトークンから注意を引き離し、情報(とノイズ)の広がりを制限し、より安定した埋め込みをもたらす。これは、間違った理由で機能しているように聞こえる。正しい行動は、最初の数個ではなく、正しいニューロンが注意を受けることだと思う。すべてをそこに詰め込むのは、ぼやけるという補完的な罪だ。ぼやけを防ぐために、スパース性の事前知識と組み合わせたアテンションの均一化を調査したい。

要するに、常に注目すべき正しいトークンがあるわけじゃないってこと。探している情報がそこにないなら、どんな巧妙なアテンションスキームでも見つけられない。そうなったときに期待できるのは、「見つからなかった」場合の値が「見つかった」場合の値と区別できることだけ。アテンションシンクを固定の「見つからなかった」値として機能させるのは、その一つの方法だね。

同じ問題に対するより良い解決策だったみたいだね。https://www.evanmiller.org/attention-is-off-by-one.html

これ、結局うまくいったのかな?あり得そうだけど、実際に検証が必要だね。

gpt-ossで使われているアテンションシンクは、君のリンクに似てるね。でも、分母に1を追加する代わりに、トレーニング可能な「ロジット」を追加してる(各ヘッドごとに異なるロジット)。

もしそれが十分だったら、めっちゃ面白いよね。これってパーセプトロンのバイアス項を思い出させるなぁ。トランスフォーマーでは、最初はバイアスがなかったのに、ネットワークがその入力の一つをバイアスとして再利用しちゃって、みんなイライラしてたんだよね。特定の入力を落とすと全体に不合理な影響が出るし、他の人たちもその入力の重みが異常に高くて、他の入力をバランス取るために使われてたから、余計に困ってた。最初は(ハンラボの人たちが)その入力が落ちないようにしようとしてたんだけど、今度は(オープンAIの人たちが)ネットワーク内に学習可能なバイアスを提供して、その入力をスキップすることにしたんだ。クラシックなパーセプトロンがやってることをやってるわけね。で、今この人が提案してるのは、バイアスをどこでも1に設定することでさらなる最適化を図るってこと。これって、絶対値にはあんまりこだわらないから、うまくいくかもしれない。結局、一番大きいのを選ぶだけだし、何だったかは気にしないからね。トレーニングでは、他の重みがバイアスによってスケールされるから、バイアスは1でもいいんだよね。光の速さを1に設定して物理計算するみたいな感じ。もしシンプルなフィードフォワードネットワークのパーセプトロンがあって、最終的に一番大きい出力を選ぶだけで絶対値は気にしないなら、もしかしたらすべてのパーセプトロンのバイアス項を1に設定して、学習から除外しても大丈夫かも。生物のニューロンではバイアスは学習可能なのかな? 活性化ポテンシャルの閾値(なんて呼ばれるか知らないけど)は、化学に依存してて、すべてのニューロンで同じじゃないんじゃない?

自分のトイ・トランスフォーマーモデルで、ViTのCLSトークンに似た「グローバル」トークンを追加したら、かなり大きな改善が見られたよ。MSリサーチの「Diff transformer」っていう別のアプローチも見たことあるよね。

単一の欠陥(または高ノルムトークン)[1]はアテンションシンクに関連してるかもしれないね。すべての高ノルムトークンが同じ方向を持っているのは興味深い。理論自体はそんなに複雑じゃなくて、トレーニング中にうまく修正できるかもしれない。[1] https://openreview.net/pdf?id=4yBnUokU2v

プロンプトで、アテンションシンクにどの情報を残すべきかヒントを与える方法ってあるのかな?

ただし、私たちはその観察について深く掘り下げなかった。ああ、皮肉だね。

ハハ、LLMで「NULL」ポインタを見つけたみたいだね。