rope笔记

DINOv3 中的 RoPE 笔记

对应代码：
dinov3/layers/rope_position_encoding.py — RopePositionEmbedding 类
dinov3/layers/attention.py — apply_rope / rope_apply / rope_rotate_half

1. RoPE 的标准公式

对长度为 D 的特征向量 $\mathbf{x}$，RoPE 把它切成 $D/2$ 个二维子向量 $[\,x_{2i},\; x_{2i+1}\,]$，每个子向量用一个角度 $\theta_i$ 做 2D 旋转：

$$ \begin{pmatrix} x_{2i}’ \ x_{2i+1}’ \end{pmatrix}

\begin{pmatrix} \cos\theta_i & -\sin\theta_i \ \sin\theta_i & \cos\theta_i \end{pmatrix} \begin{pmatrix} x_{2i} \ x_{2i+1} \end{pmatrix} $$

1.1 角度 $\theta_i$ 的几何级数构造

$\theta_i$ 不是任意选的，而是按几何级数从大到小排列（高频 → 低频），让不同维度子空间捕捉不同尺度的位置信息：

$$ \theta_i = \text{base}^{-2i/D}, \quad i = 0, 1, \ldots, D/2 - 1 $$

$\text{base}$ 是频率基数：原始 RoPE 论文（Su et al. 2021）取 $\text{base}=10000$，DINOv3 默认为 $\text{base}=100$（见 §4.3）
$i$ 越大，$\theta_i$ 越小 → 旋转越慢 → 越"低频"，对应长程依赖
$i$ 越小，$\theta_i$ 越大 → 旋转越快 → 越"高频"，对应短程依赖

1.2 $\cos$ 和 $\sin$ 的显式公式

对于位置为 $m$ 的 token，第 $i$ 对维度的实际旋转角是 $m \cdot \theta_i$，因此

$$ \boxed{\;\cos\theta_i(m) = \cos\!\left(m \cdot \theta_i\right) = \cos\!\left(\frac{2\pi\, m}{\text{base}^{2i/D}}\right)\;} $$$$ \boxed{\;\sin\theta_i(m) = \sin\!\left(m \cdot \theta_i\right) = \sin\!\left(\frac{2\pi\, m}{\text{base}^{2i/D}}\right)\;} $$

两种等价写法（与 §4 的 DINOv3 实现一致）：
频率式（原始论文）：$\theta_i = \text{base}^{-2i/D}$，角度 $= m\theta_i$
周期式（DINOv3 代码）：周期 $T_i = \text{base}^{2i/D}$，角度 $= 2\pi m / T_i$
二者完全等价：$m\theta_i = m \cdot \text{base}^{-2i/D} = 2\pi m / \text{base}^{2i/D}$（差一个 $2\pi$ 因子仅影响"一个完整旋转对应多少位置"的刻度，不改变相对编码性质）。

DINOv3 用了**“半切分"变体**——不是配对 $(x_{2i}, x_{2i+1})$，而是配对 $(x_i,\; x_{i+D/2})$。两种配对在数学上表达能力等价（见 §5）。

2. `apply_rope` 拆解

整个函数就两件事：拼接 prefix + 真正的旋转。

def apply_rope(self, q, k, rope):
    q_dtype, k_dtype = q.dtype, k.dtype
    sin, cos = rope
    rope_dtype = sin.dtype
    q = q.to(dtype=rope_dtype)
    k = k.to(dtype=rope_dtype)

    N = q.shape[-2]                     # 序列总长 = 1 (CLS) + M (storage) + HW
    prefix = N - sin.shape[-2]          # 1 + M：prefix token 数量

    # prefix 部分（CLS / storage）不参与旋转
    q_prefix = q[:, :, :prefix, :]
    q = rope_apply(q[:, :, prefix:, :], sin, cos)
    q = torch.cat((q_prefix, q), dim=-2)

    k_prefix = k[:, :, :prefix, :]
    k = rope_apply(k[:, :, prefix:, :], sin, cos)
    k = torch.cat((k_prefix, k), dim=-2)

    q = q.to(dtype=q_dtype)
    k = k.to(dtype=k_dtype)
    return q, k

sin/cos 的形状 [HW, D] 中 HW 来自 RopePositionEmbedding.forward() 输出的图像网格大小（H×W 个 patch），所以 prefix 就是”非 patch“的 token 数量（CLS + storage）。这些 token 不参与旋转，因为它们没有空间位置。

3. 真正的旋转：`rope_apply`

def rope_apply(x, sin, cos):
    return (x * cos) + (rope_rotate_half(x) * sin)

关键 trick：把 2D 旋转"伪装"成 elementwise 乘法 + 一个特殊的半切排列。

def rope_rotate_half(x):
    # x:   [ x0  x1  x2 | x3  x4  x5]
    # out: [-x3 -x4 -x5 | x0  x1  x2]
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

设 $\mathbf{x} = [x_0, \ldots, x_{D-1}]$，则 rotate_half(x)[i]：

当 $i \in [0, D/2)$：等于 $-x_{i+D/2}$
当 $i \in [D/2, D)$：等于 $x_{i-D/2}$

代入 rope_apply：

$i \in [0, D/2)$：$x'_i = x_i \cos\theta_i - x_{i+D/2} \sin\theta_i$
$i \in [D/2, D)$：$x'_i = x_i \cos\theta_i + x_{i-D/2} \sin\theta_i$

这正是"配对 $(i, i+D/2)$ 的 2D 旋转"展开后的标量形式。

4. `sin` / `cos` 的构造

RopePositionEmbedding.forward（rope_position_encoding.py:57-106）：

angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :]  # [HW, 2, D/4]
angles = angles.flatten(1, 2)  # [HW, D/2]
angles = angles.tile(2)        # [HW, D]
cos = torch.cos(angles)  # [HW, D]
sin = torch.sin(angles)  # [HW, D]

4.1 `flatten(1, 2)` — 把 2D 坐标的频率展平

coords 形状 [HW, 2]（h、w 两个坐标），periods 形状 [D/4]。广播后每个 (patch, 坐标, 频率) 三元组都得到一个角度，展平后每个 patch 有 $2 \times D/4 = D/2$ 个独立角度。

4.2 `tile(2)` — 让 sin/cos 在两个半区"成对重复”

# angles:  [θ_0, θ_1, ..., θ_{D/2-1}]              # 长度 D/2
# angles.tile(2) = [θ_0, ..., θ_{D/2-1}, θ_0, ..., θ_{D/2-1}]  # 长度 D

这样 cos[i] == cos[i+D/2]、sin[i] == sin[i+D/2]，恰好保证配对 $(i, i+D/2)$ 用同一个角度 $\theta_i$ 旋转。

4.3 周期 `periods` 的几何意义

_init_weights（rope_position_encoding.py:108-121）：

periods = self.base ** (2 * torch.arange(D//4) / (D//2))

含义：对于第 $k$ 个频率分量，旋转角 $\theta_k = 2\pi \cdot c / T_k$，其中：

$c \in [-1, 1]$ 是归一化空间坐标
$T_k = \text{base}^{2k/(D/2)}$（几何级数周期，与原始 RoPE 论文一致）

DINOv3 默认 base=100（经典实现是 10000）。

5. 整条链路：代码 ↔ 公式对照

代码片段	公式含义
`periods`	$T_k = \text{base}^{2k/(D/2)}$（几何级数）
`coords ∈ [-1, 1]` 网格	归一化空间坐标 $c \in [-1, 1]$
`angles = 2π·c/T`	每个 (坐标, 频率) 对的旋转角 $\theta$
`tile(2)`	让 $i$ 和 $i+D/2$ 共享同一 $\theta$
`rope_rotate_half`	实现配对 $(i, i+D/2)$ 的"交叉取反"
`xcos + rot(x)sin`	2D 旋转的 elementwise 形式
`prefix = N - HW`	CLS / storage token 不参与旋转

6. 具体例子（D=6）

设 $\mathbf{x} = [x_0, x_1, x_2, x_3, x_4, x_5]$，$\theta = [\theta_0, \theta_1, \theta_2]$。

cos 形状 [6]：[c_0, c_1, c_2, c_0, c_1, c_2] sin 形状 [6]：[s_0, s_1, s_2, s_0, s_1, s_2] rotate_half(x)：$[-x_3, -x_4, -x_5, x_0, x_1, x_2]$

于是：

$$ \begin{aligned} x'_0 &= x_0 c_0 - x_3 s_0 & x'_3 &= x_3 c_0 + x_0 s_0 \quad \text{（配对 } (0, 3)，\theta_0）\\ x'_1 &= x_1 c_1 - x_4 s_1 & x'_4 &= x_4 c_1 + x_1 s_1 \quad \text{（配对 } (1, 4)，\theta_1）\\ x'_2 &= x_2 c_2 - x_5 s_2 & x'_5 &= x_5 c_2 + x_2 s_2 \quad \text{（配对 } (2, 5)，\theta_2）\\ \end{aligned} $$

正是三个独立的 2D 旋转，作用在三对维度上。

7. 半切分 vs 标准公式的"等价性"

7.1 两种变体的区别

设 $D=8$，4 个独立角度 $\theta_0, \theta_1, \theta_2, \theta_3$。

标准 RoPE（Su et al. 2021）：配对相邻维度 $(x_0, x_1)_{\theta_0}, (x_2, x_3)_{\theta_1}, (x_4, x_5)_{\theta_2}, (x_6, x_7)_{\theta_3}$
DINOv3 半切分：配对相隔 $D/2$ 的维度 $(x_0, x_4)_{\theta_0}, (x_1, x_5)_{\theta_1}, (x_2, x_6)_{\theta_2}, (x_3, x_7)_{\theta_3}$

关键点：两种变体对同一像素位置编码——位置由 sin/cos 的行索引（即 patch 索引）决定，与特征向量里哪两个维度被配对无关。“参考像素位置"是相同的，真正不同的是"角度 $\theta_i$ 被分配给哪一对特征维度”。

7.2 为什么说它们在"数学上等价"？

(a) 角度的"重排"等价

标准 RoPE 的 sin/cos 排列：[θ_0, θ_0, θ_1, θ_1, θ_2, θ_2, θ_3, θ_3]（每个角度连续重复 2 次）
半切分：[θ_0, θ_1, θ_2, θ_3, θ_0, θ_1, θ_2, θ_3]（角度均匀分布，前后半重复）

两种排列都是同一组 4 个独立角度的某种排列。也就是"把 $\theta$ 装进 sin/cos 张量的方式"不同，但角度集合相同。

(b) 用一个固定置换矩阵可以把两者对齐

存在一个固定置换 $P \in \{0,1\}^{D \times D}$（与位置无关），使得

$$ \text{RoPE}_{\text{half}}(\mathbf{x}, \theta) = P \cdot \text{RoPE}_{\text{standard}}(P^T \mathbf{x}, P_\theta \theta) $$

具体做法：对输入做置换

$$ \mathbf{x}' = [x_0, x_4, x_1, x_5, x_2, x_6, x_3, x_7] $$

（即按"步长 $D/2$“交错），那么 $(x_0, x_4) \to (x'_0, x'_1)$、$(x_1, x_5) \to (x'_2, x'_3)$、…，正好变成标准配对。

(c) 关键不变性质：相对位置编码

RoPE 的核心数学性质是：

$$ (R_m \mathbf{q})^T (R_n \mathbf{k}) = \mathbf{q}^T R_{n-m} \mathbf{k} $$

两种变体都满足这个性质（因为它们都是一族正交矩阵乘法，与配对方式无关）。注意力分数只依赖于 query 和 key 的相对位置 $n-m$，而不是绝对位置。

7.3 在深度学习实践中为什么可以"互换”？

网络里紧跟着 RoPE 的是 $W_q, W_k$ 线性层。对 $W_q$ 输出维度的固定置换 $\pi$，可以等价地通过置换 $W_q$ 的行索引来实现：

$$ W_q \cdot \text{RoPE}_{\text{half}}(\mathbf{x}) = W'_q \cdot \text{RoPE}_{\text{standard}}(\mathbf{x}) $$

其中 $W'_q$ 是把 $W_q$ 的行按 $\pi$ 重排后的矩阵。

因此：

在初始化时（$W_q, W_k$ 各向同性），两种变体在统计意义上等价——它们表达的是同一族位置相关的正交变换
在训练后，如果两个模型用相同的角度 $\theta$ 集合、相同的数据训，理论上一个能做到的事情，另一个也能做到（容量相同），只是具体的 $W_q, W_k$ 矩阵会学到的"排列"不同
实际的 $W_q, W_k$ 没有任何归纳偏置要求某个特定维度子空间对应"高频角度"或"低频角度"，所以哪种排列都行

7.4 一个"非等价"的细微之处

两种变体对梯度的流动方式不同。考虑一个特征维度 $x_i$：

标准变体：$x_i$ 的梯度只会流到配对的那一维 $x_{i \pm 1}$
半切分：$x_i$ 的梯度会流到相隔 $D/2$ 的 $x_{i \pm D/2}$

这会影响：

训练早期的优化动力学（梯度通过 RoPE 回传时跨越的距离不同）
不同频率角度学习到的"语义"（可能不同维度子空间最终承载不同信息）

但这不影响"等价性"的实质——这只是优化路径差异，不是表达能力差异。

7.5 小结

两种变体不是 bit-exact 等价，但在以下三个层面等价：
表达同一族变换：都是 $D/2$ 个独立 2D 旋转，使用相同角度集合
满足 RoPE 的核心性质：注意力分数只依赖相对位置
可被吸收到 $W_q, W_k$：差异是固定的维度置换，可在训练中被线性层吸收
“参考像素位置"是一样的（都是行索引），真正不同的是"哪个特征维度配对共享一个旋转角”。选哪种是实现偏好，不是数学性质差异。

8. 一句话总结

DINOv3 的 RoPE = 2D 网格坐标 + 几何级数周期 → $\theta$ 角 → 用 “tile 2 + 切半取反” 技巧把 2D 旋转写成 elementwise。apply_rope 函数外面那层"prefix 拼接"则是为了让 CLS 和 storage token 不参与旋转（它们没有空间位置）。半切分变体与标准 RoPE 在表达能力上等价，差异只是维度配对方式，可在 $W_q, W_k$ 中被吸收。