Rotary Position Embedding (RoPE)

Rotary Position Embedding（RoPE）は、各クエリベクトルとキー・ベクトルをトークン位置に比例した角度だけ回転させることで、Transformerにトークンの順序を伝えます。自己アテンションだけでは単語の順序を区別できないため、モデルには位置情報を何らかの形で与える必要があります。RoPEでは、位置mのクエリと位置nのキーの間のアテンションスコアは、そのオフセットm − nだけに依存します。RoPEは2021年にRoFormerモデルとともに発表され、EleutherAIのGPT-JとGPT-NeoXによって普及しました。現在では、Llama、Mistral、Qwen、Gemma、DeepSeekなどのオープンな大規模言語モデル（LLM）で標準的な位置エンコーディングになっています。

RoPEの仕組み#

RoPEでは、次元数dの各クエリベクトルとキー・ベクトルを、d/2個の特徴ペアに分け、それぞれを平面上の点として扱います。位置mでは、ペアiを角度m·θᵢだけ回転させます。θᵢ = base^(−2i/d)であり、元の論文ではbaseに10,000を使います。最初のペアは速く回転して近傍の細かな順序を区別し、最後のペアはゆっくり回転して数千位置離れたトークンも区別します。クエリをmθ、キーをnθだけ回転させても、内積は差分(m − n)θだけに依存し、回転によってベクトルの長さが変わることはありません。回転は、クエリとキーの射影後に各アテンション層内で適用します。値ベクトルは回転せず、入力時のトークン埋め込みにも何も加えません。

以下のスニペットは、MetaのLlamaリファレンスコードと同様に複素数を使って回転を実装し、相対位置の性質を検証します。

import torch

torch.manual_seed(0)
dim = 64
freqs = 10000 ** (-torch.arange(0, dim, 2) / dim)  # one rotation frequency per pair of features


def rope(x, pos):
    """Rotate each consecutive feature pair of x by the angle pos * frequency."""
    pairs = torch.view_as_complex(x.reshape(-1, 2))
    return torch.view_as_real(pairs * torch.polar(torch.ones_like(freqs), pos * freqs)).flatten()


q, k = torch.randn(2, dim)
# The same query and key, three tokens apart, at two different places in the sequence
print(f"positions 5 and 2:     {rope(q, 5) @ rope(k, 2):.4f}")
print(f"positions 105 and 102: {rope(q, 105) @ rope(k, 102):.4f}")
print(f"positions 5 and 4:     {rope(q, 5) @ rope(k, 4):.4f}")  # one token apart
print(f"Vector length kept: {rope(q, 5).norm():.4f} vs {q.norm():.4f}")

位置5と2、または105と102にある場合でも、どちらのペアも3トークン離れていれば、出力は-8.2811になります。同じベクトルが1トークン離れている場合は、-11.4563になります。回転後のクエリの長さは8.3705のままです。

RoPEとその他の位置エンコーディングの比較#

• 正弦波絶対位置エンコーディング：元のTransformerでは、最初の層の前に固定の正弦・余弦ベクトルを入力埋め込みに1回加えます。相対距離は、混ぜ込まれた信号から推論する必要があります。
• 学習可能な絶対位置埋め込み：BERTとGPT-2は位置ごとに1つのベクトルを学習しますが、最大512位置と1,024位置までで、それを超える位置の埋め込みはありません。
• 相対位置バイアス：T5とALiBiは、学習可能または固定の距離ベースのバイアスをアテンションスコアに加えます。RoPEは回転自体で相対距離をエンコードし、パラメーターを追加しません。

コンテキストウィンドウの拡張#

RoPEをシーケンス長Lで学習したモデルは、Lを超える位置の回転角度を見たことがなく、それらの位置に遭遇するとアテンションの性能が低下します。コンテキスト拡張手法では、最初から再学習する代わりに角度を再スケーリングします。Position Interpolationは、新しい位置を学習済みの範囲に収め、1,000回のファインチューニングステップ以内にLlamaモデルのコンテキストウィンドウを32,768トークンに拡張しました。NTK-aware scalingとYaRNは、速く回転するペアよりも遅く回転する低周波数のペアを大きく伸ばします。YaRNは、以前の手法より学習トークン数を10分の1に抑えながら、Llama 2を128Kトークンまで拡張しました。Llama 3は事前学習中にbaseを10,000から500,000に引き上げたため、最も遅いペアのトークンごとの回転量が小さくなり、遠く離れた位置も区別できます。

ビジョンモデルにおけるRoPE#

画像には2つの座標が必要なため、ビジョントランスフォーマーは2D、つまり軸方向のRoPEを使います。各ヘッドの特徴ペアの半分はパッチの列に応じて回転し、残りの半分は行に応じて回転します。Qwen2-VLはマルチモーダルなM-RoPEに時間軸を3つ目の軸として追加し、動画に対応します。Ultralyticsパッケージで動作するMetaのSAM 2は、メモリアテンションの4つの層で軸方向RoPEを適用します。この層では、現在の動画フレームの特徴が、過去のフレームから保存された特徴にアテンションを向けます。保存されたすべてのフレームのキーには、現在のフレームのクエリと同じグリッドの角度が割り当てられるため、両者間のアテンションはパッチ間の空間オフセットに依存します。