3D RoPE

3D RoPE 把 head_dim 切成三段，分别承载时间、高度、宽度坐标——旋转机制与一维 RoPE 完全同构，代价只是位置 id 从标量变成三元组。

为什么一维编码不够

视频 token 天生带三个坐标 $(t, h, w)$ 。送进 Transformer 前要展平成一维序列，展平那一刻时空邻近关系就被破坏了：同一帧里上下相邻的两个 patch，在序列里可能隔了一整行的距离；而序列里紧挨着的两个 token，可能来自不同帧的无关位置。

旋转位置编码（RoPE）的相对位置性质只对展平后的标量下标成立，模型学到的是”序列距离”而非”时空距离”。要让注意力真正感知时空结构，必须把三个坐标各自独立编码进去——这就是 3D RoPE。

谱系上它处在中间位置：2D RoPE 处理图像 patch 的 $(h, w)$ ，3D RoPE 扩展到视频的 $(t, h, w)$ ，再往上是 MRoPE 详细介绍所代表的通用方案，用同一套三维 id 空间统一编码纯文本、图像和视频。

维度分配

核心设计决策：把 head_dim 拆成三段，每段只旋转一个坐标轴。常见分配方式为

$d i m_{t} = d - 4 ⌊ d /6 ⌋, d i m_{h} = d i m_{w} = 2 ⌊ d /6 ⌋$

三段之和恰为 $d$ ，且每段都是偶数（RoPE 按二维子空间成对旋转的硬性要求）。 $d = 192$ 时得到 64/64/64 的均分； $d = 128$ 时为 44/42/42，时间轴略多拿两对频率。

时间轴通常分到不少于空间轴的维度：一段视频的帧数跨度往往小于空间 patch 数，但时序一致性对生成质量的影响更敏感，多给频率意味着更细的时间分辨能力。

实现

每段独立生成频率、拼接后走标准 rotate_half 路径：

import torch
 
def split_dims(head_dim: int) -> tuple[int, int, int]:
    """把 head_dim 拆成 (dim_t, dim_h, dim_w)，各段均为偶数且总和为 head_dim。"""
    unit = head_dim // 6
    return head_dim - 4 * unit, 2 * unit, 2 * unit
 
 
def build_3d_freqs(pos_ids: torch.Tensor, head_dim: int, theta: float = 10000.0):
    """pos_ids: (N, 3) 的 (t, h, w) 整数坐标，返回 (N, head_dim) 的 cos/sin。"""
    halves = []
    for axis, dim_ax in enumerate(split_dims(head_dim)):
        # 每个轴有自己独立的频率阶梯，互不干扰
        inv_freq = 1.0 / (theta ** (torch.arange(0, dim_ax, 2).float() / dim_ax))
        halves.append(pos_ids[:, axis, None].float() * inv_freq)  # (N, dim_ax//2)
 
    freqs = torch.cat(halves, dim=-1)              # (N, head_dim//2)
    emb = torch.cat([freqs, freqs], dim=-1)        # (N, head_dim)
    return emb.cos(), emb.sin()
 
 
def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)
 
 
def apply_3d_rope(q: torch.Tensor, k: torch.Tensor, cos, sin):
    """q/k: (..., N, head_dim)"""
    q_out = q * cos + rotate_half(q) * sin
    k_out = k * cos + rotate_half(k) * sin
    return q_out, k_out
 
 
# 2 帧 x 3 高 x 3 宽的视频块
T, H, W, head_dim = 2, 3, 3, 192
grid = torch.stack(torch.meshgrid(
    torch.arange(T), torch.arange(H), torch.arange(W), indexing="ij"
), dim=-1).reshape(-1, 3)                          # (18, 3)
 
cos, sin = build_3d_freqs(grid, head_dim)
q = torch.randn(1, 4, grid.shape[0], head_dim)     # (B, heads, N, D)
k = torch.randn_like(q)
q, k = apply_3d_rope(q, k, cos, sin)
print(q.shape)                                     # torch.Size([1, 4, 18, 192])

cat([freqs, freqs], -1) 与 rotate_half 的配合值得留意：索引 $i$ 与 $i + d /2$ 构成旋转对，两者取自同一轴的同一个频率，所以三段拼接不会让不同轴的分量串到同一个二维子空间里去。

轴解耦

3D RoPE 最重要的性质是三个轴完全解耦。注意力分数可以分解为三段内积之和，第 $α$ 段只依赖该轴的坐标差 $Δ α$ ：

$⟨ q_{m}, k_{n} ⟩ = \sum_{α \in {t, h, w}} f_{α} (Δ α)$

时间段只看 $Δ t$ ，空间段只看 $(Δ h, Δ w)$ 。这带来几个直接好处：

性质	来源
分辨率泛化	空间旋转绑定 patch 坐标而非序列下标，改 $H / W$ 不改变同一空间偏移的编码
时长泛化	时间段独立，加帧只是把 $t$ 的取值范围拉长
位置插值	三个轴可各自缩放 id，RoPE 长度外推与插值的 PI / NTK-aware / YaRN 都能逐轴套用
零额外参数	纯函数式变换，不引入可训练权重

逐轴插值是实践中的关键能力：训练时用 256×256 短片，推理时想出 512×512 长片，可以只对空间两轴做 NTK-aware 缩放而不动时间轴，避免一刀切缩放损伤时序建模。这是一维位置编码方案做不到的粒度。

应用

3D RoPE 已被视频生成 DiT 架构广泛采用，Sora、VideoPoet、CogVideoX 等模型都在时空注意力层使用这类三轴分解的旋转编码。相比早期视频模型把时间位置做成可学习 embedding 的做法，它在不增加参数的前提下换来了任意分辨率与任意时长的推理灵活性——这对需要同时支持多种输出规格的生成模型是刚性需求。

知识花园

探索

3D RoPE

为什么一维编码不够

维度分配

实现

轴解耦

应用

关系图谱

目录

反向链接