モデル学習ノート · 第1章

位置エンコーディング:Transformer は順序をどう知る?

正弦波から RoPE まで。位置情報を入れる場所と、長い文脈を扱う際の条件を整理します。

Muxuan Liu · 自分のモデル学習ノートを整理し、式と実装例を確認した版です。

順序を別に伝える理由

「犬が猫を追う」と「猫が犬を追う」では、登場する語が似ていても意味が違います。位置情報や順序に依存するマスクを持たない全結合の自己注意は、入力を並べ替えると出力も同様に並べ替わる仕組みで、何番目かを専用に表してはいません。

工学的な問いは、計算のどこへ順序を入れるかです。入力に位置ベクトルを加える方法も、位置の組の相互作用を変える方法もあります。因果マスクも順序の制約を与えるため、すべての Transformer に順序がないと言い切るのは正確ではありません。

絶対位置:各場所に住所を付ける

絶対位置は「自分は何番目か」を表します。元の Transformer は、ネットワークへ入る前の入力埋め込みに正弦波の位置ベクトルを加えます。各層の Q・K・V へそれぞれ加えることが定義ではありません。位置ごとに学習する表を用意する方法もあります。

学習可能な表には設定した範囲があります。長くするには補間や追加学習などが必要です。位置ごとに別の行を持つことは、位置間の関係を学べないことを意味しません。相対距離を明示的には記述しない、という違いです。

正弦波:速さの異なる時計の針

異なる速さで回る複数の針を想像してください。その組み合わせで位置を表すのが直感的な見方です。二つの次元を sin と cos の組にし、周波数を次第に小さくします。式の p は位置、d は埋め込み幅、i は次元の組の番号です。

周波数は 10000^(−2i/d) です。負号を落とすと別の周波数配置になります。未知の位置でも式を計算できますが、それだけで学習時より長い文脈での性能が保証されるわけではありません。

PE(p, 2i) = sin(p · 10000^(−2i/d)) PE(p, 2i+1) = cos(p · 10000^(−2i/d))

相対位置:何個離れているかを伝える

「17 番目」より「三つ左」が役立つ場合もあります。相対位置の手法は、その距離を相互作用へ取り込みます。Shaw らの方法は key/value の計算に相対ベクトルを加えます。距離を端点を含む [−k, k] に切り詰めるなら、区分は 2k+1 個です。

相対位置表現は一つの演算ではありません。スコアに偏りを加える方法、key/value の寄与を変える方法、query と key を回転する RoPE などがあります。名前だけでなく、実際の注意の式を確かめると理解しやすくなります。

RoPE:位置を回転角へ変える

RoPE は query と key の座標の組を、位置に応じた角度で回転します。回転後の内積には相対角度が現れます。R(m)ᵀR(n) = R(n−m) という関係が、絶対位置の回転から相対的な差を得る仕組みを表します。

RoPE 自体は、密な注意の二乗の計算量を線形に変えません。無制限の文脈も自動では得られません。周波数、スケーリング、学習長、注意の実装などが影響します。位置の表現と計算の高速化は分けて考えます。

q̂ₘ = R(m)qₘ, k̂ₙ = R(n)kₙ q̂ₘᵀk̂ₙ = qₘᵀ R(n−m) kₙ

形状から実装を読む

入力は [batch, sequence, d_model]、位置表は [1, max_len, d_model] です。表は batch 方向へブロードキャストされます。register_buffer によりモデルと一緒に保存・デバイス移動されますが、学習対象のパラメータにはなりません。関連する演算すべての勾配を禁止する機能ではありません。

例は奇数・偶数両方の次元に対応します。max_len は用意した表の長さで、正弦波の式の理論的な限界ではありません。検証時には dropout を無効にし、通常は train/eval モードに応じて使います。

PyTorch の例を開く
import math
import torch
from torch import nn

class SinusoidalPositionEncoding(nn.Module):
    """Input: [batch, sequence, d_model]."""
    def __init__(self, d_model, max_len=5000, dropout=0.1):
        super().__init__()
        if d_model < 1 or max_len < 1:
            raise ValueError('d_model and max_len must be positive')
        pos = torch.arange(max_len, dtype=torch.float32)[:, None]
        freq = torch.exp(
            torch.arange(0, d_model, 2, dtype=torch.float32)
            * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros(max_len, d_model)
        pe[:, 0::2] = torch.sin(pos * freq)
        pe[:, 1::2] = torch.cos(pos * freq[:d_model // 2])
        self.register_buffer('pe', pe[None, :, :])
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        if x.ndim != 3 or x.shape[-1] != self.pe.shape[-1]:
            raise ValueError('Expected [batch, sequence, d_model]')
        if x.shape[1] > self.pe.shape[1]:
            raise ValueError('Sequence exceeds configured max_len')
        return self.dropout(x + self.pe[:, :x.shape[1]].to(dtype=x.dtype))