sharin

Appearance

位置エンコーディングとRoPE ​

実装: llm/rope/ / 実行: go test ./llm/rope/

attention はトークンの順序を見ないので、位置情報を外から注入する必要がある。mini-GPT では位置ベクトルを足す絶対位置方式を使ったが、今の主流は Q と K を位置に応じて回転させる RoPE だ。回転どうしの内積は角度の差だけで決まるため、スコアが相対位置のみに依存する。この章では両方式を実装して相対性をテストで確かめ、長コンテキスト化の入口の位置補間まで作る。

この章で作るもの ​

attention は全トークンの内積を取るだけで、並び順をどこにも使っていない。"cat sat" と "sat cat" は、位置情報がなければ同じに見える。mini-GPT では「位置 i 用の学習ベクトルを埋め込みに足す」方式でこれを補ったが、この方式には弱点があり、実物の Llama や Qwen は別の方式に移っている。それが RoPE(rotary position embedding)だ。

発想の転換は、位置を「足す」のではなく「回す」ことにある。Q と K の各次元ペアを 2 次元平面と見なし、位置 pos のトークンは角度 pos·θ だけ回転させる。回転された 2 本のベクトルの内積は角度の差だけで決まるという幾何の性質から、attention スコアは 2 つのトークンの間隔のみに依存するようになる。

位置 0     位置 1     位置 2      各ペア (x0,x1) を角度 pos·θ で回転
   ↑         ↗         →
   │        ╱
   ●───    ●───      ●───

内積( R(m)·q, R(n)·k ) = 角度差 (m−n)·θ だけの関数
→ (位置9のq, 位置5のk) と (位置109のq, 位置105のk) は同じスコア
RoPE の仕組み。次元をペアに分け、各ペアを位置に比例した角度で回転させる。位置 m の Q と位置 n の K の内積は角度差 (m−n)·θ だけで決まるので、文全体が前後にずれても注目のパターンは変わらない

順に見ていく。

  1. 位置は足すのではなく回す: Q と K を次元ペアごとに角度 pos·freq で回転させる。V は回さない
  2. スコアが相対位置だけに依存: 回転の内積は角度差のみで決まる。「5 語前を見る」パターンが文のどこでも同じに働く
  3. 周波数の階層: 先頭ペアは高周波で近距離を細かく、末尾ペアは低周波で遠距離を大まかに測る。1 つのベクトルに複数の目盛りを持つ

① 絶対位置: 足す方式とその弱点 ​

まず古典から。原論文 "Attention Is All You Need" の位置エンコーディングは、位置ごとに sin/cos のパターンを作って埋め込みに足す:

go

// Sinusoidal は "Attention Is All You Need" の絶対位置エンコーディングを返す。
// 各位置 pos の偶数次元に sin(pos/10000^(2i/d))、奇数次元に cos を置き、
// これをトークン埋め込みに足す。mini-GPT の学習型位置埋め込みの学習不要版。
func Sinusoidal(seqLen, dim int) [][]float64 {
	pe := make([][]float64, seqLen)
	for pos := 0; pos < seqLen; pos++ {
		row := make([]float64, dim)
		for i := 0; i+1 < dim; i += 2 {
			angle := float64(pos) / math.Pow(10000, float64(i)/float64(dim))
			row[i] = math.Sin(angle)
			row[i+1] = math.Cos(angle)
		}
		pe[pos] = row
	}
	return pe
}

位置 0 と位置 1 で違うベクトルが足されるので、attention は位置を区別できるようになる。mini-GPT の学習型位置埋め込みは、この表を学習で獲得する変種にあたる。

だが「足す」方式には構造的な弱点が 2 つある。1 つ目は、内積が絶対位置に依存することだ。「直前の語を見る」という関係は位置 5 と 4 の間でも位置 105 と 104 の間でも同じはずだが、足された位置ベクトルは位置ごとに違うので、モデルは同じ関係を位置ごとに学び直すことになる。2 つ目は、学習時に見た最大長を超えた位置に対応できないことだ。学習型なら位置 2049 のベクトルがそもそも存在しない。

② RoPE: 回す方式 ​

RoPE はこの 2 つを同時に解く。実装は短い:

go

// RoPE は次元をペア (x0,x1), (x2,x3), ... に分け、ペア i を
// 角度 pos·freqs[i] だけ回転させる。freqs は先頭ペアほど高周波
// (1 位置で大きく回る = 近距離を細かく見る)、末尾ほど低周波
// (ゆっくり回る = 遠距離の大まかな位置を持つ)。
type RoPE struct {
	dim   int
	freqs []float64
}

// New は dim 次元(偶数)用の RoPE を作る。freqs[i] = 10000^(-2i/dim)。
func New(dim int) (*RoPE, error) {
	if dim <= 0 || dim%2 != 0 {
		return nil, errors.New("rope: dim must be positive and even")
	}
	freqs := make([]float64, dim/2)
	for i := range freqs {
		freqs[i] = math.Pow(10000, -2*float64(i)/float64(dim))
	}
	return &RoPE{dim: dim, freqs: freqs}, nil
}

// Apply はベクトル x を位置 pos の回転にかけた新しいベクトルを返す。
// 回転は直交変換なので長さを変えず、同じ pos なら常に同じ回転になる。
// attention では Q と K にだけ適用する(V は回さない)。
func (r *RoPE) Apply(x []float64, pos int) []float64 {
	return r.applyAt(x, float64(pos))
}

func (r *RoPE) applyAt(x []float64, pos float64) []float64 {
	out := make([]float64, len(x))
	copy(out, x)
	for i := 0; i*2+1 < len(x) && i < len(r.freqs); i++ {
		theta := pos * r.freqs[i]
		sin, cos := math.Sin(theta), math.Cos(theta)
		a, b := x[i*2], x[i*2+1]
		out[i*2] = a*cos - b*sin
		out[i*2+1] = a*sin + b*cos
	}
	return out
}

見どころは 2 つある。1 つは周波数の設計だ。freqs[i] = 10000^(-2i/dim) により、先頭のペアは 1 位置で約 1 ラジアン回る高周波、末尾のペアはほとんど回らない低周波になる。時計の秒針と時針のように、1 つのベクトルが目盛りの違う複数の回転を同時に持ち、近距離の細かい間隔も遠距離の大まかな間隔も表せる。

もう 1 つは適用先だ。回すのは Q と K だけで、V は回さない。位置が関与すべきなのは「どこに注目するか」の採点であって、注目先から運んでくる内容ではないからだ。

相対性は実装のテストで固定している。位置 9 の Q と位置 5 の K の内積を取り、両方を 1、10、500 ずらしても内積は変わらない。一方で間隔そのものを変えれば内積は変わる。「文が前にずれても注目パターンが不変」という性質が、恒等式としてコードで確かめられる。

③ 長コンテキストへ: 位置補間 ​

RoPE にも学習レンジの問題は残る。学習時に位置 4096 までしか見ていないモデルに位置 8000 を与えると、経験のない角度の組み合わせが現れて性能が崩れる。最も素朴な対策が位置補間(position interpolation)で、位置を factor で割ってから回す:

go

// ApplyInterpolated は位置補間(position interpolation)つきの適用。
// 位置を factor で割ってから回すことで、学習時より長い系列でも
// 角度が学習済みのレンジに収まる。長コンテキスト化の最も素朴な形。
func (r *RoPE) ApplyInterpolated(x []float64, pos int, factor float64) []float64 {
	return r.applyAt(x, float64(pos)/factor)
}

位置 8000 を factor 4 で補間すると位置 2000 と同じ回転になり、角度は学習済みのレンジに収まる。そのかわり位置の目盛りが 4 倍粗くなるので、近距離の解像度は下がる。実物ではこの粗さを周波数ごとに調整する改良(NTK-aware スケーリング、YaRN)が使われており、高周波ペアは補間せず近距離の解像度を守り、低周波ペアだけ補間して遠距離を伸ばす。

動かす ​

下のデモは、この実装の回転をそのまま描いている。「回転を見る」では 1 つのペアが位置とともに回る様子と、周波数の違うペアが別々の速さで回ることを確かめられる。「相対性」では位置 m の Q と位置 n の K の内積を表示しながら全体を後ろへずらし、絶対位置が変わってもスコアが動かないことを見る。

デモRoPE(回転位置埋め込み)位置 0
回転を見る相対性
ペア0(高周波)freq 1.000 / 角度 0°
ペア1(中間)freq 0.100 / 角度 0°
ペア3(低周波)freq 0.001 / 角度 0°

位置 0 では全ペアが回転ゼロ(恒等)。ここから位置が進むごとに、各ペアが自分の周波数ぶんだけ回る

1 / 8

RoPE は Q と K の次元ペアを位置に比例した角度で回す。回転の内積は角度差だけで決まるので、 attention スコアは 2 トークンの間隔のみに依存する。「5 語前を見る」パターンが 文のどこにいても同じに働くのはこの性質による。

設計の観点 ​

  • なぜ Q/K だけ回すか: 位置が効くべきは注目の採点(Q·K)だけ。V まで回すと出力の内容自体が位置で歪む
  • KV キャッシュとの相性: 各位置の回転は固定なので、キャッシュ済みの K を後から回し直す必要がない。生成で 1 トークン足すたびに新しい位置の回転を新トークンにだけ適用すればよい
  • 外挿と補間の区別: 学習レンジ外の位置をそのまま使うのが外挿(崩れやすい)、レンジ内に押し込むのが補間(解像度を犠牲に安定)。NTK/YaRN は周波数ごとに補間量を変える折衷
  • ALiBi との比較: attention スコアに距離ペナルティを直接足す方式もある(ALiBi)。実装は最も簡単で外挿に強いが、表現力で RoPE が優勢となり主流は RoPE に収束した
  • 長コンテキストの費用: 位置の表現を伸ばしても、attention の n² コストと KV キャッシュのメモリは別問題として残る(次の attention 変種・推論高速化の章で扱う)

メリット・デメリットと実例 ​

方式相対性レンジ外への強さ実例
学習型絶対位置なし不可(ベクトルが無い)GPT-2、GPT-3、mini-GPT
Sinusoidalなし(理論上は外挿可)弱いTransformer 原論文
RoPEあり補間・NTK/YaRN で拡張Llama 系、Qwen、DeepSeek、Mistral
ALiBiあり(距離ペナルティ)強いMPT、BLOOM

裏どり:

  • RoFormer(2021): RoPE の原論文。中国の Zhuiyi Technology 発で、のちに Llama が採用して事実上の標準になった
  • Llama 系の長コンテキスト化: Llama 2 の 4k から Code Llama の 100k への拡張は RoPE の周波数基底の変更(base 10000 → 1000000)で実現された。回転の設計だけでコンテキスト長が桁で変わる実例
  • Position Interpolation(2023): Meta の論文。少量の追加学習と補間で 2k → 32k に拡張できることを示した
  • YaRN(2023): 周波数ごとに補間戦略を変える改良。オープンモデルの長コンテキスト化で広く使われる

簡略化したこと ​

  • 1 ベクトルへの適用のみ: 実物は (batch, head, seq, dim) のテンソルへ一括適用する。回転の中身は同じ
  • NTK-aware / YaRN 未実装: 一様な位置補間まで。周波数別の補間は解説に留めた
  • attention への組み込みなし: attention の実装は位置情報を持たないまま、mini-GPT は足す方式のままにしてある。RoPE を Q/K に挟む位置だけ本文で示した
  • 複素数表現なし: 論文の複素数形式でなく、同値な実数ペアの 2 次元回転で書いた

参考資料 ​