位置エンコーディングの理論と各種手法を解説

Transformer の Self-Attention 機構は入力トークンの順序を区別できないため、位置情報を外部から注入する必要があります。この「位置エンコーディング(Positional Encoding)」の設計は、モデルの長文処理能力や汎化性能に直結する重要な研究課題です。

原論文の正弦波位置エンコーディングから、学習可能な位置エンコーディング、相対位置エンコーディング、RoPE(Rotary Position Embedding)、ALiBi(Attention with Linear Biases)まで、各手法の数学的原理を厳密に導出し、それぞれの長所・短所を比較します。

本記事の内容

  • Self-Attention の置換不変性の証明と位置情報の必要性
  • 正弦波位置エンコーディングの性質(相対位置の内積表現、線形変換可能性の証明)
  • 学習可能な位置エンコーディング
  • 相対位置エンコーディング(Shaw et al.)
  • RoPE(回転行列による位置注入、内積が相対位置のみに依存する証明)
  • ALiBi(線形バイアス)
  • 各手法の比較と長文対応能力
  • Python での各種位置エンコーディングの実装と可視化

前提知識

この記事を読む前に、以下の記事を読んでおくと理解が深まります。

なぜ位置情報が必要か — Self-Attention の置換不変性

置換不変性の厳密な証明

Self-Attention が入力の順序を区別できないことを厳密に証明します。

入力 $\bm{X} = [\bm{x}_1, \bm{x}_2, \dots, \bm{x}_T]^T \in \mathbb{R}^{T \times d}$ に対し、Self-Attention は

$$ \text{SA}(\bm{X}) = \text{softmax}\left(\frac{\bm{X}\bm{W}^Q (\bm{X}\bm{W}^K)^T}{\sqrt{d_k}}\right)\bm{X}\bm{W}^V $$

任意の置換 $\pi$ に対応する置換行列 $\bm{P}_\pi \in \{0, 1\}^{T \times T}$ を考えます。$\bm{P}_\pi$ は各行・各列にちょうど1つの1を持つ直交行列で、$\bm{P}_\pi^T \bm{P}_\pi = \bm{I}$ を満たします。

$\bm{P}_\pi \bm{X}$ は $\bm{X}$ の行を $\pi$ で並べ替えた行列です。Self-Attention に $\bm{P}_\pi \bm{X}$ を入力すると

$$ \text{SA}(\bm{P}_\pi \bm{X}) = \text{softmax}\left(\frac{\bm{P}_\pi \bm{X}\bm{W}^Q (\bm{P}_\pi \bm{X}\bm{W}^K)^T}{\sqrt{d_k}}\right)\bm{P}_\pi \bm{X}\bm{W}^V $$

$\bm{Q}’ = \bm{P}_\pi \bm{X}\bm{W}^Q = \bm{P}_\pi \bm{Q}$, $\bm{K}’ = \bm{P}_\pi \bm{K}$, $\bm{V}’ = \bm{P}_\pi \bm{V}$ として

$$ \bm{Q}'{\bm{K}’}^T = \bm{P}_\pi \bm{Q} (\bm{P}_\pi \bm{K})^T = \bm{P}_\pi \bm{Q}\bm{K}^T \bm{P}_\pi^T $$

ここで $\bm{S} = \bm{Q}\bm{K}^T / \sqrt{d_k}$ とし、softmax を行方向に適用する操作を $\sigma$ と書くと

$$ \sigma(\bm{P}_\pi \bm{S} \bm{P}_\pi^T) = \bm{P}_\pi \sigma(\bm{S}) \bm{P}_\pi^T $$

この等式は、$\bm{P}_\pi \bm{S} \bm{P}_\pi^T$ の $(\pi(i), \pi(j))$ 成分が $S_{ij}$ であり、行 $\pi(i)$ に softmax を適用した結果は行 $i$ に softmax を適用した結果を $\pi$ で並べ替えたものに等しいことから導かれます。

したがって

$$ \begin{align} \text{SA}(\bm{P}_\pi \bm{X}) &= \bm{P}_\pi \sigma(\bm{S}) \bm{P}_\pi^T \cdot \bm{P}_\pi \bm{V} \\ &= \bm{P}_\pi \sigma(\bm{S}) (\bm{P}_\pi^T \bm{P}_\pi) \bm{V} \\ &= \bm{P}_\pi \sigma(\bm{S}) \bm{V} \\ &= \bm{P}_\pi \, \text{SA}(\bm{X}) \end{align} $$

つまり $\text{SA}(\bm{P}_\pi \bm{X}) = \bm{P}_\pi \, \text{SA}(\bm{X})$ が成り立ちます。入力の並べ替えに対して出力が同じように並べ替わるだけであり、Self-Attention は位置の絶対的な情報を一切利用していません。これを同変性(equivariance)と呼びます。

特に、各位置の出力だけを見ると、同じ集合の入力であれば順序によらず同じ出力が得られます(集合に対する不変性)。自然言語では語順が意味を決定するため、何らかの方法で位置情報を注入する必要があります。

正弦波位置エンコーディング

定義

Vaswani et al.(2017)が提案した正弦波位置エンコーディングは

$$ \begin{align} PE(pos, 2i) &= \sin\left(\frac{pos}{10000^{2i/d}}\right) \\ PE(pos, 2i+1) &= \cos\left(\frac{pos}{10000^{2i/d}}\right) \end{align} $$

ここで $d = d_{\text{model}}$, $pos = 0, 1, \dots, T-1$, $i = 0, 1, \dots, d/2 – 1$ です。

$\omega_i = 1/10000^{2i/d}$ と定義すると、次元ペア $(2i, 2i+1)$ は周波数 $\omega_i$ の正弦波・余弦波を形成します。

性質1: 相対位置の内積表現

$PE(pos)$ と $PE(pos + k)$ の内積を計算します。

$$ \begin{align} \langle PE(pos), PE(pos+k) \rangle &= \sum_{i=0}^{d/2-1} \Big[ \sin(\omega_i \cdot pos)\sin(\omega_i(pos+k)) + \cos(\omega_i \cdot pos)\cos(\omega_i(pos+k)) \Big] \end{align} $$

積和の公式 $\cos(A – B) = \cos A \cos B + \sin A \sin B$ より

$$ \sin(\omega_i \cdot pos)\sin(\omega_i(pos+k)) + \cos(\omega_i \cdot pos)\cos(\omega_i(pos+k)) = \cos(\omega_i \cdot k) $$

したがって

$$ \langle PE(pos), PE(pos+k) \rangle = \sum_{i=0}^{d/2-1} \cos(\omega_i \cdot k) $$

この結果は絶対位置 $pos$ に依存せず、相対位置 $k$ のみの関数です。これは位置エンコーディングの内積が相対的な距離を反映するという望ましい性質です。

さらに、$k$ が増加すると $\cos(\omega_i \cdot k)$ は振動しながら減衰する傾向があり(異なる周波数の $\cos$ の和はデコヒーレンスを起こす)、位置が離れるほど内積は小さくなります。

性質2: 線形変換可能性の証明

任意の固定オフセット $k$ に対して、$PE(pos + k)$ が $PE(pos)$ の線形変換で表せることを証明します。

次元ペア $(2i, 2i+1)$ に注目します。三角関数の加法定理より

$$ \begin{align} \sin(\omega_i(pos+k)) &= \sin(\omega_i \cdot pos)\cos(\omega_i \cdot k) + \cos(\omega_i \cdot pos)\sin(\omega_i \cdot k) \\ \cos(\omega_i(pos+k)) &= \cos(\omega_i \cdot pos)\cos(\omega_i \cdot k) – \sin(\omega_i \cdot pos)\sin(\omega_i \cdot k) \end{align} $$

行列形式で書くと

$$ \begin{pmatrix} PE(pos+k, 2i) \\ PE(pos+k, 2i+1) \end{pmatrix} = \underbrace{\begin{pmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{pmatrix}}_{\bm{R}_k^{(i)}} \begin{pmatrix} PE(pos, 2i) \\ PE(pos, 2i+1) \end{pmatrix} $$

$\bm{R}_k^{(i)}$ は $\omega_i k$ の角度の回転行列であり、$pos$ に依存しません。

全次元をまとめると、ブロック対角行列

$$ \bm{M}_k = \text{diag}\left(\bm{R}_k^{(0)}, \bm{R}_k^{(1)}, \dots, \bm{R}_k^{(d/2-1)}\right) \in \mathbb{R}^{d \times d} $$

を用いて

$$ PE(pos + k) = \bm{M}_k \cdot PE(pos) $$

と書けます。$\bm{M}_k$ は $k$ のみに依存する行列であり、相対位置 $k$ の情報が線形変換として表現できることを意味します。モデルはこの線形構造を利用して相対位置関係を学習できます。

性質3: 異なる位置のエンコーディングの一意性

$d$ が十分大きいとき、異なる位置 $pos \neq pos’$ に対して $PE(pos) \neq PE(pos’)$ が保証されます。これは、$\omega_i$ が異なる値を取るため、$d/2$ 個の周波数の正弦波の組み合わせが一意の「指紋」を形成するためです。

学習可能な位置エンコーディング

学習可能な位置エンコーディングでは、位置 $pos$ に対応する $d$ 次元ベクトル $\bm{p}_{pos} \in \mathbb{R}^d$ を学習可能なパラメータとして定義します。

$$ \bm{P} = [\bm{p}_0, \bm{p}_1, \dots, \bm{p}_{T_{\max}-1}]^T \in \mathbb{R}^{T_{\max} \times d} $$

入力埋め込み $\bm{x}_{pos}$ に位置エンコーディングを加算します。

$$ \tilde{\bm{x}}_{pos} = \bm{x}_{pos} + \bm{p}_{pos} $$

利点: * データから最適な位置表現を直接学習できる * 正弦波のような固定パターンに制約されない

欠点: * 学習時に見た最大系列長 $T_{\max}$ を超える位置に外挿できない * パラメータ数が $T_{\max} \times d$ 増加する

BERT、GPT-2、ViT などで採用されています。実験的には、正弦波との性能差は小さいことが多いですが、外挿能力において正弦波が優れるケースがあります。

相対位置エンコーディング(Shaw et al., 2018)

Shaw et al. は、絶対位置ではなく相対位置を直接モデル化するアプローチを提案しました。

定式化

通常の Self-Attention のスコア計算

$$ e_{ij} = \frac{\bm{x}_i \bm{W}^Q (\bm{x}_j \bm{W}^K)^T}{\sqrt{d_k}} $$

に相対位置情報を追加します。

$$ e_{ij} = \frac{\bm{x}_i \bm{W}^Q (\bm{x}_j \bm{W}^K + \bm{a}_{ij}^K)^T}{\sqrt{d_k}} $$

ここで $\bm{a}_{ij}^K \in \mathbb{R}^{d_k}$ は相対位置 $i – j$ に対応する学習可能なベクトルです。同様に、バリューの計算にも相対位置情報を追加します。

$$ \bm{z}_i = \sum_j \alpha_{ij} (\bm{x}_j \bm{W}^V + \bm{a}_{ij}^V) $$

相対位置はクリップされます。

$$ \bm{a}_{ij}^K = \bm{w}_{\text{clip}(i-j, -C, C)}^K, \quad \text{clip}(x, a, b) = \max(a, \min(b, x)) $$

$C$ はクリッピング距離で、$|i – j| > C$ の位置は全て同じベクトルで表現されます。これにより、パラメータ数は $O(C \cdot d_k)$ に抑えられます。

Attention スコアの展開

スコアを展開して、各項の役割を明確にしましょう。

$$ \begin{align} e_{ij} &= \frac{1}{\sqrt{d_k}} (\bm{x}_i \bm{W}^Q)(\bm{x}_j \bm{W}^K + \bm{a}_{ij}^K)^T \\ &= \frac{1}{\sqrt{d_k}} \left[ \underbrace{(\bm{x}_i \bm{W}^Q)(\bm{x}_j \bm{W}^K)^T}_{\text{content-content}} + \underbrace{(\bm{x}_i \bm{W}^Q)(\bm{a}_{ij}^K)^T}_{\text{content-position}} \right] \end{align} $$

第1項はトークン内容間の類似度、第2項はクエリ内容と相対位置の相互作用を表します。

RoPE(Rotary Position Embedding)

Su et al.(2021)が提案した RoPE は、回転行列を用いて位置情報を注入する手法です。内積が自然に相対位置のみに依存するという優れた性質を持ちます。

基本アイデア

RoPE の核心は、位置 $m$ のクエリ $\bm{q}$ と位置 $n$ のキー $\bm{k}$ に対して、回転行列 $\bm{R}_m$, $\bm{R}_n$ を乗じることです。

$$ \tilde{\bm{q}}_m = \bm{R}_m \bm{q}, \quad \tilde{\bm{k}}_n = \bm{R}_n \bm{k} $$

Attention スコアの内積は

$$ \tilde{\bm{q}}_m^T \tilde{\bm{k}}_n = \bm{q}^T \bm{R}_m^T \bm{R}_n \bm{k} $$

ここで回転行列の性質 $\bm{R}_m^T = \bm{R}_{-m}$ と $\bm{R}_{-m}\bm{R}_n = \bm{R}_{n-m}$ より

$$ \tilde{\bm{q}}_m^T \tilde{\bm{k}}_n = \bm{q}^T \bm{R}_{n-m} \bm{k} $$

内積が相対位置 $n – m$ のみに依存することが分かります。

回転行列の具体形

$d$ 次元ベクトルを $d/2$ 個の2次元ペアに分割し、各ペアに異なる角速度の回転を適用します。

位置 $m$ の回転行列は

$$ \bm{R}_m = \begin{pmatrix} \bm{R}_m^{(0)} & & \\ & \bm{R}_m^{(1)} & \\ & & \ddots \\ & & & \bm{R}_m^{(d/2-1)} \end{pmatrix} $$

各 $2 \times 2$ ブロックは

$$ \bm{R}_m^{(i)} = \begin{pmatrix} \cos(m\theta_i) & -\sin(m\theta_i) \\ \sin(m\theta_i) & \cos(m\theta_i) \end{pmatrix} $$

ここで $\theta_i = 10000^{-2i/d}$ は正弦波位置エンコーディングと同じ周波数です。

内積が相対位置のみに依存する証明

RoPE を適用したクエリとキーの内積を厳密に導出します。

2次元ペア $(2i, 2i+1)$ に注目します。位置 $m$ のクエリの $i$ 番目のペアを $\bm{q}^{(i)} = (q_{2i}, q_{2i+1})^T$、位置 $n$ のキーの $i$ 番目のペアを $\bm{k}^{(i)} = (k_{2i}, k_{2i+1})^T$ とします。

RoPE 適用後のペアは

$$ \tilde{\bm{q}}_m^{(i)} = \begin{pmatrix} q_{2i}\cos(m\theta_i) – q_{2i+1}\sin(m\theta_i) \\ q_{2i}\sin(m\theta_i) + q_{2i+1}\cos(m\theta_i) \end{pmatrix} $$

$$ \tilde{\bm{k}}_n^{(i)} = \begin{pmatrix} k_{2i}\cos(n\theta_i) – k_{2i+1}\sin(n\theta_i) \\ k_{2i}\sin(n\theta_i) + k_{2i+1}\cos(n\theta_i) \end{pmatrix} $$

$i$ 番目のペアの内積は

$$ \begin{align} (\tilde{\bm{q}}_m^{(i)})^T \tilde{\bm{k}}_n^{(i)} &= [q_{2i}\cos(m\theta_i) – q_{2i+1}\sin(m\theta_i)][k_{2i}\cos(n\theta_i) – k_{2i+1}\sin(n\theta_i)] \\ &\quad + [q_{2i}\sin(m\theta_i) + q_{2i+1}\cos(m\theta_i)][k_{2i}\sin(n\theta_i) + k_{2i+1}\cos(n\theta_i)] \end{align} $$

第1項を展開します。

$$ \begin{align} &q_{2i}k_{2i}\cos(m\theta_i)\cos(n\theta_i) – q_{2i}k_{2i+1}\cos(m\theta_i)\sin(n\theta_i) \\ &- q_{2i+1}k_{2i}\sin(m\theta_i)\cos(n\theta_i) + q_{2i+1}k_{2i+1}\sin(m\theta_i)\sin(n\theta_i) \end{align} $$

第2項を展開します。

$$ \begin{align} &q_{2i}k_{2i}\sin(m\theta_i)\sin(n\theta_i) + q_{2i}k_{2i+1}\sin(m\theta_i)\cos(n\theta_i) \\ &+ q_{2i+1}k_{2i}\cos(m\theta_i)\sin(n\theta_i) + q_{2i+1}k_{2i+1}\cos(m\theta_i)\cos(n\theta_i) \end{align} $$

第1項と第2項を加算すると、多くの交差項が相殺されます。

$q_{2i}k_{2i}$ の係数:

$$ \cos(m\theta_i)\cos(n\theta_i) + \sin(m\theta_i)\sin(n\theta_i) = \cos((m-n)\theta_i) $$

$q_{2i+1}k_{2i+1}$ の係数:

$$ \sin(m\theta_i)\sin(n\theta_i) + \cos(m\theta_i)\cos(n\theta_i) = \cos((m-n)\theta_i) $$

$q_{2i}k_{2i+1}$ の係数:

$$ -\cos(m\theta_i)\sin(n\theta_i) + \sin(m\theta_i)\cos(n\theta_i) = \sin((m-n)\theta_i) $$

$q_{2i+1}k_{2i}$ の係数:

$$ -\sin(m\theta_i)\cos(n\theta_i) + \cos(m\theta_i)\sin(n\theta_i) = -\sin((m-n)\theta_i) $$

まとめると

$$ (\tilde{\bm{q}}_m^{(i)})^T \tilde{\bm{k}}_n^{(i)} = (q_{2i}k_{2i} + q_{2i+1}k_{2i+1})\cos((m-n)\theta_i) + (q_{2i}k_{2i+1} – q_{2i+1}k_{2i})\sin((m-n)\theta_i) $$

これは相対位置 $m – n$ のみの関数であり、絶対位置 $m$, $n$ には依存しません。

全次元の内積は

$$ \tilde{\bm{q}}_m^T \tilde{\bm{k}}_n = \sum_{i=0}^{d/2-1} (\tilde{\bm{q}}_m^{(i)})^T \tilde{\bm{k}}_n^{(i)} $$

各ペアが相対位置のみに依存するため、全体の内積も相対位置 $m – n$ のみに依存します。$\blacksquare$

RoPE の効率的な実装

回転行列を陽に構築する代わりに、複素数の積として実装できます。$\bm{q}^{(i)}$ を複素数 $q_{2i} + q_{2i+1}j$ と見なすと

$$ \tilde{q}_m^{(i)} = (q_{2i} + q_{2i+1}j) \cdot e^{jm\theta_i} = (q_{2i} + q_{2i+1}j)(\cos(m\theta_i) + j\sin(m\theta_i)) $$

実部と虚部を分離すると、先ほどの回転行列の適用と一致します。この複素数表現により、要素ごとの乗算で効率的に計算できます。

ALiBi(Attention with Linear Biases)

Press et al.(2022)が提案した ALiBi は、位置エンコーディングを入力に加算するのではなく、Attention スコアに線形バイアスを直接加算する手法です。

定式化

ALiBi では、Attention スコアに位置依存のバイアス項を加えます。

$$ \text{Attention}(\bm{Q}, \bm{K}, \bm{V}) = \text{softmax}\left(\frac{\bm{Q}\bm{K}^T}{\sqrt{d_k}} + m \cdot \bm{B}\right)\bm{V} $$

バイアス行列 $\bm{B}$ は

$$ B_{ij} = -|i – j| $$

つまり、相対距離に比例したペナルティを与えます。$m$ はヘッドごとのスカラーで、$h$ ヘッドの場合

$$ m_h = \frac{1}{2^{8h/H}}, \quad h = 1, 2, \dots, H $$

と等比数列で設定されます($H$ は総ヘッド数)。例えば8ヘッドの場合、$m$ の値は $\frac{1}{2}, \frac{1}{4}, \frac{1}{8}, \dots, \frac{1}{256}$ です。

直感的理解

ALiBi は「遠い位置ほど Attention スコアにペナルティを与える」という非常にシンプルな仕組みです。

  • ヘッドごとに異なる $m$ を使うことで、局所的な注目(大きな $m$)と大局的な注目(小さな $m$)を同時に捉える
  • 学習可能なパラメータが追加されない
  • 入力に位置ベクトルを加算しないため、埋め込み空間を圧迫しない

長文外挿能力

ALiBi の最大の利点は学習時よりも長い系列への外挿能力です。学習時に長さ $L$ までしか見ていなくても、推論時に $L’ > L$ の系列に対して自然に対応できます。これは、バイアス $-|i-j|$ の構造が系列長に依存しないためです。

対照的に、正弦波 PE は理論的には外挿可能ですが実際には性能が劣化しやすく、学習可能 PE は学習範囲外の位置に全く対応できません。

各手法の比較

特性 正弦波 PE 学習可能 PE Shaw et al. RoPE ALiBi
追加パラメータ なし $T_{\max} \times d$ $2C \times d_k$ なし なし
位置情報の注入先 入力に加算 入力に加算 Attention スコア Q, K に回転 Attention スコア
絶対/相対 絶対(相対性あり) 絶対 相対 相対 相対
長文外挿 理論的に可能 不可能 クリップ範囲内 良好 非常に良好
採用例 Transformer (原論文) BERT, GPT-2, ViT Transformer-XL LLaMA, Qwen BLOOM, MPT
実装の簡便さ 容易 容易 やや複雑 中程度 非常に容易

Python での実装と可視化

各種位置エンコーディングの実装

import numpy as np
import matplotlib.pyplot as plt

def sinusoidal_pe(max_len, d_model):
    """正弦波位置エンコーディング"""
    PE = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
    PE[:, 0::2] = np.sin(position * div_term)
    PE[:, 1::2] = np.cos(position * div_term)
    return PE

def learned_pe(max_len, d_model, seed=42):
    """学習可能な位置エンコーディング(ランダム初期化で模擬)"""
    np.random.seed(seed)
    return np.random.randn(max_len, d_model) * 0.02

def rope_apply(q, k, positions_q, positions_k, d_model):
    """RoPE を Q, K に適用する"""
    # 周波数の計算
    theta = 10000.0 ** (-np.arange(0, d_model, 2) / d_model)  # (d/2,)

    def rotate(x, positions):
        """位置に応じた回転を適用"""
        x_rot = np.zeros_like(x)
        for idx, pos in enumerate(positions):
            angles = pos * theta  # (d/2,)
            cos_a = np.cos(angles)
            sin_a = np.sin(angles)
            # 偶数次元と奇数次元に分離
            x_even = x[idx, 0::2]
            x_odd = x[idx, 1::2]
            x_rot[idx, 0::2] = x_even * cos_a - x_odd * sin_a
            x_rot[idx, 1::2] = x_even * sin_a + x_odd * cos_a
        return x_rot

    q_rot = rotate(q, positions_q)
    k_rot = rotate(k, positions_k)
    return q_rot, k_rot

def alibi_bias(T, n_heads):
    """ALiBi バイアス行列を生成"""
    # 各ヘッドの傾き
    slopes = 1.0 / (2.0 ** (8.0 * np.arange(1, n_heads + 1) / n_heads))

    # 相対距離行列
    positions = np.arange(T)
    rel_dist = -np.abs(positions[:, np.newaxis] - positions[np.newaxis, :])  # (T, T)

    # 各ヘッドのバイアス
    biases = np.zeros((n_heads, T, T))
    for h in range(n_heads):
        biases[h] = slopes[h] * rel_dist

    return biases


# パラメータ設定
max_len = 64
d_model = 32
n_heads = 4

位置エンコーディングの比較可視化

import numpy as np
import matplotlib.pyplot as plt

def sinusoidal_pe(max_len, d_model):
    PE = np.zeros((max_len, d_model))
    position = np.arange(max_len)[:, np.newaxis]
    div_term = np.exp(np.arange(0, d_model, 2) * (-np.log(10000.0) / d_model))
    PE[:, 0::2] = np.sin(position * div_term)
    PE[:, 1::2] = np.cos(position * div_term)
    return PE

max_len = 64
d_model = 32

# 正弦波 PE の生成
sin_pe = sinusoidal_pe(max_len, d_model)

fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# (1) 正弦波 PE のヒートマップ
im1 = axes[0, 0].imshow(sin_pe, cmap='RdBu', aspect='auto', interpolation='nearest')
axes[0, 0].set_xlabel('Dimension')
axes[0, 0].set_ylabel('Position')
axes[0, 0].set_title('Sinusoidal PE')
plt.colorbar(im1, ax=axes[0, 0])

# (2) 正弦波 PE の内積行列(相対位置依存性)
dot_sin = sin_pe @ sin_pe.T
im2 = axes[0, 1].imshow(dot_sin, cmap='viridis', aspect='auto')
axes[0, 1].set_xlabel('Position j')
axes[0, 1].set_ylabel('Position i')
axes[0, 1].set_title('Sinusoidal PE: dot(PE(i), PE(j))')
plt.colorbar(im2, ax=axes[0, 1])

# (3) 特定の次元での正弦波の様子
positions = np.arange(max_len)
for dim_idx in [0, 4, 8, 16]:
    axes[1, 0].plot(positions, sin_pe[:, dim_idx],
                    label=f'dim={dim_idx}', linewidth=1.5)
axes[1, 0].set_xlabel('Position')
axes[1, 0].set_ylabel('PE value')
axes[1, 0].set_title('Sinusoidal PE: Selected Dimensions')
axes[1, 0].legend()
axes[1, 0].grid(True, alpha=0.3)

# (4) 相対位置 k に対する内積の値
ref_pos = 0
inner_products = [np.dot(sin_pe[ref_pos], sin_pe[k]) for k in range(max_len)]
axes[1, 1].plot(range(max_len), inner_products, 'b-', linewidth=2)
axes[1, 1].set_xlabel('Relative Position k')
axes[1, 1].set_ylabel('dot(PE(0), PE(k))')
axes[1, 1].set_title('Inner Product vs Relative Position')
axes[1, 1].grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

RoPE の効果の可視化

import numpy as np
import matplotlib.pyplot as plt

def rope_attention_scores(seq_len, d_model, seed=42):
    """RoPE を適用した Attention スコアを計算"""
    np.random.seed(seed)
    theta = 10000.0 ** (-np.arange(0, d_model, 2) / d_model)

    # ランダムな Q, K
    Q = np.random.randn(seq_len, d_model) * 0.1
    K = np.random.randn(seq_len, d_model) * 0.1

    # RoPE 適用
    Q_rot = np.zeros_like(Q)
    K_rot = np.zeros_like(K)
    for pos in range(seq_len):
        angles = pos * theta
        cos_a = np.cos(angles)
        sin_a = np.sin(angles)
        Q_rot[pos, 0::2] = Q[pos, 0::2] * cos_a - Q[pos, 1::2] * sin_a
        Q_rot[pos, 1::2] = Q[pos, 0::2] * sin_a + Q[pos, 1::2] * cos_a
        K_rot[pos, 0::2] = K[pos, 0::2] * cos_a - K[pos, 1::2] * sin_a
        K_rot[pos, 1::2] = K[pos, 0::2] * sin_a + K[pos, 1::2] * cos_a

    # Attention スコア
    scores_no_rope = Q @ K.T / np.sqrt(d_model)
    scores_rope = Q_rot @ K_rot.T / np.sqrt(d_model)
    return scores_no_rope, scores_rope

def alibi_bias(T, n_heads):
    slopes = 1.0 / (2.0 ** (8.0 * np.arange(1, n_heads + 1) / n_heads))
    positions = np.arange(T)
    rel_dist = -np.abs(positions[:, np.newaxis] - positions[np.newaxis, :])
    biases = np.zeros((n_heads, T, T))
    for h in range(n_heads):
        biases[h] = slopes[h] * rel_dist
    return biases

seq_len = 32
d_model = 64

scores_no_rope, scores_rope = rope_attention_scores(seq_len, d_model)

fig, axes = plt.subplots(2, 2, figsize=(14, 10))

# (1) RoPE なしの Attention スコア
im1 = axes[0, 0].imshow(scores_no_rope, cmap='RdBu', aspect='auto',
                         vmin=-2, vmax=2)
axes[0, 0].set_title('Attention Scores (No PE)')
axes[0, 0].set_xlabel('Key Position')
axes[0, 0].set_ylabel('Query Position')
plt.colorbar(im1, ax=axes[0, 0])

# (2) RoPE ありの Attention スコア
im2 = axes[0, 1].imshow(scores_rope, cmap='RdBu', aspect='auto',
                         vmin=-2, vmax=2)
axes[0, 1].set_title('Attention Scores (with RoPE)')
axes[0, 1].set_xlabel('Key Position')
axes[0, 1].set_ylabel('Query Position')
plt.colorbar(im2, ax=axes[0, 1])

# (3) ALiBi バイアス(各ヘッド)
n_heads = 4
biases = alibi_bias(seq_len, n_heads)
for h in range(n_heads):
    axes[1, 0].plot(range(seq_len), biases[h, 0, :],
                    label=f'Head {h+1} (m={1/(2**(8*(h+1)/n_heads)):.4f})',
                    linewidth=2)
axes[1, 0].set_xlabel('Key Position (Query at pos 0)')
axes[1, 0].set_ylabel('Bias')
axes[1, 0].set_title('ALiBi Bias per Head')
axes[1, 0].legend(fontsize=8)
axes[1, 0].grid(True, alpha=0.3)

# (4) ALiBi バイアス行列(ヘッド1)
im4 = axes[1, 1].imshow(biases[0], cmap='viridis', aspect='auto')
axes[1, 1].set_title(f'ALiBi Bias Matrix (Head 1, m={1/(2**(8*1/n_heads)):.4f})')
axes[1, 1].set_xlabel('Key Position')
axes[1, 1].set_ylabel('Query Position')
plt.colorbar(im4, ax=axes[1, 1])

plt.tight_layout()
plt.show()

RoPE の相対位置依存性の検証

import numpy as np
import matplotlib.pyplot as plt

def verify_rope_relative(d_model=64, max_offset=50, seed=42):
    """RoPE の内積が相対位置のみに依存することを数値的に検証"""
    np.random.seed(seed)
    theta = 10000.0 ** (-np.arange(0, d_model, 2) / d_model)

    # 固定された q, k ベクトル
    q = np.random.randn(d_model)
    k = np.random.randn(d_model)

    def apply_rope_1d(x, pos):
        """1つのベクトルに RoPE を適用"""
        x_rot = np.zeros_like(x)
        angles = pos * theta
        x_rot[0::2] = x[0::2] * np.cos(angles) - x[1::2] * np.sin(angles)
        x_rot[1::2] = x[0::2] * np.sin(angles) + x[1::2] * np.cos(angles)
        return x_rot

    # 様々な絶対位置で相対位置 k の内積を計算
    offsets = range(max_offset)
    results = {}

    for base_pos in [0, 10, 25, 50]:
        inner_prods = []
        for k_offset in offsets:
            q_rot = apply_rope_1d(q, base_pos)
            k_rot = apply_rope_1d(k, base_pos + k_offset)
            inner_prods.append(np.dot(q_rot, k_rot))
        results[base_pos] = inner_prods

    return offsets, results

offsets, results = verify_rope_relative()

plt.figure(figsize=(10, 6))
for base_pos, inner_prods in results.items():
    plt.plot(offsets, inner_prods, linewidth=2, label=f'base_pos={base_pos}')

plt.xlabel('Relative Position k', fontsize=12)
plt.ylabel('Inner Product', fontsize=12)
plt.title('RoPE: Inner Product vs Relative Position\n'
          '(Lines overlap, confirming dependence only on relative position)',
          fontsize=13)
plt.legend(fontsize=11)
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()

この図では、異なる絶対位置(base_pos = 0, 10, 25, 50)に対して、相対位置 $k$ の関数として内積をプロットしています。全ての曲線が完全に重なることから、RoPE を適用した内積が相対位置のみに依存することが数値的に確認できます。

まとめ

本記事では、Transformer における位置エンコーディングの各種手法について、理論的な導出から実装まで解説しました。

  • Self-Attention の置換不変性: 置換行列を用いた厳密な証明により、位置情報の注入が不可欠であることを示した
  • 正弦波 PE: 相対位置の内積表現と線形変換可能性という優れた数学的性質を持つ
  • 学習可能 PE: データ駆動で最適化されるが外挿能力に欠ける
  • 相対位置エンコーディング(Shaw et al.): Attention スコアに相対位置情報を直接注入する
  • RoPE: 回転行列により内積が自然に相対位置のみに依存し、LLaMA 等で広く採用されている
  • ALiBi: 線形バイアスという最もシンプルなアプローチで、優れた長文外挿能力を実現する

次のステップとして、以下の記事も参考にしてください。