TikTok RoPE：坐标对旋转与相对位置点积

题意与假设

这道练习题保留了面经中值得训练的核心约束，并将需要澄清的条件单独列出。现场拿到变体时，先和面试官确认这些条件。

• 训练每head正偶数D，相邻实坐标配对，频率base^(−2j/D)，绝对位置显式传入。
• 代码是NumPy数学参考，不是优化GPU kernel；不把长上下文外推效果当必然。

解题路线

1. 第j对坐标用角position·frequency_j旋转，二维矩阵为[[cos,−sin],[sin,cos]]。Q/K同规则处理，V通常不需这份旋转，具体模型另确认。
2. 旋转是正交变换，保每对范数；R(m)ᵀR(n)=R(n−m)，故旋转后Q/K点积依相对position差进入表达，而非只添加一个绝对位置向量。
3. 不同频率让不同坐标对有不同旋转周期；base与D决定频谱。所有位置共同平移时成对Q/K点积不变，是可独立核验的不变量。
4. 配对布局不可混淆：相邻(0,1),(2,3)和half-split(0,D/2)等是不同张量约定，转换需与权重布局一致；不能一侧换布局便认为模型等价。
5. cache存已按对应绝对位置旋转的K，新query要用继续增长的offset；padding/packed sequence的position IDs也要对应各实际序列，不只依batch列号。
6. 旋转允许计算更大位置不等于模型在未训练长度仍可靠，频率缩放/插值改变相位，需要额外训练/评估。数值精度及很大角的计算也需检查。
7. 参考验证复数乘exp(iθ)、范数保持、共同位置平移和错误shape；只看cos/sin输出外观无法发现坐标对错。

手动走一遍

二维x=[1,0]、角π/2输出[0,1]，范数仍1。Q位置m、K位置n同时加100，其旋转点积仍取决n−m，不应改变。

Python 参考实现

下面的代码按上文假设实现，适合本地运行并追加自己的测试用例。

代码区可横向滚动；使用键盘时，先聚焦代码区再按方向键查看。

import numpy as np


def rope(x, positions, base=10000.0):
    """[T,D] real rows, adjacent coordinate pairs, explicit absolute positions."""
    x = np.asarray(x, dtype=np.float64)
    positions = np.asarray(positions, dtype=np.float64)
    if x.ndim != 2 or x.shape[1] == 0 or x.shape[1] % 2:
        raise ValueError("positive even width required")
    if positions.shape != (len(x),) or not np.isfinite(base) or base <= 1:
        raise ValueError("one position per row and finite base > 1 required")
    if not np.isfinite(x).all() or not np.isfinite(positions).all():
        raise ValueError("finite rows and positions required")
    frequencies = base ** (-np.arange(0, x.shape[1], 2) / x.shape[1])
    angles = positions[:, None] * frequencies[None, :]
    c, s = np.cos(angles), np.sin(angles)
    result = np.empty_like(x)
    result[:, 0::2] = x[:, 0::2] * c - x[:, 1::2] * s
    result[:, 1::2] = x[:, 0::2] * s + x[:, 1::2] * c
    return result

复杂度

T×D参考时间/输出O(TD)，预计算cos/sin也O(TD)；模型多层多头与cache布局另计。

易错点与追问

容易失分的边界

• Q/K配对规则不同。
• V必定也旋转。
• 能算更长position就保证长上下文质量。

继续深挖

• 半拆布局如何转换？
• 不同频率缩放如何评估？
• packed sequence位置如何管理？