每日基础课

一句话定义：RoPE，也叫旋转位置编码，是一种把“位置信息”编码进 Query 和 Key 的方法；它通过让向量的每两个维度按位置旋转一个角度，使注意力天然感知 token 之间的相对距离。口诀：位置不相加，向量转一下；QK 一相乘，距离就显形。生活类比：想象一排人按顺序站队。普通词向量只知道“这个人是谁”，不知道“他站第几个”。RoPE 的做法不是给每个人贴一个位置标签，而是让每个人根据自己站的位置，身体朝向转一个角度。第 1 个人转一点，第 10 个人转多一点。两个人互相看时，他们朝向之间的角度差，就暗含了“他们相隔多远”。必要公式：假设一个向量的一对维度是：x_pair = [x_1, x_2]。在位置 m 上，RoPE 做旋转：RoPE(x, m) = [x_1 cos(mθ) - x_2 sin(mθ), x_1 sin(mθ) + x_2 cos(mθ)]。这就是二维平面里的旋转。符号解释：x 是某个 token 的 Query 或 Key 向量；m 是这个 token 的位置，比如第 0、1、2、3 个；θ 是这一对维度对应的旋转频率；cos 和 sin 控制旋转角度；mθ 表示“位置越靠后，转得越多”。如果向量维度很高，比如 head_dim = 8，就把它拆成 4 对：[0,1]、[2,3]、[4,5]、[6,7]，每一对用不同的 θ 旋转。低维对转得快，高维对转得慢。更完整地说：对第 i 对维度，常用 θ_i = 10000^(-2i/d)，d 是每个 attention head 的维度，i 是第几对维度。它解决什么问题：Transformer 的注意力本身不懂顺序。输入“狗咬人”和“人咬狗”，如果不加位置信息，自注意力只看到一堆 token，不知道谁在前谁在后。对音频、视频也是一样。模型可能知道这一帧有嘴部运动，那一帧有声音，但不知道哪个先发生、哪个后发生、间隔多久。没有 RoPE 会怎样：模型很难稳定区分顺序；长序列上位置关系容易混乱；注意力只能靠内容相似度，缺少“距离感”；在说话人日志中，可能更难判断某段语音前后衔接、说话人切换边界、跨片段依赖。RoPE 的关键好处：它不是简单告诉模型“这是第几个”，而是让 Q 和 K 的点积直接包含相对位置信息。直观理解：位置 m 的 Q 和位置 n 的 K 计算注意力时，结果会自然依赖 m - n，也就是两者相隔多远。所以它特别适合 Transformer 的注意力机制。

练习目标：用 PyTorch 手写一个最小版 RoPE，并把它接到一个单头 self-attention 里。输入：随机 token 表示 x，形状是 [batch, seq_len, dim]。输出：attention 输出 out，形状仍是 [batch, seq_len, dim]；同时打印 RoPE 前后的 Query 片段，确认旋转发生了。验收标准：代码能直接运行；输出形状正确；q 和 q_rope 数值不同；attention 权重每一行加起来约等于 1。思路：第一步，准备输入 x。第二步，用线性层生成 Q、K、V。第三步，为每个位置生成 sin/cos。第四步，把 Q、K 每两个维度一组做旋转。第五步，正常计算注意力。完整可运行代码如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

torch.manual_seed(42)

def build_rope_cache(seq_len, dim, device=None, base=10000):
    """
    生成 RoPE 需要的 cos 和 sin。
    seq_len: 序列长度
    dim: 每个 head 的维度，必须是偶数
    返回:
        cos: [seq_len, dim/2]
        sin: [seq_len, dim/2]
    """
    assert dim % 2 == 0, "RoPE 要求 dim 是偶数，因为要两两成对旋转。"

    # 每一对维度对应一个频率
    # inv_freq shape: [dim/2]
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2, device=device).float() / dim))

    # position shape: [seq_len]
    positions = torch.arange(seq_len, device=device).float()

    # angles shape: [seq_len, dim/2]
    angles = positions[:, None] * inv_freq[None, :]

    cos = torch.cos(angles)
    sin = torch.sin(angles)
    return cos, sin


def apply_rope(x, cos, sin):
    """
    对 x 应用 RoPE。
    x shape: [batch, seq_len, dim]
    cos/sin shape: [seq_len, dim/2]
    返回:
        x_rotated shape: [batch, seq_len, dim]
    """
    batch, seq_len, dim = x.shape
    assert dim % 2 == 0, "最后一维 dim 必须是偶数。"

    # 拆成偶数维和奇数维
    # x_even, x_odd shape: [batch, seq_len, dim/2]
    x_even = x[..., 0::2]
    x_odd = x[..., 1::2]

    # 让 cos/sin 可以广播到 batch 维
    # cos, sin: [1, seq_len, dim/2]
    cos = cos.unsqueeze(0)
    sin = sin.unsqueeze(0)

    # 二维旋转公式：
    # new_even = x_even * cos - x_odd * sin
    # new_odd  = x_even * sin + x_odd * cos
    x_rot_even = x_even * cos - x_odd * sin
    x_rot_odd = x_even * sin + x_odd * cos

    # 交错拼回原来的维度顺序
    x_rotated = torch.empty_like(x)
    x_rotated[..., 0::2] = x_rot_even
    x_rotated[..., 1::2] = x_rot_odd

    return x_rotated


class TinySelfAttentionWithRoPE(nn.Module):
    def __init__(self, dim):
        super().__init__()
        assert dim % 2 == 0, "为了演示 RoPE，dim 需要是偶数。"
        self.dim = dim
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, dim, bias=False)
        self.v_proj = nn.Linear(dim, dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(self, x):
        """
        x shape: [batch, seq_len, dim]
        """
        batch, seq_len, dim = x.shape

        q = self.q_proj(x)
        k = self.k_proj(x)
        v = self.v_proj(x)

        cos, sin = build_rope_cache(seq_len, dim, device=x.device)

        q_rope = apply_rope(q, cos, sin)
        k_rope = apply_rope(k, cos, sin)

        # attention scores shape: [batch, seq_len, seq_len]
        scores = torch.matmul(q_rope, k_rope.transpose(-1, -2)) / math.sqrt(dim)

        # attention weights shape: [batch, seq_len, seq_len]
        attn = F.softmax(scores, dim=-1)

        # out shape: [batch, seq_len, dim]
        out = torch.matmul(attn, v)
        out = self.out_proj(out)

        return out, attn, q, q_rope


def main():
    batch = 2
    seq_len = 6
    dim = 8

    # 假设这是 2 条音频/视频 token 序列，每条 6 个 token，每个 token 8 维
    x = torch.randn(batch, seq_len, dim)

    model = TinySelfAttentionWithRoPE(dim)
    out, attn, q, q_rope = model(x)

    print("input x shape:", x.shape)
    print("output out shape:", out.shape)
    print("attention shape:", attn.shape)

    print("\nRoPE 前 q[0, 1, :4]:")
    print(q[0, 1, :4])

    print("\nRoPE 后 q_rope[0, 1, :4]:")
    print(q_rope[0, 1, :4])

    print("\n检查 attention 每一行是否加起来约等于 1：")
    print(attn[0].sum(dim=-1))

    print("\n第 0 个样本的 attention 矩阵：")
    print(attn[0])


if __name__ == "__main__":
    main()