1 基本原理

1.1 注意機構の概念

注意機構は、深層学習モデルが入力情報のうち重要な部分に焦点を当てる仕組みである。従来のエンコーダ・デコーダモデルでは、固定長の文脈ベクトルが情報のボトルネックとなる問題があったが、注意機構は各タイムステップでエンコーダの全隠れ状態を参照し、重み付き和を計算することで柔軟な情報利用を可能にする。これにより、モデルはタスクに応じて入力の各部分への注意配分を学習できる。

1.2 自己注意の定義

自己注意機構は、同一シーケンス内の要素間で注意を計算する手法である。従来の注意機構が異なるシーケンス間(例:エンコーダとデコーダ)で用いられるのに対し、自己注意は一つのシーケンス内で各位置が他の全位置との関係を直接モデル化する。これにより、シーケンスの長距離依存関係を並列に捉えることができ、再帰的処理に伴う逐次的な計算制約から解放される。

1.3 クエリ・キー・バリューの仕組み

自己注意では、入力ベクトルを3つの異なる線形変換によってクエリキーバリューに射影する。クエリは注目する側の要素、キーは注目される側の要素の特徴を表し、バリューは実際の情報内容を保持する。クエリとキーの内積によって類似度を計算し、そのスコアソフトマックス関数正規化することで注意重みを得る。最後に、この重みでバリューを加重和し、出力を生成する。この機構により、各要素は他の全要素と動的に関連付けられる。

2 数学的定式化

2.1 スケーリング内積注意

最も広く使われる自己注意の計算方法は、スケーリング内積注意である。入力が行列\(Q\)(クエリ)、\(K\)(キー)、\(V\)(バリュー)で与えられるとき、出力は次のように定義される: \[ \text{Attention}(Q, K, V) = \text{softmax}\left( \frac{QK^\top}{\sqrt{d_k}} \right) V \] ここで\(d_k\)はキーの次元数であり、スケーリング係数\(\sqrt{d_k}\)は内積値が大きくなりすぎてソフトマックス関数の勾配が消失するのを防ぐ役割を果たす。この計算は行列演算で効率的に実装でき、GPUでの並列化に適している。

2.2 ソフトマックス関数と重み付け

ソフトマックス関数は、注意スコア(内積値)を確率分布に変換する。各行に対して適用され、各位置のキーに対する注意重みが非負かつ合計1になるように正規化される。この正規化された重み行列をバリュー行列に掛け合わせることで、各クエリ位置に対応する加重和が計算され、出力シーケンスが得られる。重みの分布はモデルがどの入力要素に注目すべきかを示し、解釈可能性の手がかりとなる。

2.3 マルチヘッド自己注意

2.3.1 ヘッドの分割と結合

マルチヘッド自己注意は、単一の注意機構を複数並列に実行する拡張である。クエリ、キー、バリューをそれぞれ\(h\)個の低次元部分に分割(ヘッド)し、各ヘッドで独立にスケーリング内積注意を計算する。全てのヘッドの出力を連結した後、再度線形変換することで最終的な出力を得る。これにより、モデルは異なる位置の異なる表現部分空間に同時に注目できるようになる。

2.3.2 パラメータ効率

マルチヘッド注意では、ヘッド数\(h\)と各ヘッドの次元\(d_v\)の積が元の次元\(d_{\text{model}}\)と一致するように設定されることが多い。つまり、総パラメータ数は単一ヘッドの場合とほぼ変わらず、パラメータ効率が高い。また、ヘッドごとに異なる注意パターンを学習できるため、表現力が向上する。実際のTransformerモデルでは、\(h=8\)や\(h=16\)などが標準的に用いられる。

3 主要なバリエーション

3.1 位置エンコーディング

自己注意機構は、入力シーケンスの順序情報を自然に持たない(並列処理のため位置の概念がない)。そこで、位置情報を埋め込むための位置エンコーディングが必須となる。

3.1.1 絶対位エンコーディング

絶対位置エンコーディングは、シーケンス内の各絶対位置に対して一意なベクトルを生成し、入力埋め込みに加算する。代表的なものに、Transformer原論文で用いられた正弦波・余弦波による固定エンコーディングがある。これは学習可能なパラメータを持たず、任意の長さのシーケンスに対応可能だが、絶対位置のみを表現するため、相対的な位置関係を直接モデル化しない。

3.1.2 相対位置エンコーディング

相対位置エンコーディングは、クエリとキーの間の相対的な距離や方向に基づいて注意スコアを調整する。例えば、Transformer-XLやT5、DeBERTaなどで採用されており、絶対位置よりも一般化性能が高いとされる。具体的には、注意重みの計算時にバイアス項として相対位置情報を加える方法や、キーの埋め込みに相対位置表現を組み込む方法がある。

3.2 効率的自己注意

標準の自己注意はシーケンス長\(n\)に対して\(O(n^2)\)の計算量とメモリを必要とするため、長いシーケンスを扱う際のボトルネックとなる。そこで、様々な効率的なバリエーションが提案されている。

3.2.1 スパース注意

スパース注意は、全ペアではなく一部の要素間のみに注意を限定する手法である。例えば、固定の間隔(固定ストライド)や局所的な窓(局所注意)、あるいは学習された注意パターンに基づいてスパースな結合を生成する。代表的なものに、Longformerのスライディングウィンドウ注意+グローバル注意、BigBirdのランダム注意+局所注意+グローバル注意の組み合わせがある。これにより計算量を\(O(n \sqrt{n})\)や\(O(n \log n)\)に削減できる。

3.2.2 線形注意

線形注意は、ソフトマックス関数をカーネル関数で近似し、注意行列の計算を線形時間で行う手法である。代表的なものに、Linformer(低ランク近似)、Performer(正準線形注意、FAVOR+)、Reformer(LSH注意:局所性鋭敏ハッシュを用いた近似)などがある。これらは注意行列を陽に計算せず、クエリとキーの積の順序を入れ替えることで\(O(n)\)の計算量を達成する。ただし、精度が若干低下する場合がある。

4 応用分野

4.1 自然言語処理

4.1.1 BERTとGPT

BERTはTransformerのエンコーダ部を積み重ね、双方向の自己注意を用いてマスク言語モデリングと次文予測で事前学習を行う。GPTはデコーダ部のみを積み重ね、自己回帰的な生成に特化する。両者ともマルチヘッド自己注意を中核とし、文脈を深く理解した表現を獲得し、様々なNLPタスクで高い性能を示した。

4.1.2 機械翻訳

自己注意機構はシーケンス・ツー・シーケンス翻訳において、エンコーダ内での文脈理解とデコーダ内での生成過程の両方で長距離依存を捉える。従来のRNNベースの翻訳モデルを大きく凌駕し、Google翻訳など商用システムにも採用された。特に、入力文と出力文のアライメントを注意重みとして可視化できる点が利点である。

4.2 コンピュータビジョン

4.2.1 Vision Transformer

ViT(Vision Transformer)は、画像をパッチに分割し、各パッチを単語のように扱ってTransformerエンコーダに入力する。標準の自己注意機構により、画像全体のグローバルな関係を捉えることができる。CNNのような局所受容野に頼らず、大規模データでの学習によりCNNを超える性能を達成した。その後、DeiT、Swin Transformerなどの改良版が登場している。

4.2.2 画像生成

自己注意は画像生成モデルにも導入されている。例えば、GANのジェネレータに自己注意層を追加することで、テクスチャや全局構造の一貫性が向上することが示された(Self-Attention GAN)。また、拡散モデルにおいても、U-Net内のアテンション機構が高品質な画像合成に寄与している。

4.3 その他の領域

自己注意機構は、音声処理(音声認識、音声合成)、マルチモーダル学習(画像キャプション、ビデオ理解)、科学的計算(タンパク質構造予測AlphaFold2)など多岐にわたる分野で応用されている。特にAlphaFold2では、Evoformerと呼ばれるアーキテクチャで自己注意を用いて進化的情報と残基間の相互作用をモデル化し、大きな成果を挙げた。

5 利点と限界

5.1 並列計算と長距離依存

自己注意の最大の利点は、シーケンス全体の並列計算が可能な点である。RNNのような逐次処理が必要ないため、GPUでの高速な学習・推論が実現する。また、任意の距離にある要素間の直接的な相互作用を1層でモデル化できるため、長距離依存関係の学習が容易である。例えば、文の先頭と末尾の単語が直接関係する場合でも、注意重みが大きくなるように学習できる。

5.2 計算量とメモリの課題

標準の自己注意はシーケンス長\(n\)に対して\(O(n^2)\)の計算量・メモリを要する。そのため、非常に長いシーケンス(例:長文、高解像度画像、長い動画)を扱う際には現実的でない。また、各注意ヘッドのパラメータ数も多いため、モデル全体のパラメータ数が膨大になりがちである。この問題に対処するために、効率的自己注意のバリエーションが提案されている。

5.3 解釈可能性

自己注意の注意重みは、可視化することでモデルが入力のどの部分に着目しているかを示す手がかりとなる。しかし、注意重みが必ずしも因果関係や意味的な重要度と一致するとは限らない。また、モデルが複雑になるほど注意パターンを解釈するのは難しい。近年、注意重みだけでなく、勾配や摂動を用いた説明手法が研究されている。

6 実装と最適化

6.1 フレームワークでの実装例

主要な深層学習フレームワーク(PyTorch、TensorFlow、JAX)では、自己注意機構の実装が標準ライブラリに用意されている。例えば、PyTorchのnn.MultiheadAttentionモジュールはクエリ、キー、バリューの線形変換とマルチヘッド注意計算をカプセル化している。Transformerモデルを構築する際には、これらの高レベルAPIを用いることで容易に実装できる。また、Hugging FaceのTransformersライブラリでは、BERTやGPTなど事前学習済みモデルが豊富に提供されている。

6.2 高速化手法

自己注意の計算を高速化・省メモリ化するための手法が多数開発されている。代表的なものとして、混合精度学習(FP16/BF16)、FlashAttention(IO意識型注意計算によるメモリアクセス最適化)、さらにはカスタムカーネル(CUDAによるfused attention)などがある。FlashAttentionは、従来の注意計算がメモリバウンドである問題を解決し、注意行列を明示的に格納せずに計算を行うことで、大幅な速度向上とメモリ削減を実現した。また、効率的自己注意のバリエーションを専用カーネルで実装することで、長いシーケンスでも実用的な速度を達成できる。