マルチヘッド・アテンションは、深層学習におけるトランスフォーマーアーキテクチャの中核的な構成要素である。これは、クエリ、キー、バリューのそれぞれに学習された射影を用いて、複数のアテンション演算を並列に計算することで、基本的なアテンション機構を拡張する。これにより、モデルは入力シーケンス内の構文依存関係、共参照、位置的手がかりなど、さまざまな種類の関係やパターンを、複数の表現部分空間にわたって捉えることができる。これらの並列アテンションヘッドの出力は連結され、線形射影されて最終結果が生成され、モデルは異なる位置の異なる表現部分空間からの情報に共同で注意を向けることができる。
この概念は、Googleの研究者であるヤコブ・ウシュコレイト、ウカシュ・カイザー、ニキ・パルマーらによる2017年の論文「Attention Is All You Need」で導入された。この論文は、再帰や畳み込みを排し、アテンション機構のみに依存するトランスフォーマーを提案し、系列変換タスクにおける学習速度と性能の大幅な向上をもたらした。マルチヘッド・アテンションは以来、GPT、BERT、T5などの大規模言語モデルの基盤要素となり、自然言語処理、コンピュータビジョン、音声認識で広く使用されている。
メカニズム
アテンション機構では、各入力トークンはベクトル(しばしばトークン埋め込みと呼ばれる)として表現される。トークンのシーケンスに対して、モデルはクエリ(Q)、キー(K)、バリュー(V)の3つの行列を計算する。クエリとキー間のアテンションスコアは、通常、ドット積をキー次元の平方根でスケーリングして計算される。これらのスコアはソフトマックス関数で正規化され、アテンション重みが生成され、これを用いてバリューの加重和が計算される。
マルチヘッド・アテンションは、このプロセスを複数回並列に実行し、それぞれが元のクエリ、キー、バリューの異なる学習済み線形射影を用いる。例えば、8ヘッドの場合、モデルは8つの別々のアテンション出力を計算し、それぞれがシーケンスの異なる側面に焦点を当てる。出力は連結され、最終的な線形層を通過する。この設計により、モデルは異なる表現部分空間からの情報に注意を向けることができ、多様な言語パターンを捉えるのに特に有用である。
「ハード」な重みが学習中に設定されるのに対し、アテンション重みは「ソフト」であり、順伝播中に計算されるため、入力ごとに変化する。これにより、モデルは入力の関連部分に動的に焦点を当てることができる。機械翻訳におけるアライメントの例が示すように、「I love you」をフランス語に翻訳する際、モデルは「je」を生成するときに「I」に、「t'」を生成するときに「you」に、「aime」を生成するときに「love」に高いアテンション重みを割り当てる。
歴史
アテンション機構は、リカレントニューラルネットワーク(RNN)の弱点に対処するために開発された。RNNは文末の情報を優先し、初期の単語の重要性を減衰させる傾向がある。初期のアテンション機構、例えばBahdanauスタイルの加法的アテンション(2014年)やLuongスタイルの乗法的アテンション(2015年)は、機械翻訳のためのエンコーダー・デコーダーRNNアーキテクチャに組み込まれた。これにより、デコーダーは最終的な隠れ状態のみに依存するのではなく、すべてのエンコーダー隠れ状態に直接アクセスできるようになった。
大きな進展は自己アテンションによってもたらされた。これは、入力シーケンス内の各要素が他のすべての要素に注意を向けることで、モデルがグローバルな依存関係を捉えることを可能にする。このアイデアはトランスフォーマーの中心であり、再帰をアテンション機構に完全に置き換えた。トランスフォーマーの並列化可能なアーキテクチャは学習の大幅な高速化を可能にし、BERT、T5、生成的プレトレーニング済みトランスフォーマー(GPT)などのモデルでの採用につながった。
変種
アテンションにはいくつかの変種が開発されており、多くはソフトな重みを実装している。これらには以下が含まれる:
- 高速重みプログラマー(1992年):ある「遅い」ニューラルネットワークが、外積を通じて別のニューラルネットワークの「高速」な重みを出力する。これは後に「線形化自己アテンション」と改名された。
- Bahdanauスタイルのアテンション(加法的アテンション):フィードフォワードネットワークを用いてアライメントスコアを計算する。
- Luongスタイルのアテンション(乗法的アテンション):デコーダーとエンコーダーの状態間のドット積を用いる。
- 位置的アテンションおよび因子分解位置的アテンション:位置情報をアテンション計算に組み込む。
- 畳み込みニューラルネットワークのための空間的およびチャネル的アテンション:空間次元または特徴チャネルに作用する。
これらの変種は、エンコーダー側の入力を再結合して各ターゲット出力への効果を再分配し、しばしばドット積の相関行列を用いて再重み付け係数を計算する。
最適化
アテンション行列のサイズは入力トークン数の2乗に比例するため、長いシーケンスではメモリを大量に消費する可能性がある。Flash Attentionは、メモリ要件を削減し、精度を犠牲にすることなく効率を高める実装である。アテンション計算をGPUの高速なオンチップメモリに収まる小さなブロックに分割し、大きな中間行列を格納する必要性を減らす。
Metaによって開発されたFlexAttentionは、ユーザーがソフトマックス前にアテンションスコアを変更できるようにし、最適なアテンションアルゴリズムを動的に選択するアテンションカーネルであり、カスタムアテンションパターンに柔軟性を提供する。
応用
アテンションは、自然言語処理、コンピュータビジョン、音声認識で広く使用されている。NLPでは、質問応答や要約などのタスクで文脈理解を向上させる。ビジョンでは、視覚的アテンションがモデルが関連する画像領域に焦点を当てるのを助け、物体検出や画像キャプションを強化する。
コンピュータビジョンでは、ビジョントランスフォーマー(ViT)が画像パッチにマルチヘッド・アテンションを適用する。アテンションマップは、アテンションスコアをヒートマップとして可視化するもので、ViTモデルの意思決定プロセスを検査する日常的な方法となっている。深い層ほど、より意味的に有意義な可視化を示す傾向がある。アテンションロールアウトは、連続するアテンションマップのドット積を計算することで、全層のアテンションスコアを組み合わせる再帰的アルゴリズムである。
ビジョントランスフォーマーは通常、自己教師あり方式で学習されるため、アテンションマップは一般にクラスに敏感ではない。分類ヘッドが取り付けられると、クラス判別アテンションマップ(CDAM)が、アテンションマップとクラストークンに関する勾配を組み合わせて、クラス固有の説明を提供する。
影響
マルチヘッド・アテンションは、生成AIと大規模言語モデルの台頭に大きく貢献してきた。これにより、モデルは長いシーケンスを効率的に処理し、複雑な依存関係を捉えることができ、OpenAI、Anthropic、Google DeepMindなどの組織が開発するシステムの重要な構成要素となっている。このアーキテクチャはハードウェア設計にも影響を与えており、NVIDIAやCerebrasなどの企業がアテンションベースの計算に最適化されたチップを開発している。
アテンションの効率性と解釈可能性の向上に関する研究は続いている。スパースアテンション、線形アテンション、カーネルベースの手法などの技術は、アテンションの2次複雑性を低減することを目指しており、解釈可能性ツールはモデルが何を学習するかを理解するのに役立つ。マルチヘッド・アテンションは、機械学習と人工知能における活発な研究分野であり続けている。