松尾研究所テックブログ
🧠

Attention再入門 is all you need

Taisei Ozaki
に公開
機械学習
transformer
Attention
llm
tech

Attention再入門 is all you need

松尾研究所の尾崎です．25卒でデータサイエンティストをやっています．

Attention機構は，"Attention is all you need"論文で一気に脚光を浴びて以来，昨今のAIブームを支えるLLM(transformer)の根幹的技術です．今回はそんなAttention機構が登場以来，どういう方向で進化してきたのかを整理して，皆さんの興味を発掘したり，普段何気なく使っている技術の裏側を学ぶきっかけにしたり，していただきたいと思い，本記事執筆に至っています．

※本記事は社内で行った勉強会からの抜粋となっております．

1. Attention進化のマトリクス（俯瞰図）

現在のLLMの進化は、上記マトリクスの「3つの対象」と「2つのアプローチ」の掛け合わせで整理できると考えています．

では本題に入る前に，まずはいくつかのキーワードを復習しましょう．

そもそもattentionとは？

Attention機構とは，一言で言えば「入力のどこに注目するかを学習する仕組み」です．これが昨今のLLM分野の発展の始祖になっています．

Attention機構については解説されている記事がたくさんあります．ので，ここではあまり深堀ませんが，簡単にキーワードだけ．

従来のRNNやLSTMは，入力トークン（単語や文字）のシーケンスを順番に処理するため，長い文脈の依存関係を捉えるのが苦手でした．Attentionはこの問題を解決します．入力系列の各要素に対して「どの程度注目するか」を重み付けし，重要な情報を選択的に抽出できるのです．

具体的には，まず以下の3つのベクトル Query（Q），Key（K），Value（V） を入力トークンから作ります：

  1. Query: 各入力トークンが「何を探しているか」を表すベクトル
  2. Key: 各要素の「ラベル」のようなベクトル
  3. Value: 実際に取り出したい「情報」を表すベクトル

その後QueryとKeyの類似度（内積）を計算し，softmaxで正規化することで「注目度(Attention weight)」を得ます．この注目度でValueを重み付けして合計したものが出力です．数式で表すと以下のようになります．

\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V

\sqrt{d_k} で割っているのは，次元が大きいと内積の値が極端に大きくなり，softmaxの勾配が消失するのを防ぐためです．展開して書いていくとセクション冒頭のスライドの末尾にあるような形になります．

multi-head attentionとは？

言語には，文法的な関係，意味的な関係，代名詞の参照など，複数の関係性が同時に存在します．単一のAttention（Q, K, V）だと単語間の特定の関係性しか表現できませんが，QKVのパターンを複数用意すれば異なる種類の関係性を同時に捉えられます．これがMulti-head attentionで，文字通り複数のAttentionを並列に実行します．

Multi-Head Attention = Concat(head_1, head_2, ..., head_h) × W^O

各ヘッドは異なる「視点」を持ちます．例えば：

  • あるヘッドは主語と動詞の関係に注目
  • 別のヘッドは形容詞と修飾される名詞の関係に注目
  • また別のヘッドは代名詞とその参照先の関係に注目

実装上は，Q/K/Vをヘッド数で分割し，それぞれで独立にAttentionを計算した後，結果を連結します．GPT-4クラスのモデルでは，96〜128ものヘッドが使われていると言われています．

Attentionの計算コストが膨大な理由とは？

Attentionの計算量は O(n^2 \cdot d) です．ここで n は系列長，d は埋め込み次元です．

すべてのトークンペアの関係性を計算するため，トークン長の2乗の組み合わせの分だけ計算量は増えます．つまり，n^2です．

例えば，4,000トークンの文書を処理する場合：

  • 計算するペアの数 = 4,000 \times 4,000 = 16,000,000（1,600万ペア）

100,000トークンに拡張すると：

  • 計算するペアの数 = 100,000 \times 100,000 = 10,000,000,000（100億ペア）

Attention行列（n \times n）を明示的に保持する必要があるので，メモリ使用量も同様に O(n^2) でスケールします．1Millionトークンの場合，メモリ使用量は2TBになるそう．．．

KVcacheとは？

LLMの推論（テキスト生成）は 自己回帰的 に行われます．つまり，1トークンずつ順番に生成し，生成したトークンを入力に追加して次のトークンを予測します．

ここで問題が生じます．新しいトークンを生成するたびに，過去のすべてのトークンに対するKey（K）とValue（V）を再計算するのは無駄です．

KVキャッシュ は，この再計算を避けるために，過去のトークンのKとVをメモリに保持しておく仕組みです．

ステップ1: "春" → KV計算 → キャッシュに保存
ステップ2: "春　は" → "は"だけKV計算，キャッシュと結合
ステップ3: "春　は　曙" → "曙"だけKV計算，キャッシュと結合

これにより，各ステップの計算量が O(n^2) から O(n) に削減されます．

ただし，KVキャッシュはメモリを消費します．モデルが大きく，系列長が長いほど，キャッシュサイズは膨大になります．この「KVキャッシュのメモリ問題」が，後述するGQAやPagedAttentionといった最適化技術の動機となっています．

2. 計算量の削減

前述の通り，Attention機構には計算量の爆発という功罪があります．しかし実際のボトルネックは単純な演算回数だけではなく，GPUメモリ内のデータ移動（IO） にもあります．

なぜIOがボトルネックになるのか？

GPUのメモリには階層構造があります．計算コア（テンソルコア）に近いほど高速ですが容量が小さくなります．

  • SRAM（レジスタやL1/L2キャッシュ）: 計算コアに直結した超高速メモリ．実際の演算はここで行われる
  • HBM（High Bandwidth Memory）: 大容量だが遅いメインメモリ．データの保管場所

標準的なAttention実装では，巨大な n \times n のAttention行列を一度HBMに書き出し，softmaxを計算し，再びHBMから読み込んで出力を計算します．この「HBMへの往復」が，H100のような最新GPUでさえ計算ユニットが計算をしていない時間を増やしてしまう原因になります．データを運ぶのが計算速度に追いつかない 状態になってしまうわけです．

FlashAttention

この問題に対する根本的な解法が FlashAttention です．

FlashAttentionの核心は「タイリング（tiling）」と呼ばれるテクニックです．巨大なAttention行列を一度に作るのではなく，SRAMに収まる小さなブロック単位で計算を完結させます．具体的には：

  1. Query, Key, Valueをブロック単位で読み込む - SRAMに乗るサイズに分割
  2. ブロック内でAttentionを計算 - 中間結果はSRAM上で保持
  3. オンラインsoftmaxで正規化を「流れ作業」で実行 - 通常のsoftmaxは全要素を見て最大値や合計を求める必要がありますが，オンラインsoftmaxはブロックごとに正規化係数を更新しながら最終的に同じ結果を得る手法です

FlashAttention 3 (FA3) では，H100 GPUのハードウェア特性をさらに深く活用しています．H100には TMA（Tensor Memory Accelerator） という専用ユニットがあり，メモリ転送とテンソル演算を非同期に並列実行できます．FA3はこれを活かし，次のブロックのデータ転送を行いながら現在のブロックの計算を進める「パイプライン化」を実現しました．結果として，H100上でAttentionの実効スループットを理論値の75%以上まで引き上げています．

ただし注意が必要なのは，FlashAttentionは計算量自体は O(n^2) のままという点です．あくまでメモリ効率の最適化であり，計算オーダーを変えているわけではありません．

O(n^2) からの脱却

もう一つのアプローチは，そもそも O(n^2) の計算を回避する構造への移行です．大きく分けて Sparse Attention と Linear Attention の2つの方向性があります．

Sparse Attention：注目すべき場所を絞る

Attention行列は実際はほぼゼロの部分，つまり当該トークンとあまり関係のないトークンもかなりあります．それは無視してもいいのではないかと考えられるわけです．

Sparse Attention は，すべてのペアを計算するのではなく，重要そうな箇所だけを選択的に計算します．

代表的な手法としては以下があります：

  • Longformer: 各トークンは近傍のトークン（局所的なウィンドウ）と，特定の「グローバルトークン」（[CLS]など）にのみAttentionを向けます．これにより計算量は O(n) に削減されます．

  • BigBird: Longformerのアイデアに加え，ランダムに選んだトークンへのAttentionも追加します．局所性 + グローバル性 + ランダム性の組み合わせで，理論的にもチューリング完全性を保証しています．

  • Native Sparse Attention (NSA): 最近提案された手法で，学習によってスパースパターンを動的に決定します．ハードコードされたパターンではなく，入力に応じて「どこを見るべきか」を学習するのが特徴です．

Sparse Attentionの計算量は設計によりますが，典型的には O(n \log n) や O(n \sqrt{n}) に抑えられます．

Linear Attention：カーネルトリックで計算順序を変える

Linear Attention は，より数学的なアプローチで O(n) を実現します．

ポイントは計算順序です．標準Attentionでは (QK^T) を先に計算するため n \times n の行列が生じますが，Linear Attentionでは (\phi(K)^T V) を先に計算することで，d \times d の行列で済みます（d は埋め込み次元）．系列長 n が大きくなっても，計算量は O(n \cdot d^2) でトークン長に対して線形にスケールします．

代表的な手法：

  • Performer: softmaxの指数関数部分を，ランダムに初期化した三角関数（sin, cos）の組み合わせで近似します．これにより \phi(Q) と \phi(K) を独立に計算でき，上記の計算順序の変更が可能になります．
  • Linear Transformer: ReLU や ELU などのシンプルな活性化関数を \phi として使用します．

ただし，Linear Attentionには精度面での課題もあるようです．softmax の鋭いピークを作る能力が失われるため，特定のタスクでは性能が落ちることがあります．

中間まとめ①

両方重要ではありますが，FlashAttentionは，LLMの学習や推論において非常に重要なツールの一つになっています．知っておいて損はないかなと思います．モデルの精度を一切落とさずに2-4倍の高速化が得られるため，ほぼデファクトスタンダードとなっている印象です．

一方でインストール周りやバージョン選択で他のツールとコンフリクトすることも多々あるのが玉に瑕です．当然GPUやcudaのバージョンともかなり密接な関係です．その点は調べて進めていくしかないですね．逆にいうと，どのエラーがわからないエラーのときに，もしかして"flash-attn周りじゃね？"と思って調べていると当たりだったりみたいなのはよくあります．

この辺はGPUはじめチップ関連の分野にLLMの側からつながる入口です．アルゴリズムだけでなく，ハードウェアの知識もちゃんとわかるサイエンティストになりたいですね．キャッチアップを続けたいです．

3. KVキャッシュの処理

Attention機構を持つtransformerモデルの推論時の最大の壁は「メモリ容量」です．どれだけ計算が速くなっても，KVキャッシュがGPUメモリに乗り切らなければ，生成を続けることはできません．

この壁に対し，現在は データの持ち方（共有・圧縮） と 管理方法（動的割当） の両面から攻めるのが主流となっています．

なぜKVキャッシュがボトルネックになるのか？

冒頭説明したように，過去に計算したKeyとValueをキャッシュしておくことで，重複計算を避けることができます．しかし，このキャッシュのサイズは馬鹿になりません．

例えば，70B級のモデルを考えてみましょう．

  • 層数: 80層
  • KVヘッド数: 8
  • ヘッド次元: 128
  • シーケンス長: 4096トークン
  • バッチサイズ: 32

このとき，KVキャッシュに必要なメモリは：

80 \times 8 \times 2 \times 128 \times 4096 \times 32 \times 2 \text{bytes} \approx 42 \text{GB}

これはモデルの重みとは別に必要です．しかも，シーケンス長やバッチサイズに比例して膨らみ続けます．100Kトークンの長文対応などしようものなら，キャッシュだけで数百GBに達してしまいます．

アプローチ1：データ量の削減（GQA / MQA）

まずはキャッシュするデータ量そのものを減らすアプローチです．

Multi-Query Attention (MQA)

2019年にGoogleが提案したMQAは，全てのQuery headで1つのK/Vペアを共有するという大胆な手法です．Attentionの本質は「QueryがどのKeyに注目するか」を学習することなので，異なるQueryが同じKey-Valueプールを参照しても，それぞれのQueryは異なる注目パターンを学習できます．これにより，KVキャッシュのサイズを \frac{1}{h}（hはhead数）に削減できます．

ただし，当然ながらこれだけパラメータを削れば精度低下のリスクがあります．Key-Valueが共通化されることで，各headが独自の「視点」を持つ能力が制限されてしまいます．

Grouped Query Attention (GQA)

GQAはMQAとMHA（標準のMulti-Head Attention）の中間を取った手法です．全headで1つではなく，グループごとにKVを共有します．例えば32個のQuery headを8グループに分け，4 Query headごとに1つのKVペアを共有します．

これによりKVキャッシュを\frac{1}{groups}に削減しつつ，MHAに近い表現力を維持できるようになりました．現在のLlama，Mistralなどのアーキテクチャが明らかになっている主要モデルはほぼ全てGQAを採用しています．

Multi-head Latent Attention (MLA)

MQAやGQAが「KVヘッドを共有する」アプローチだったのに対し，MLA（Multi-head Latent Attention） は異なる発想でKVキャッシュを削減します．低ランク圧縮です．

MLAはDeepseekで導入されたことがよく知られていますが，KeyとValueを低次元の潜在空間に圧縮してからキャッシュすることが本質です．推論時には，この圧縮された潜在ベクトルを元のサイズに復元してAttention計算を行います．

数学的には，従来の W_{KV} 行列を低ランク行列分解 W_{DKV} \cdot W_{UKV} に置き換えています．W_{DKV} で一度低次元化し，キャッシュするのはこの圧縮された潜在ベクトル c_{KV} のみにします．推論時に W_{UKV} で元の次元に戻してK, Vを復元します．

MLAの利点は，GQAやMQAのようにヘッドを共有するのではなく，各Queryヘッドが独自のK, Vを持てる，つまりモデルパラメータを削っているわけではないので，学習によって情報保存できるパラメータは実質的にMHAと同じという点です．

またMLAはMHAをわずかに上回る性能を示すことも報告されています[1]．低次元圧縮によりノイズが減り，精度に寄与している可能性があります．さらにMHA2MLAのような変換ツールも登場しており，手法としての注目度も高まっていますね．

アプローチ2：メモリ管理の最適化（PagedAttention）

次はメモリの断片化です．

従来transformersなどの推論エンジンでは，各リクエストに対して「最大シーケンス長分のメモリ」を事前に確保していました．例えば最大4096トークンの設定なら，実際に100トークンしか生成しなくても4096トークン分のメモリを占有し続けます．

この無駄を解消するのがPagedAttentionです．vllmでお馴染みですね．

OSの仮想メモリから着想を得たアプローチ

PagedAttentionは，OSのページング機構をKVキャッシュに適用したものとコンセプトベースでイメージできます．

  1. KVキャッシュを固定サイズの「ブロック」に分割
  2. 各ブロックを論理アドレスと物理アドレスのマッピングで管理
  3. 必要に応じて動的にブロックを割り当て・解放

これによりメモリ使用効率が劇的に向上し，実際に使っている分だけメモリを消費することになります．またvllm serverがわかりやすいですが，同じGPUでより多くのリクエストを処理可能になる強力な技術です．

KV Cache Quantization

MLAは次元圧縮でしたが，量子化でもいいんではないかという議論も当然ありまして，KVキャッシュをFP16からINT8やINT4に量子化することで，さらにメモリを削減する推論手法もvllmやOllamaなどはサポートしています．モデル本体の量子化と組み合わせることで，コンシューマーレベルGPUのかなり小さなメモリで大規模モデルを動かせるようになります．

中間まとめ②

KVキャッシュの最適化は，2025年~2026年のLLM開発では必須の技術となっています．特にGQA/MLAあたりは，もはや標準技術になりつつありますね．vllmのPagedAttentionも，推論ライブラリとしての確固たるものを築きつつあります．それだけこの分野が重要であり，かつ技術的なチャレンジが多いことの裏返しでもあります．当然キャッチアップを続けたいですね．

4. 長文脈対応

長文といえば最初に1Millionトークンのサポートを可能にしたGeminiが連想されますね．ただ1Millionトークンを超える「超長文」を扱うには，数式上の工夫と，メモリ拡張のためにマルチGPU/マルチノードを繋ぐ通信の工夫がセットで必要です．長文脈対応は，昨今のTest Time Scalingを支える根幹技術であり，o1やGeminiなどが「長く考える」ことを可能にしている重要なピースでもあります．

なぜ長文脈が難しいのか？

何度も言って恐縮ですが，Attentionの計算量はO(n^2)でスケールするため，トークン数nが大きくなると急激に計算量が爆発します．これに加えて，位置埋め込み（Positional Encoding）の問題があります．単純なAttentionは文章内の位置情報を追えませんし，位置情報エンコーディングをしたとしても物理的に離れた場所にある各トークンとの関係性の違いを表現することは難しいです．また学習時に見たことのない長さのシーケンスに対して，位置情報をどう表現するか？というのも課題です．

長文脈対応には，上記2つの課題それぞれに対するアプローチがあります．

アプローチ1：計算量・メモリへの対応（並列化とメモリ分散）

単一GPUのメモリには限界があります．100Kトークンのシーケンスを扱おうとすると，KVキャッシュだけで数十GBに達することも．これを解決するのがシーケンス並列化（Sequence Parallelism）の技術です．

Ring Attention

Ring Attentionは，巨大なシーケンスをGPUごとに分割し，輪（リング）状にデータを回しながら計算する手法です．

GPU 0: [トークン 0-1000]    →  計算しながら次のGPUへKVを送信
GPU 1: [トークン 1001-2000] →  GPU 0からKVを受信しつつ計算
GPU 2: [トークン 2001-3000] →  GPU 1からKVを受信しつつ計算
...（リング状に循環）

ポイントは，通信と計算のオーバーラップです．あるGPUが計算している間に，別のGPUにデータを送信することで，通信のオーバーヘッドを隠蔽します．これにより，理論上はGPUの数に比例してコンテキスト長を伸ばせるようになります．

Striped Attention

Ring Attentionをさらに改良した手法がStriped Attentionです．シーケンスをGPU間で「ストライプ状」に分散させることで，因果マスク（causal mask）による負荷の不均衡を解消し，より効率的な並列処理を実現します．

アプローチ2：位置埋め込みへの対応（位置エンコーディングの拡張）

RoPE（Rotary Position Embedding）とその進化

現代のLLMで広く使われているRoPEは，トークンの相対位置を回転行列（sinやcosの成分）で表現する手法です．ただし，学習時のコンテキスト長を超えると何が起こってしまうかというと，たとえば4kトークンで学習した場合，4kトークン先が一番遠いトークンとして学習されるわけですが，このときの角度を \theta とした場合，もし5kトークンが入力されるとそのトークンに付される角度はナイーヴに考えると \frac{5}{4}\theta となるわけですが，モデルは \frac{5}{4}\theta など見たことないので，位置情報が正しく表現できなくなって性能が劣化します．

  • LongRoPE / LongRoPE2: 学習範囲外の長さでも位置情報を失わないよう，周波数成分を動的にリスケールします．具体的には，低周波成分（長距離の相対位置を担当）と高周波成分（近距離の相対位置を担当）を異なる比率でスケーリングし，学習時と推論時のギャップを埋めます．

  • YaRN (Yet another RoPE extensioN): RoPEの周波数をNTK-aware（Neural Tangent Kernel を考慮した）方法で補間し，ファインチューニングなしで文脈長を拡張可能にします．NTK-awareとは，高周波成分と低周波成分で異なる補間戦略を取ることで，単純な線形補間よりも情報の損失を抑える手法です．Llama 2をベースに，4Kから128Kへの拡張で良好な結果を示しています．

YaRNは追加学習なしで既存モデルの文脈長を拡張できるという点で非常に重要です．o1やDeepSeek-R1のようなReasoningモデルは，推論時に長い思考過程（Chain of Thought）を生成するため，事前学習時よりも長いコンテキストを扱う必要がありますし，「長くしゃべることで性能があがる」Test time scalingの成立のためにも前提的に32k以上の長いトークン列を取り扱うことができるようにならないといけません．YaRNのような技術があることで，学習時4Kトークンで訓練されたモデルでも，推論時に数十Kトークンの思考を学習無しで展開できるようになります．

より詳しい技術については以下のABEJAの服部さんのテックブログも参照ください．

中間まとめ③

長文脈対応は，昨今のLLMの能力向上を支えるtest time scalingにおける重要な技術です．ここも無視できませんね．ただYaRNの登場で非常に簡単にモデルを長文脈対応にできるようになったのは大きいです．実装もシンプルですし，既存モデルのファインチューニングなしで長文脈対応が可能になるのは大きなメリットです．ただYaRNやって終わりというわけではなく，長文理解のためにはAttention機構に関わる工夫以外にも多くの要素が必要です．やっぱりキャッチアップは続けたいです．

5. さらなる展望

ここまでAttention機構の効率化を見てきましたが，最新モデルでは，Attention以外の仕組みも組み合わせることで，効率と精度の両立を図るアプローチが注目されています．

SSM（Mamba系）は長文脈を O(n) で処理できますが，複雑な関係性の把握は苦手．一方，Attentionは関係把握に優れますが計算コストが高い．そこで Jamba や Zamba2 のようなハイブリッドモデルでは，浅い層を SSM 主体にし，深い層で Attention を配置することで，効率と精度の両立を図っています．最近だとQwen3-nextやNemotron-3-nanoはまさにこのアプローチを採用していますね．

「Attentionを使う場所を選ぶ」という発想の転換も，効率化の重要なアプローチです．

全体まとめ

Attentionの進化，構造化すると結構面白いなと思いますし，まだまだ進化の余地がある分野だなとも思います．

私がこのテーマに惹かれるのは，ここに理論と実装の美しい交差点があるからです．数学的に優れたアイデアだけでは不十分で，GPUのメモリ階層やCUDAカーネルの最適化といった「泥臭い」エンジニアリングが不可欠です．逆に，実装の工夫だけでも限界があり，Mambaのような根本的なアーキテクチャの再考が必要になることもあります．

そして何より，この分野は今まさに動いている．2024年から2025年にかけて，本記事で紹介した技術の多くが発表・実装されました．2026年今年は，また新たな手法がスタンダードになっているかもしれません．非常に楽しみです．

結局キャッチアップ続いていきます．頑張りましょう．．．
ということで．長文お付き合いいただき，ありがとうございました．

参考文献・リンク集

基礎論文

計算量の削減

KVキャッシュの処理

長文脈対応

ハイブリッドアーキテクチャ

ツール・ライブラリ

脚注

  1. DeepSeek-V2論文（https://arxiv.org/abs/2405.04434）のTable 4において，MLAを用いたモデルがMHAベースのモデルと同等以上の性能を達成しつつ，KVキャッシュを93.3%削減していることが示されています． ↩︎

GitHubで編集を提案
松尾研究所テックブログ
松尾研究所テックブログ

株式会社松尾研究所のテックブログです。

Discussion