Attention Sinks

アテンションシンクとは、多くのヘッドやレイヤーでアテンション重みの大きな割合を受け取る一方、クエリには関係しないトークンです。多くの場合、シーケンスの最初のトークンが該当します。MIT、Meta AI、Carnegie Mellon UniversityのGuangxuan Xiaoらは、ICLR 2024で発表した「Efficient Streaming Language Models with Attention Sinks」でこの現象を命名しました。同氏らは、Llama 2、MPT、Falcon、Pythiaなどの大規模言語モデル (LLMs)では、スライディングKVキャッシュによって先頭のトークンが破棄されると生成が破綻し、4つのトークンを保持すると安定した生成が回復することを示しました。

シンクが形成される理由#

自己注意では、各クエリのスコアをソフトマックスで重みに変換し、それらの合計は常に1になります。コンテキスト内にヘッドと関連するトークンがない場合も、重みの割り当て先が必要です。因果的な言語モデルでは、先頭のトークンが後続のすべてのトークンから見えるため、学習によって、ヘッドが必要としない重みの既定の割り当て先になります。Llama 2を分析したStreamingLLMでは、下から2つを除くすべてのレイヤーで、全ヘッドのアテンションが初期トークンに集中していました。ソフトマックスにシンクロジットを加えると、この効果を確認できます。

import torch

torch.manual_seed(0)
values = torch.randn(6, 4)
# A query that matches none of the 6 tokens: every score is close to zero
scores = torch.tensor([0.1, -0.2, 0.0, 0.2, -0.1, 0.0])

# Softmax must hand out a total weight of 1, so the output mixes all six values
weights = scores.softmax(dim=-1)
print(f"No sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

# A sink logit joins the softmax but carries no value, so it absorbs the spare weight
sink = torch.tensor([3.0])
weights = torch.cat([sink, scores]).softmax(dim=-1)[1:]
print(f"With sink: weight on tokens {weights.sum():.2f}, output norm {(weights @ values).norm():.2f}")

シンクがない場合、ヘッドは無関係な6つの値に重みをすべて分散させ、ノルム0.83の出力を返します。シンクロジットを3にすると、シンクが重みの77%を占め、出力は0.19まで縮小するため、このクエリに対してヘッドはほとんど機能しません。OpenAIのオープンウェイトモデルgpt-ossでは、各アテンションヘッドに学習可能なシンクロジットが組み込まれています。これにより、出力に寄与しない余剰の重みを各ヘッドが割り当てる場所を持てます。

StreamingLLMによるストリーミング生成#

無期限に生成を続けるモデルは、いずれコンテキストウィンドウを超え、すべてのトークンをKVキャッシュに保持できなくなります。直近のトークンだけを保持する通常のウィンドウアテンションでは、先頭のトークンが外れた時点で破綻します。StreamingLLMは、初期トークン4つのキーと値を、直近のトークンを保持するローリングウィンドウの隣に永続的に保持します。また、位置は元のテキスト内の位置ではなく、キャッシュ内の位置に基づいて割り当て、各ステップで回転位置埋め込み (RoPE)を再適用します。

この方法は、400万トークン以上にわたって言語モデリングを安定させ、新しいトークンごとにキャッシュを再計算するスライディングウィンドウより最大22.2倍高速に動作しました。ただし、ウィンドウによってモデルが想起できる範囲は制限されます。途中で破棄されたトークンは完全に失われるため、ストリーミングチャットボットやAIエージェントが記憶できるのは、直近のやり取りだけです。著者らはまた、各サンプルの先頭に、学習可能な専用のプレースホルダートークンを1つ置いてモデルを事前学習させました。このトークンが単独でシンクの役割を担いました。

アテンションシンクと関連用語の比較#

• スライディングウィンドウアテンション： 各トークンが参照できる範囲を直近の近傍に制限します。全アテンションで学習したモデルに推論時に適用すると、シンクがウィンドウから外れた時点で破綻します。
• コンテキスト劣化： プロンプトがコンテキストウィンドウ内に収まっていても、長くなるにつれて精度が低下する現象です。

ビジョンモデルにおけるシンク#

ビジョントランスフォーマー (ViTs)にも、関連する現象が見られます。「Vision Transformers Need Registers」で、MetaとInriaの研究者は、教師あり学習および自己教師あり ViTにおける情報量の少ない背景パッチ上に、内部計算へ転用されたノルムの異常に大きなトークンを発見しました。入力に少数の追加レジスタトークンを加えることで、その役割を吸収し、特徴マップとアテンションマップからアーティファクトを除去し、高密度予測の精度を向上させました。

Ultralyticsパッケージで実行できるMetaのSAM 2は、アンカーとウィンドウを保持する同様の方式で動画メモリを管理します。新しいフレームごとに、メモリアテンションは、動画全体を通して保持されるユーザープロンプト付きフレームのメモリバンクエントリーに加えて、直近6フレームのエントリーを読み取ります。そのため、動画の長さにかかわらず、フレームあたりのアテンション計算量は一定に保たれます。