2D回転がTransformerの長文脈問題を解決した理由：RoPEのメカニクス重視の考察

Attention機構は位置エンコーディングに2つの要件を求めます：

  1. 相対距離の感度：トークン ii がトークン jj に注意を払う際には、グローバルな文脈内での位置ではなく、主に2つのトークンの距離（ i−ji - j ）に依存すべきです。
  2. 特徴量の保持：位置情報の注入は、モデルが学習した意味的埋め込み特徴を破壊したり損なったりしてはなりません。

元の絶対位置エンコーディングは、静的なsin/cos波を入力埋め込みに直接加算していました：

xi=ei+pi \mathbf{x}_i = \mathbf{e}_i + \mathbf{p}_i

これにより、モデルは意味的意味（ ei\mathbf{e}i ）と位置（ pi\mathbf{p}_i ）を分離するためにパラメータ容量を浪費せざるを得ませんでした。さらに、最大系列長 N=2048N = 2048 で学習されたモデルにとって、位置 20492049 は未学習のベクトル p2049\mathbf{p}{2049} となり、生成の即時崩壊を引き起こしていました。

その後のアプローチでは、相対バイアス項 bi,jb_{i,j} を直接アテンションマトリクス QKT+BQK^T + B に加算する方法が試みられました。機能的には有効でしたが、 N×NN \times N 行列の修正はFlashAttentionのようなカーネルレベルのGPU融合を阻害し、巨大なメモリオーバーヘッドを生み出しました。

そこで登場したのが Rotary Position Embedding (RoPE)（Su et al., 2021）です。位置ベクトルを加算したり、アテンションマトリクスを修正したりする代わりに、RoPEは注意計算の前にQueryとKeyベクトルを2D部分平面内で回転させます。

幾何学：回転下での内積

注意を相対的にするためには、位置 mm におけるベクトル x\mathbf{x} に適用するエンコード関数 R(x,m)R(\mathbf{x}, m) を考えます。これにより、位置 mm のQueryと位置 nn のKeyの内積が相対オフセット（ m−nm - n ）のみに依存するようにします：

⟨R(q,m),R(k,n)⟩=g(q,k,m−n) \langle R(\mathbf{q}, m), R(\mathbf{k}, n) \rangle = g(\mathbf{q}, \mathbf{k}, m - n)

ベクトルノルムを保持しつつ角度シフトをエンコードするにはどうすればよいでしょうか？ 複素空間での回転です。

2D平面において、ベクトル x=[x1,x2]T\mathbf{x} = [x_1, x_2]^T を角度 mθm\theta で回転させることは、直交回転行列で表現されます：

RΘ,m2=(cos⁡mθ−sin⁡mθ sin⁡mθcos⁡mθ) R_{\Theta, m}^2 = \begin{pmatrix} \cos m\theta & -\sin m\theta \ \sin m\theta & \cos m\theta \end{pmatrix}

位置 mm で回転されたQueryと位置 nn で回転されたKeyの内積を計算すると：

(RΘ,m2q)T(RΘ,n2k)=qT(RΘ,m2)TRΘ,n2k=qTRΘ,n−m2k \left(R_{\Theta, m}^2 \mathbf{q}\right)^T \left(R_{\Theta, n}^2 \mathbf{k}\right) = \mathbf{q}^T \left(R_{\Theta, m}^2\right)^T R_{\Theta, n}^2 \mathbf{k} = \mathbf{q}^T R_{\Theta, n - m}^2 \mathbf{k}

なぜなら RT(m)R(n)=R(n−m)R^T(m) R(n) = R(n - m) であるため、絶対位置 mm と nn は完全に相殺されます。 結果として得られる注意の重みは、距離（ m−nm - n ）の関数となります。

dd 次元のベクトルに対して、RoPEはチャネルを d/2d/2 個の2D平面ペアに分割し、各ペアに異なる回転周波数 θi\theta_i を適用します：

Θ=θi=10000−2(i−1)/d,i∈[1,2,…,d/2] \Theta = {\theta_i = 10000^{-2(i-1)/d}, \quad i \in [1, 2, \dots, d/2]}

高性能ベクトル化PyTorch実装

実際には、すべてのヘッドとトークンに対して完全なブロック対角回転行列を明示的に構築するのは非効率です。複素数表現または実数値スライストリックを用いて、RoPEを効率的に実装できます：

RΘ,mx=x⊙cos⁡(mΘ)+x~⊙sin⁡(mΘ) R_{\Theta, m} \mathbf{x} = \mathbf{x} \odot \cos(m\Theta) + \tilde{\mathbf{x}} \odot \sin(m\Theta)

ここで x~=[−x2,x1,−x4,x3,… ]\tilde{\mathbf{x}} = [-x_2, x_1, -x_4, x_3, \dots] です。

以下は、クリーンで本番環境向けのPyTorch実装です：


python
import torch
import torch.nn as nn

class RotaryPositionalEmbedding(nn.Module):
    """
    Rotary Position Embedding (RoPE) for multi-head attention.
    Applies 2D plane rotations across d_model / 2 feature pairs.
    """
    def __init__(self, dim: int, max_seq_len: int = 8192, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.base = base

        # Compute frequency bands theta_i = 10000^(-2(i-1)/d)
        inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Precompute cosine and sine cache for max_seq_len
        self._build_cache(max_seq_len)

    def _build_cache(self, max_seq_len: int):
        t = torch.arange(max_seq_len, dtype=self.inv_freq.dtype)
        # Outer product: [seq_len, dim / 2]
        freqs = torch.outer(t, self.inv_freq)

        # Duplicate frequencies to match feature dimensions: [seq_len, dim]
        emb = torch.cat((freqs, freqs), dim=-1)

        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def _rotate_half(self, x: torch.Tensor) -> torch.Tensor:
        """Splits vector in half and negates/swaps halves: [-x2, x1]"""
        x1 = x[..., : self.dim // 2]
        x2 = x[..., self.dim // 2 :]
        return torch.cat((-x2, x1), dim=-1)

    def forward(self, x: torch.Tensor, seq_len: int) -> torch.Tensor:
        """
        Args:
            x: Tensor of shape [batch_size, num_heads, seq_len, head_dim]
            seq_len: Current sequence length
        Returns:
            Tensor with RoPE applied to Query or Key
        """
        cos = self.cos_cached[:seq_len, :].unsqueeze(0).unsqueeze(1) # [1, 1, seq_len, dim]
        sin = self.sin_cached[:seq_len, :].unsqueeze(0).unsqueeze(1) # [1, 1, seq_len, dim]

        # Apply elementwise rotation formula
        return (x * cos) + (self._rotate_half(x) * sin)

Enter fullscreen mode Exit fullscreen mode