HTTM: Head-wise Temporal Token Merging for Faster VGGT论文学习

论文详解：HTTM: Head-wise Temporal Token Merging for Faster VGGT

一、研究背景与核心问题

1. 核心模型：VGGT 的优势与瓶颈

VGGT 简介 ：Visual Geometry Grounded Transformer（视觉几何接地 Transformer）是 3D 场景重建领域的突破性模型，能够一次性联合推断所有关键 3D 属性（相机姿态、深度、密集几何），无需复杂的视觉几何后处理，性能超越当时的 SOTA 方法。
核心瓶颈 ：VGGT 的关键设计是 "帧级注意力 + 全局注意力" 交替结构。在全局注意力层中，所有视图的 token 需参与 "全对全" 注意力计算 ------ 即使是小场景，token 序列长度也可能超过 2 万，对于中大型场景的长序列输入（如千帧级），这会导致极高的延迟（latency），成为实际应用的主要障碍。

2. 现有解决方案的不足

论文针对两类主流加速方法的局限性进行了分析：

稀疏注意力方法（如 Sparse Transformer、BigBird）：依赖 LLM/VLM 中 "注意力分数集中于少数 token" 的高稀疏性，但 VGGT 的注意力分布更偏向低分值（图 2 对比 VGGT 与 Llama 3.1），稀疏性极低，这类方法加速效果有限。
Token 合并方法 （如 ToMe、FastVGGT）：
1. 采用 "所有注意力头统一合并" 的策略，导致解合并后 token 特征重复，引发 "特征坍缩"，限制模型表达能力；
2. 未适配 VGGT 的 "时空相似性模式"，合并开销高（全局相似度计算为 O (N²)）；
3. FastVGGT 直接沿用 2D 视觉任务的合并逻辑，未利用 VGGT 的头级别相似性，性能受限。

3. 核心洞察：VGGT 的 Token 相似性模式

论文通过分析发现，VGGT 的 token 相似性源于两大因素，共同形成 "时空耦合" 的独特结构：

RoPE 架构效应 ：VGGT 在每一层都重新应用旋转位置编码（RoPE），而非仅在输入层添加固定编码：
- 全局注意力层：RoPE 增强 token 的空间区分度（ distant locations 相似度低）；
- 帧注意力层：RoPE 对每帧应用相同编码，强化相邻帧相同空间区域的时间一致性。
输入相似度效应 ：
- 空间冗余：高冗余输入（如墙面）的 token 空间相似度显著高于低冗余输入（如杂乱物体）；
- 时间连续性：连续相似帧的 token 在 "帧长度偏移的副对角线" 上呈现高相似度，非连续帧则该结构弱化。

二、核心方法：HTTM（Head-wise Temporal Token Merging）

HTTM 是专为 VGGT 全局注意力层设计的无训练 token 合并方法，通过三大核心创新 + 自适应异常值过滤，在降低延迟的同时保持重建质量，具体设计如下：

1. 头级别 Token 合并（Head-wise Token Merging）

解决问题：现有方法 "所有头统一合并" 导致特征坍缩。
设计逻辑：每个注意力头独立进行 token 合并，基于自身的相似模式，避免跨头共享合并规则。
具体流程 ：
1. 对输入 token 序列\(x \in \mathbb{R}^{N×d}\)进行多头投影，得到\(Q, K, V \in \mathbb{R}^{h×N×d_{head}}\)（h为头数，\(d_{head}=d/h\)）；
2. 每个头i的\(Q^{(i)}, K^{(i)}, V^{(i)}\)独立合并：将 token 分为 "源 token（src）" 和 "目标 token（dst）"，计算 src 与 dst 的余弦相似度，将 top-r 个 src 合并到最相似的 dst（V 的合并跟随 K 的匹配，保证 KV 一致性）；
3. 解合并时，原始 token 的输出为其所属合并 token 的输出拷贝，不同头的合并结果拼接后，特征更具多样性和互补性。

2. 时间重排序与块级别合并（Temporal Reordering and Merging）

解决问题：全局相似度计算开销高（O (N²d_head)），直接空间分块会错过时间相关的高相似 token。
设计逻辑：通过 "时间重排序" 将高相似的时空 token 聚合到固定大小的块中，块内合并降低开销（O (N)）。
具体流程 ：
1. 时间重排序：将\(n_t\)帧中大小为\(n_s\)的空间块堆叠，形成\(n_b = n_s×n_t\)的 "时空块"（块内均为高相似 token）；
2. 块内合并：仅在时空块内进行 token 合并，避免全局匹配；
3. 解合并后重排序：恢复 token 的原始输入顺序，保证模型后续计算的正确性。
优势：既捕捉了空间局部性（块内空间相邻 token），又利用了时间一致性（块内跨帧对应 token），合并质量与开销达到平衡。

3. 自适应异常值过滤（Adaptive Outlier Filtering）

解决问题：固定块大小和合并比例会导致低相似块中的 "异常 token"（原始 token 与合并 token 偏差大）被强制合并，损害重建质量。
设计逻辑：动态过滤偏差过大的异常 token，避免其影响合并效果。
具体流程 ：
1. 初始合并后，计算所有原始 token 到其合并 token 的 L2 偏差（未合并 token 偏差为 0）；
2. 跨所有头标记 top-d% 的大偏差 token 为异常值，生成二进制掩码\(M_o\)；
3. 更新合并 token：减去异常值的贡献，异常值不再参与合并；
4. 高效实现：自定义 CUDA 核，通过块级并行和实时偏差计算降低开销。

4. 第一帧锚定（First-frame Anchoring，误差缓解优化）

解决问题：相机运动较大时，VGGT 的姿态估计误差显著。
设计逻辑：将第一帧所有 token 设为 "参考 token（dst）"，后续帧的高相似 token 合并到这些参考 token，稳定参考帧，减少漂移。
优势：第一帧 token 占比不足 1%，锚定开销可忽略，却能大幅降低相机姿态估计误差。

开始详细研究

论文中很多不懂的概念。开始仔细研究：

稀疏注意力

传统 Transformer 的 "全对全注意力"（比如 VGGT 的全局注意力层）有个致命问题：

假设有 N 个 token（比如文本、图像 patch、3D 点云特征），每个 token 都要和其他所有 N 个 token 计算相似度，计算量是 N×N（O (N²)）。
当 N 很大时（比如长文档、千帧视频、大场景 3D 重建），O (N²) 的计算量会爆炸 ------ 比如 N=1 万时，要计算 1 亿次相似度，GPU 根本扛不住，这就是 "长序列瓶颈"。

稀疏注意力的本质的是：大多数场景下，token 只和少数几个其他 token 有强依赖（比如文本中 "苹果" 只和 "吃""红色" 相关，图像中相邻像素更相关），没必要和所有 token 互动。通过 "筛选关键互动对象"，减少计算量，实现加速。

不管是 Sparse Transformer 还是 BigBird，核心都遵循 3 步：

定义 "稀疏规则"：确定每个 token 只和哪些 "邻居 token" 互动（比如 "只看左右 5 个 token""随机选 10 个远距离 token"）。
只计算稀疏互动：只计算每个 token 与其 "邻居 token" 的注意力分数，不做全对全计算。
保持效果不跌：筛选的 "邻居 token" 要覆盖关键依赖，避免因为少算而导致模型性能下降。

关键优势：计算量从 O (N²) 降到 O (N×K)（K 是每个 token 的邻居数，通常 K 固定为几十 / 几百），当 N 很大时，O (N×K)≈O (N)，几乎线性增长，速度大幅提升。

稀疏注意力的问题

对vggt没啥用。因为vggt里面的token对其他的token都挺关联的，很难说它就和特定的几个token特别关联，其他的就不关联。因此不方便特意筛选几个。

Token 合并

不管是 ToMe 还是 FastVGGT，核心都遵循 "找相似→合并→解合并" 的三步流程，目的是在不丢关键信息的前提下，减少参与注意力计算的 Token 数量：

选 Token：把输入 Token 序列分成两类 ------"目标 Token（dst，保留不合并）" 和 "源 Token（src，待合并）"，通常按固定比例划分（比如 25% dst、75% src）。
算相似：计算所有 src 和 dst 的相似度（常用余弦相似度），找到每个 src 最匹配的 dst（"最佳搭档"）。
合并与解合并：把 src 合并到其最佳匹配的 dst（比如取特征平均值），用少量合并后的 Token 做注意力计算；计算完成后，再把合并后的 Token "解合并" 回原始长度，保证后续网络层的输入维度一致。

关键目标：让合并后的 Token 能替代原始 Token 的核心信息，同时减少 Token 总数（比如从 N 减到 M，M<N），让注意力计算量从 O (N²) 降到 O (M²)，实现加速。

dst和src怎么选的，每个论文可能都不一样：

1. ToMe（2D ViT）

采样方法：随机采样或均匀间隔采样。
比例：默认 dst=25%，src=75%（比如把 ViT 的 Token 序列从 196 个压缩到 49 个，压缩率 75%）。
原因：2D 图像的 Token 是空间像素块，无时间顺序，随机 / 均匀采样足够覆盖所有空间区域，不用复杂规则。

2. FastVGGT（3D VGGT 适配）

采样方法：基础是 "均匀间隔采样"，核心优化是 "第一帧锚定"（强制第一帧 Token 为 dst）。
比例：dst≈34%（最终 Token 序列压缩到原来的 34%），比 ToMe 略高 ------ 因为 3D 重建对 Token 信息更敏感，不敢选太少 dst。
原因：3D 序列有时间连续性，第一帧作为参考帧，其 Token 能稳定相机姿态，避免合并后漂移；均匀间隔采样则覆盖后续帧的空间信息。

Token 合并的问题

1. 所有头统一合并 → 特征坍缩，限制表达能力

具体问题：ToMe、FastVGGT 让所有注意力头用 "同一套相似度规则" 合并 Token------ 比如头 1 关注空间特征、头 2 关注纹理特征，但两者必须按相同的 src-dst 匹配关系合并。
为什么会坍缩：不同注意力头的核心作用是捕捉不同维度的特征（比如有的看局部、有的看全局），统一合并会强制所有头的 Token 特征 "趋同"。解合并后，不同头输出的 Token 特征高度重复，相当于模型丢失了多维度表达能力，这就是 "特征坍缩"。
通俗类比：就像把 5 个不同职能的团队（对应 5 个注意力头），按同一标准裁员合并，最后所有团队都只剩 "全能但平庸" 的成员，无法发挥各自的专业优势。

2. 未适配 VGGT 时空相似性 → 合并开销高（O (N²)）

VGGT 的关键特性：之前讲过，VGGT 的 Token 相似性是 "时空耦合" 的 ------ 同一空间区域的 Token（空间局部性）、相邻帧的对应 Token（时间一致性）相似度最高（源于 RoPE 和输入相似性）。
ToMe/FastVGGT 的问题：它们采用 "全局相似度计算"------ 不管 Token 是空间相邻还是跨帧对应，都要计算所有 src 和 dst 的相似度，复杂度仍是 O (N²)（比如 N=2 万时，要计算 4 亿次相似度）。
核心矛盾：VGGT 的高相似 Token 本来集中在 "时空局部"，但这些方法没利用这个规律，仍做全局计算，导致合并阶段的开销居高不下（比如 FastVGGT 的匹配开销占比达 80% 以上，表 4 显示其匹配 latency 达 2.31s，而 HTTM 仅 0.12s）。

HTTM的做法

HTTM相当于针对vggt特性优化的token合并做法。以前的token合并是把token合并完了再去做多头注意力那一套流程。而HTTM是先分多头注意力，把原本很高维度的token变成多个低维的token（比如说原本4096维度的token被8个头分，变成512维度的token。不过这个地方还没有进行真正的注意力计算，只是先分个片而已）然后在这些低维度的token上各自做token合并，然后在做多头注意力的标准流程（各自算个注意力再拼接回去啥的）