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

最新推荐文章于 2025-12-10 09:11:08 发布

原创最新推荐文章于 2025-12-10 09:11:08 发布 · 922 阅读

8 ·

CC 4.0 BY-SA版权

文章标签：

#学习

论文详解：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合并，然后在做多头注意力的标准流程（各自算个注意力再拼接回去啥的）