Gosuke Miyashita about feed

Pyxisの開発: アテンション

21 July 2026

前回の記事 でRMSNormを実装した。今回はトランスフォーマーの中核であるアテンションを実装した。対応するPRは #19 。

アテンションとは

ここまでの処理では、各トークンのベクトルはそのトークン単体の情報しか持っていない。embeddingテーブルから取り出したベクトルはトークンの種類だけで決まっていて、前後にどんなトークンが並んでいるかは反映されていない。アテンションは、各トークンが文脈中の他のトークンの情報を取り込むための仕組み。「文脈中のどのトークンにどれだけ注目するか」を計算して、注目先のトークンの情報を自分のベクトルに混ぜ込む。

アテンション自体は2014年にBahdanauらが 機械翻訳のために提案 したもので、2017年の "Attention Is All You Need" がアーキテクチャ全体をアテンション中心に組み立てた。これがトランスフォーマーで、現在のLLMの基礎になっている。

Q/K/V

アテンションでは、各トークンのベクトルにそれぞれ別の重み行列（ wq / wk / wv ）を掛けて、Q・K・Vという3種類のベクトルを作る。あるトークンのQと全トークンのKの内積を取ると、トークン間の関連度に相当するスコアが得られる。このスコアをsoftmax（数値の列を合計1の重みに変換する関数）にかけて、その重みで全トークンのVを加重平均したものが、そのトークンの新しいベクトルになる。

Q・K・Vの名前はそれぞれQuery・Key・Valueの略だが、これはこの構造がなぜ機能するかを人間が理解するための後付けの解釈で、動作の本質は「入力に別々の重みを掛けて得たベクトルの内積を関連度として使う」というだけ。 wq / wk / wv はモデルの訓練時に、次トークン予測の精度が上がるように学習された結果としてそういう役割を担うようになっている。式で書くとこうなる。 

Attention(Q, K, V) = softmax(QKᵀ / sqrt(head_dim)) V

内積を sqrt(head_dim) で割っているのは、次元数が大きくなるほど内積の値も大きくなり、softmaxの出力が極端な分布（ほぼ1つのトークンだけに全部の重みが集中する）になってしまうのを防ぐため。scaled dot-product attentionと呼ばれる。

因果マスク

上の式のままだと、各トークンは自分の前後の全トークンにアテンションできる。 QKᵀ は各トークンのQと全トークンのKの内積を取るので、 seq_len × seq_len のスコア行列になる。

因果マスク適用前のスコア行列。全てのマスにsijが埋まっている

たとえばQ 1 （1番目のトークンのQuery）は、自分より後ろのK 2 /K 3 /K 4 /K 5 とも内積を取ってしまう。ここには「前しか見ない」という制約は入っていない。

しかしLLMのテキスト生成では、各トークンは自分より後ろのトークンを参照してはいけない。生成時点ではまだ存在しないトークンだから、参照しようがない。学習時と推論時で条件を揃えるため、学習時も未来のトークンは見せない、というルールにする。

そのために使うのが因果マスク（causal masking）。位置 pos のトークンは位置 0..=pos のトークンだけにアテンションする、という制約を入れる。具体的には、スコア行列の上三角部分（s ij で j > i のところ）を -∞ にしておく。 -∞ はsoftmaxを通すと0になるので、未来のトークンの重みは0になり、実質的にアテンションの対象から除外される。

因果マスクを適用したスコア行列。対角線とその下は通常のスコア（sij）、上三角部分は-∞で潰されている

マルチヘッドアテンション

ここまでは「1つのトークンから1つのQベクトルができる」という前提で話してきた。実際には、そのQベクトル（Qwen3-1.7Bだとhidden_dim = 2048次元）を、head_dim = 128次元ずつ16個に分割して使う。分割したそれぞれを「ヘッド」と呼ぶ。K・Vも同じように分割される。

各ヘッドで独立にアテンション（ softmax(QKᵀ / sqrt(head_dim)) V ）を計算して、最後に全ヘッドの結果を連結して1つのベクトルに戻す。連結後にさらに出力用の重み行列（ wo ）を掛けて、最終的な出力ができる。

なぜ分割するのかというと、実験的にその方が性能が上がるから。 "Attention Is All You Need" で提案された仕組みで、なぜ機能するかは「ヘッドごとに違う観点でトークン間の関連を捉えられるから」という直感的な説明がよくされるが、これは後付けの解釈で、実際のところは分かっていない。ヘッドの多くは冗長で削除しても性能が落ちない、という 報告 もある。

GQA

素直に実装すると、Q・K・Vを全部同じヘッド数（Qwen3-1.7Bなら16個）作ることになる。この構成をMulti-Head Attention（MHA）と呼ぶ。

Qwen3はさらに一歩進んで、GQA（Grouped Query Attention、 Ainslieらが2023年に提案 ）という構成を使っている。Qwen3-1.7BではQが16ヘッドあるのに対して、K/Vは8ヘッドしかない。2つのQヘッドが同じ1つのK/Vヘッドを共有する（Qヘッド0と1がK/Vヘッド0を、Qヘッド2と3がK/Vヘッド1を、という具合に）。

なぜK/Vのヘッド数を減らすかというと、K/Vはメモリを食うから。特に推論時にはKVキャッシュ（過去のトークンのK/Vを保存して再利用する仕組み、後のPRで実装）で大きなメモリを消費するので、K/Vのヘッド数を減らすことでメモリ使用量を抑えられる。

RoPE

アテンションの計算は内積の集まりなので、そのままではトークンの並び順が結果に影響しない。「AがBを」でも「BがAを」でも同じ計算結果になってしまう。なので位置情報を何らかの形でベクトルに埋め込む必要があって、Qwen3ではRoPE（Rotary Position Embedding、 Suらが2021年に提案 ）が使われている。

RoPEは、QとKのベクトルをトークンの位置に応じた角度で回転させる手法。ベクトルの要素を2つずつペアにして、各ペアを2次元平面上の点とみなして回転する。回転角は「位置 × 周波数」で決まり、周波数はペアごとに異なる。周波数の基準となる定数がrope_thetaで、Qwen3では1000000。

内積は2つのベクトルの相対的な角度差に依存するので、位置に応じて回転させておくと、QとKの内積に位置の差（相対位置）の情報が乗る。これがRoPEの仕組み。

なお、このPRの実装は隣り合う要素（0番目と1番目、2番目と3番目、...）をペアにして回転していたが、HuggingFaceの参照実装は前半と後半の要素（0番目とhead_dim/2番目、...）をペアにする方式で、このズレは後のPRで修正している。

実装 

pub struct Attention {
    wq: Vec<f32>,
    wk: Vec<f32>,
    wv: Vec<f32>,
    wo: Vec<f32>,
    num_q_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
    rope_theta: f32,
}

wq / wk / wv がQ/K/Vを作るための重み行列で、Qwen3では model.layers.0.self_attn.q_proj.weight のような名前でsafetensorsファイルに格納されている。 wo は出力用の重み行列で、全ヘッドの結果を連結したものに掛けて最終的な出力を作る。

forward() の流れはこうなる。

  1. 各位置の入力ベクトルに wq / wk / wv を掛けてQ/K/Vを作る
  2. QとKにRoPEを適用する
  3. 各位置・各Qヘッドについて、自分以前の位置のKとの内積を取り、 sqrt(head_dim) で割ってスコアにする
  4. スコアをsoftmaxで重みに変換する
  5. 重みでVを加重平均する
  6. 全ヘッドの結果を連結して wo を掛ける

スコア計算の部分を抜粋するとこうなっている。 

let mut scores = vec![f32::NEG_INFINITY; seq_len];

for (key_pos, score) in scores.iter_mut().enumerate().take(pos + 1) {
    let k_start = key_pos * kv_dim + kv_head * self.head_dim;
    let dot = (0..self.head_dim)
        .map(|dim| queries[q_start + dim] * keys[k_start + dim])
        .sum::<f32>();
    *score = dot / scale;
}

softmax(&mut scores);

因果マスクは、スコアを NEG_INFINITY で初期化して、自分以前の位置（ take(pos + 1) ）だけスコアを計算する、という形で実装している。 NEG_INFINITY はsoftmaxを通すと0になるので、未来のトークンの重みは0になる。

GQAは、Qヘッドの番号から対応するK/Vヘッドの番号を求めることで実現している。 

let kv_head = q_head * self.num_kv_heads / self.num_q_heads;

Qが16ヘッド、K/Vが8ヘッドなら、Qヘッド0と1がK/Vヘッド0を、Qヘッド2と3がK/Vヘッド1を使う。

これで、各トークンのベクトルに、文脈中の他のトークンの情報を関連度に応じて混ぜ込む仕組みができた。次回はトランスフォーマー層のもうひとつの構成要素であるFFN（SwiGLU）の実装について書く。