Transformerアーキテクチャの数理的理解

序論：Transformerの革命的意義

Transformerアーキテクチャは、Vaswani et al.（2017）の論文「Attention Is All You Need」において提案され、深層学習と自然言語処理の両分野に根本的な変革をもたらした。リカレントニューラルネットワーク（RNN）や畳み込みニューラルネットワーク（CNN）に依存しない、注意機構のみに基づくこのアーキテクチャは、系列変換タスクにおいて卓越した性能を示すとともに、大規模並列計算との親和性により、数千億パラメータ規模のモデル訓練を実現可能にした。

本稿では、Transformerの各構成要素を数理的に詳細に分析し、その理論的性質を体系的に解説する。特に、自己注意機構の計算構造、位置エンコーディングの設計原理、多頭注意の表現力、層正規化の役割、そして計算複雑性の分析と効率化手法に焦点を当てる。読者は線形代数、確率論、および深層学習の基礎的知識を有することを前提とする。

Transformerの登場以降、このアーキテクチャはNLPのみならず、コンピュータビジョン（Vision Transformer; Dosovitskiy et al., 2021）、音声処理（Whisper; Radford et al., 2023）、タンパク質構造予測（AlphaFold 2; Jumper et al., 2021）、さらには強化学習（Decision Transformer; Chen et al., 2021）に至るまで、機械学習の広範な領域に浸透した。この普遍的適用性の数理的根拠を理解することは、現代のAI研究者にとって必須の素養である。

自己注意機構の数理的構造

自己注意機構（self-attention mechanism）は、Transformerの中核をなす演算であり、入力系列の各要素が他のすべての要素との関連度を計算し、重み付き和を出力する機構である。数理的には、以下のように定式化される。

入力系列 X ∈ ℝ^{n×d}（n: 系列長、d: 埋め込み次元）に対して、クエリ（Q）、キー（K）、バリュー（V）の三つの射影行列を用いて変換を行う：

Q = XW_Q , K = XW_K , V = XW_V

ここで W_Q, W_K ∈ ℝ^{d×d_k}、W_V ∈ ℝ^{d×d_v} は学習可能なパラメータ行列である。Scaled Dot-Product Attentionは以下で定義される：

Attention(Q, K, V) = softmax(QK⊤ / √d_k) V

スケーリング因子 1/√d_k の導入は、ドット積の分散を制御するために不可欠である。Q と K の各要素が平均0、分散1の独立な確率変数であると仮定すると、ドット積 q⊤k の分散は d_k となる。d_k が大きい場合、スケーリングなしではドット積の値が極端に大きくなり、softmax関数が飽和して勾配が消失する。√d_k による除算は、この分散を1に正規化する操作であり、学習の安定性を保証する。

注意行列 A = softmax(QK⊤ / √d_k) ∈ ℝ^{n×n} は、各行が確率分布を形成する二重確率的行列（doubly stochastic matrixではなく、行確率的行列）であり、入力系列の各位置が他の位置に配分する注意の重みを表す。この行列は、系列要素間の軟判定的な情報ルーティングを実現する。

情報理論的観点からは、注意行列の各行は、条件付きエントロピーの観点から解釈可能である。注意が一様分布に近い場合（高エントロピー）、各位置は系列全体から均等に情報を集約する。逆に、注意が特定の少数位置に集中する場合（低エントロピー）、選択的な情報抽出が行われる。Clark et al.（2019）の分析は、BERTの異なるヘッドが異なるエントロピー特性を示すことを明らかにした。

自己注意は、カーネル法との数学的接続を有する。Tsai et al.（2019）およびKatharopoulos et al.（2020）は、softmax注意をカーネル関数の特殊ケースとして解釈し、softmax(QK⊤) = φ(Q)φ(K)⊤ なる特徴写像 φ の存在を議論した。この観点は、線形注意（linear attention）の理論的基盤を提供し、計算複雑性の低減につながる。

多頭注意の表現力と解析

多頭注意（Multi-Head Attention; MHA）は、複数の独立した注意ヘッドを並列に適用し、それぞれが異なる部分空間で注意パターンを学習することを可能にする機構である。数理的には以下のように定義される：

MultiHead(Q, K, V) = Concat(head₁, …, head_h) W_O

where head_i = Attention(XW_Q^i, XW_K^i, XW_V^i)

ここで h はヘッド数、W_Q^i ∈ ℝ^{d×d_k}、W_K^i ∈ ℝ^{d×d_k}、W_V^i ∈ ℝ^{d×d_v}、W_O ∈ ℝ^{hd_v×d} は各ヘッドおよび出力射影の学習可能パラメータである。典型的には d_k = d_v = d/h と設定され、多頭注意全体の計算量は単頭注意と同等に保たれる。

多頭注意の表現力に関する理論的分析は、Transformer研究の重要な一分野を構成する。Voita et al.（2019）は、訓練後のBERTにおいて、多くのヘッドが冗長であり、剪定可能であることを実験的に示した。Michel et al.（2019）は、推論時に大部分のヘッドを除去しても性能低下が限定的であることを報告した。これらの知見は、多頭注意が表現力の冗長性を内包していることを示唆するが、訓練時の正則化効果やロバスト性への寄与は否定されていない。

Grouped Query Attention（GQA; Ainslie et al., 2023）は、キーとバリューを複数のヘッドで共有することにより、推論時のKVキャッシュのメモリ使用量を削減する手法である。Multi-Query Attention（MQA; Shazeer, 2019）はGQAの極端なケースであり、すべてのヘッドが単一のキー・バリューペアを共有する。LLaMA 2（Touvron et al., 2023）およびGemini（Google, 2024）等の最新モデルでGQAが採用されている。

注意ヘッドの役割分析は、Transformerの解釈可能性研究において中心的テーマである。構文的注意ヘッド（syntactic attention heads）は、特定の文法関係（主語-動詞一致、依存関係等）に対応する注意パターンを学習することが報告されている（Clark et al., 2019; Htut et al., 2019）。また、誘導ヘッド（induction heads; Olsson et al., 2022）は、文脈内学習（in-context learning）の実現に寄与する重要な計算パターンとして同定されている。

位置エンコーディングの設計原理

自己注意機構は本質的に順列不変（permutation invariant）であり、入力系列の順序情報を保持しない。したがって、位置エンコーディング（positional encoding）の導入が、系列の順序構造を捕捉するために不可欠となる。

元のTransformerでは、正弦関数に基づく固定的位置エンコーディングが採用された：

PE(pos, 2i) = sin(pos / 10000^{2i/d})

PE(pos, 2i+1) = cos(pos / 10000^{2i/d})

この設計の理論的根拠は、任意の固定オフセット k に対して、PE(pos+k) が PE(pos) の線形変換として表現可能であるという性質にある。すなわち、相対位置情報が線形演算によって抽出可能であり、注意機構がこれを学習できる。

学習可能な位置埋め込み（learnable positional embeddings）は、各位置に独立なパラメータベクトルを割り当てる手法であり、BERT、GPT系列で採用されている。この手法は柔軟であるが、訓練時の最大系列長を超える位置への外挿（extrapolation）が困難であるという制約がある。

Rotary Position Embedding（RoPE; Su et al., 2024）は、現在のLLMにおいて最も広く採用されている位置エンコーディング手法である。RoPEは、クエリとキーのベクトルに回転行列を適用することで、相対位置情報をドット積に直接エンコードする。具体的には、2次元の部分空間ごとに角度 θᵢ = pos × 10000^{-2i/d} の回転を適用する：

f(x, pos) = R(pos·θ) x

RoPEの優れた性質は、内積 f(q, m)⊤ f(k, n) が相対位置 m-n のみに依存することであり、これは相対位置エンコーディングの望ましい特性を自然に満たす。LLaMA（Touvron et al., 2023）、PaLM（Chowdhery et al., 2023）等で採用されている。

ALiBi（Attention with Linear Biases; Press et al., 2022）は、注意スコアに位置依存の線形バイアスを加算する手法であり、位置エンコーディングを注意計算に直接統合する。外挿性能において優れた特性を示し、BLOOM（BigScience, 2023）等で採用されている。

長系列への対応は、位置エンコーディング研究の重要な方向性である。YaRN（Peng et al., 2023）は、RoPEの周波数成分を操作することで、訓練時の文脈長を超えた推論を可能にする手法を提案した。NTK-aware scaling、Dynamic NTK等の手法が提案され、128K〜1Mトークン規模の文脈長をサポートするモデルが実現されている。

フィードフォワードネットワークとMLPの役割

Transformerの各層は、自己注意機構に加えて、位置単位のフィードフォワードネットワーク（Feed-Forward Network; FFN）を含む。標準的なFFNは二層のMLPとして定義される：

FFN(x) = W₂ · σ(W₁x + b₁) + b₂

ここで W₁ ∈ ℝ^{d×d_ff}、W₂ ∈ ℝ^{d_ff×d} であり、中間次元 d_ff は典型的には 4d に設定される。活性化関数 σ として、元のTransformerではReLUが用いられたが、現在のLLMではGELU（Gaussian Error Linear Unit; Hendrycks & Gimpel, 2016）またはSwiGLU（Shazeer, 2020）が標準的である。

SwiGLU活性化関数は以下で定義される：

SwiGLU(x) = (xW₁ ⊙ Swish(xW_gate)) W₂

ここで ⊙ は要素積、Swish(x) = x · sigmoid(βx) である。SwiGLUはGELUと比較して一貫した性能向上を示し（Shazeer, 2020）、LLaMA、PaLM、Mistral等の主要モデルで採用されている。ただし、追加のゲート射影行列 W_gate を必要とするため、パラメータ数と計算量が増加する。

FFNの役割に関する理論的理解は、近年大きく進展している。Geva et al.（2021）は、FFNの各ニューロンが特定の概念やパターンに対応するキー・バリューメモリとして機能するという解釈を提案した。すなわち、FFNの第一層がキー（入力パターンの検出）、第二層がバリュー（対応する出力の生成）として作用し、全体として連想記憶（associative memory）を実現する。この知見は、モデルの知識編集（knowledge editing; Meng et al., 2022）や事実の局在化（fact localization）研究の基盤となっている。

Mixture of Experts（MoE; Shazeer et al., 2017; Fedus et al., 2022）は、FFN層を複数の「エキスパート」に分割し、各入力に対してルーティング機構が少数のエキスパートのみを活性化する手法である。これにより、パラメータ数を増大させつつ計算量を一定に保つことが可能となり、Switch Transformer（Fedus et al., 2022）、Mixtral（Jiang et al., 2024）等のモデルで採用されている。

層正規化と残差接続の理論

Transformerの各サブレイヤー（自己注意およびFFN）は、残差接続（residual connection; He et al., 2016）と層正規化（Layer Normalization; Ba et al., 2016）を伴う。これらの構成要素は、深層Transformerの訓練安定性において決定的な役割を果たす。

層正規化は、各トークンの特徴ベクトルに対して、統計量に基づく正規化を行う：

LayerNorm(x) = γ ⊙ (x − μ) / √(σ² + ε) + β

ここで μ と σ² は特徴次元に沿った平均と分散、γ と β は学習可能なスケール・シフトパラメータ、ε は数値安定化のための小定数である。バッチ正規化（Ioffe & Szegedy, 2015）と異なり、層正規化はバッチサイズに依存せず、自己回帰生成時にも一貫して適用可能である。

Pre-LayerNorm（Pre-LN）とPost-LayerNorm（Post-LN）の配置は、訓練動特性に大きな影響を与える。元のTransformerのPost-LN配置（y = LayerNorm(x + SubLayer(x))）は、深層モデルにおいて学習率に対する敏感性が高く、訓練不安定の原因となることが報告されている（Xiong et al., 2020）。Pre-LN配置（y = x + SubLayer(LayerNorm(x))）は、残差経路の勾配をより安定的に伝播させ、訓練の安定性を改善する。ほとんどの現代LLMはPre-LNを採用している。

RMSNorm（Root Mean Square Layer Normalization; Zhang & Sennrich, 2019）は、層正規化の簡略化版であり、平均の除去（re-centering）を省略し、RMSによる正規化のみを行う：

RMSNorm(x) = γ ⊙ x / √(Σᵢ xᵢ² / d + ε)

RMSNormはLayerNormと同等以上の性能を示しつつ計算コストを低減するため、LLaMA、Mistral等で採用されている。

効率的注意機構と計算複雑性

標準的な自己注意の計算複雑性は O(n²d) であり、系列長 n に対して二次的にスケールする。この制約は、長系列処理における主要なボトルネックであり、効率化のための多様なアプローチが提案されている。

Sparse Attention（Child et al., 2019）は、注意行列の非ゼロ要素をストライドパターンや局所パターンに制限することで、計算量を O(n√n) に削減する。Longformer（Beltagy et al., 2020）は、局所注意とグローバルトークンへの注意を組み合わせた手法であり、O(n) の計算量を達成する。BigBird（Zaheer et al., 2020）は、局所注意、グローバル注意、およびランダム注意を統合し、疎注意パターンがチューリング完全であることを証明した。

Linear Attention（Katharopoulos et al., 2020）は、softmax注意をカーネル近似によって線形化する手法であり、計算量を O(nd²) に削減する。すなわち、系列長に対して線形にスケールするが、次元 d に対する依存性が二次的となる。Random Feature Attention（Peng et al., 2021）は、ランダム特徴写像を用いてsoftmaxカーネルを近似し、高品質な線形注意を実現する。

Flash Attention（Dao et al., 2022; Dao, 2023）は、アルゴリズムレベルの最適化により、注意計算のメモリアクセスパターンを改善する手法である。GPUのSRAMとHBM間のデータ転送を最小化するタイリング戦略により、注意行列 n×n を明示的に具現化（materialize）することなく、厳密な注意計算を O(n²d) の計算量で実行しつつ、メモリ使用量を O(n) に削減する。Flash Attention 2（Dao, 2023）は、さらなる並列化とワーク分割の最適化により、理論的FLOPS効率の70%以上を達成する。

Ring Attention（Liu et al., 2023）は、複数デバイス間でKVブロックをリング状に通信しながら注意を計算する手法であり、デバイス数に比例した系列長の延伸を可能にする。この手法はFlash Attentionと組み合わせることで、100万トークン規模の文脈長を実現する。

KVキャッシュと推論最適化

自己回帰生成における推論効率は、KVキャッシュ（Key-Value Cache）の管理に大きく依存する。自己回帰生成では、各ステップで新しいトークンを1つ生成し、過去のすべてのキー・バリューベクトルを再利用する。KVキャッシュのメモリ使用量は O(n × L × 2 × h × d_k)（n: 系列長、L: 層数、h: ヘッド数）であり、長系列やバッチ処理時に支配的なメモリ消費となる。

GQAおよびMQA（前述）は、KVヘッド数を削減することでキャッシュサイズを圧縮する構造的手法である。PagedAttention（Kwon et al., 2023）は、vLLMフレームワークにおいて実装された手法であり、KVキャッシュをOSのページングシステムに類似した方式で管理することで、メモリの断片化を解消し、バッチスループットを大幅に向上させる。

推論時の計算は、プレフィル（prefill）フェーズとデコード（decode）フェーズの二段階からなる。プレフィルフェーズでは入力プロンプト全体を並列処理してKVキャッシュを構築し、デコードフェーズでは逐次的にトークンを生成する。デコードフェーズはメモリ帯域幅律速（memory-bandwidth bound）であり、GPUの演算能力が十分に活用されない。Speculative Decoding（Leviathan et al., 2023）は、小規模ドラフトモデルによる候補生成と大規模モデルによる検証を組み合わせ、デコードの実効速度を向上させる手法である。

理論的展望と今後の方向性

Transformerの理論的理解は急速に深化しつつあるが、多くの根本的問題が未解決のまま残されている。表現力の観点では、Yun et al.（2020）がTransformerが任意の連続系列変換関数を近似可能であること（汎近似性）を証明した。しかし、有限の深さ・幅のTransformerがどの程度の関数クラスを効率的に表現可能であるかについては、より精密な特徴づけが必要である。

Transformerの計算能力に関しては、Pérez et al.（2021）がTransformerがチューリング完全であることを示したが、これは無限精度演算を前提とした理論的結果である。有限精度のTransformerの計算能力は、回路複雑性理論の枠組みで分析されており、Merrill & Sabharwal（2023）は、定数深さのTransformerがTC⁰回路クラスに対応することを示した。

State Space Models（SSM; Gu et al., 2022; Gu & Dao, 2023）は、Transformerの代替アーキテクチャとして注目を集めている。Mamba（Gu & Dao, 2023）は、選択的状態空間モデル（selective SSM）に基づき、系列長に対して線形の計算複雑性を達成しつつ、Transformerに匹敵する性能を示す。Jamba（Lieber et al., 2024）やZamba等のハイブリッドアーキテクチャは、TransformerとSSMの利点を統合する試みである。

これらの研究は、「注意機構は本当に必要か」という根本的問いに対する回答を模索しており、次世代の系列モデリングアーキテクチャの方向性を決定づけるものとなるだろう。Transformerの数理的理解は、新しいアーキテクチャの設計指針を提供するとともに、既存モデルの改善可能性を示唆する点で、引き続き重要な研究課題である。