Why 2D Rotations Solved Transformer Long-Context: A Mechanics-First Look at RoPE

Attention 需要位置編碼提供兩項特性：

  1. 相對距離敏感度：Token ii 對 Token jj 的注意力應主要取決於兩者之間的距離（ i−ji - j ），而非在上下文窗口中的絕對位置。
  2. 特徵保留：注入位置資訊時，不應破壞或扭曲模型已學習到的語義嵌入特徵。

原始的絕對位置編碼是將靜態正弦/餘弦波直接加到輸入嵌入上：

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} ，導致生成過程立即崩潰。

後續方法嘗試直接在注意力矩陣 QKT+BQK^T + B 中加入相對偏差項 bi,jb_{i,j} 。雖然在功能上有效，但修改 N×NN \times N 矩陣破壞了 FlashAttention 等 GPU 核心融合，同時帶來龐大的記憶體開銷。

接著登場的是 Rotary Position Embedding (RoPE)（Su et al., 2021）。RoPE 並非添加位置向量或修補注意力矩陣，而是先對 Query 與 Key 向量在 2D 子平面上進行旋轉，再計算注意力。

The Geometry: Inner Products Under Rotation

為了讓注意力具有相對性，我們希望有一個編碼函數 R(x,m)R(\mathbf{x}, m) ，應用於位置 mm 的向量 x\mathbf{x} ，使得位置 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]}

High-Performance Vectorized PyTorch Implementation

實際上，為每個 head 與 token 顯式建構完整的區塊對角旋轉矩陣效率不佳。我們可利用複數表示或實數切片技巧來高效實作 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