优质令牌狩猎:视觉几何Transformer令牌选择的搭便车指南 Good Token Hunting: A Hitchhiker's Guide to Token Selection for Visual Geometry Transformers
通过两阶段分层token选择加速视觉几何Transformer,推理时间降低85%以上
前置知识
视觉几何Transformer (Visual Geometry Transformers)
一类用于多视图3D重建的Transformer架构,能够从前向传播中联合预测多个3D属性(如相机参数、点图、深度图等)。代表模型包括VGGT、π3、MapAnything、Depth Anything 3等。这类模型的核心特征是包含全局注意力层,允许跨视图的信息聚合,但这也导致了计算复杂度的二次增长。
本文工作的直接对象,理解其架构和计算瓶颈是阅读论文的基础。论文中的所有实验都是基于VGGT和π3这两个代表性的视觉几何Transformer模型进行的。
全局注意力计算复杂度
在Transformer中,全局注意力的计算复杂度与序列长度的平方成正比。对于视觉几何Transformer,如果有N个输入帧,每帧有L个token,则全局注意力的复杂度为 $O(N^2L^2)$。这意味着当输入图像数量增加时,推理成本会呈二次增长,成为主要的计算瓶颈。
这是本文要解决的核心问题。论文开篇就明确指出这个复杂度问题,所有的方法设计都是为了在保持性能的同时降低这个复杂度。
K-center问题与最远点采样 (Farthest Point Sampling, FPS)
K-center是一个经典的聚类问题,目标是在给定点集中选择K个中心点,使得所有点到最近中心的距离最大值最小化。由于这是NP-hard问题,常用贪婪算法近似求解,最远点采样就是其中一种启发式方法:每次迭代选择距离当前已选集合最远的点。
论文中用于帧间选择的核心算法。通过FPS选择多样化的帧,确保在有限的token预算内最大化场景覆盖,这是实现高效加速的关键技术之一。
注意力模式的熵
注意力分布的熵是衡量注意力集中程度的指标。如果注意力均匀分布在所有token上,熵值接近最大值;如果注意力集中在少数几个token上,熵值较低。归一化熵定义为 $H_{norm} = \frac{\sum_{h,q} H(h,q)}{H \cdot Q \cdot H_{max}}$,其中 $H_{max} = \log(NL)$ 是最大可能熵。
论文中用于分析不同层注意力特性的关键工具。通过计算每层的归一化熵,作者发现早期层的注意力分布比较均匀(熵接近1),而中后期层的注意力更加集中(出现峰值),这直接指导了层自适应的帧内选择策略设计。
绝对轨迹误差 (ATE) 和相对位姿误差 (RPE)
相机位姿估计的标准评估指标。ATE(Absolute Trajectory Error)衡量估计轨迹与真实轨迹之间的整体偏差;RPE(Relative Pose Error)衡量连续帧之间的相对位姿误差,分为旋转误差(RPE-rot)和平移误差(RPE-trans)。这些指标越低表示位姿估计越准确。
论文中所有实验的主要评估指标。通过在不同数据集上的ATE和RPE表现,作者证明了其方法在加速的同时保持了甚至改善了重建质量。
研究动机
视觉几何Transformer已经成为多视图3D重建的强大架构,能够在单次前向传播中联合预测多个3D属性。然而,这些模型的计算成本随着输入序列长度呈二次增长,这是由于模型内部的全局注意力层导致的。当处理的帧数增加时,全局注意力的计算复杂度达到 $O(N^2L^2)$,其中N是输入帧数,L是每帧的token数。这个二次增长的复杂度限制了模型的可扩展性和效率,特别是在处理大规模场景时(如500帧的序列),推理时间会变得不可接受(VGGT在500帧时需要288秒)。
本文的目标是本文的目标是通过限制每个查询在全局注意力层中交互的key/value token数量来加速视觉几何Transformer,同时保持甚至改善重建性能。具体来说,作者希望在有限的token预算下,通过智能的token选择策略,实现显著的加速效果(目标是将500帧场景的推理时间降低85%以上),并且在多个任务(相机位姿估计、3D点云重建、视频深度估计)和数据集上保持与基线模型相当的性能。
与已有工作不同的是,与现有解决方案不同,本文采用了两阶段分层的token选择策略。现有方法如FastVGGT、SparseVGGT直接从所有帧中选择token,或者需要昂贵的模型重新训练(如LiteVGGT需要多GPU多天的训练)。而本文的创新之处在于将token选择分解为两个层次:首先在帧级别进行选择,确定应该保留哪些帧;然后在帧内进行选择,进一步丢弃冗余的token。这种分层设计不仅计算高效,而且不需要重新训练模型,是一种通用且易于应用的加速方案。此外,本文还系统性地探索了不同的token选择策略,发现基于多样性的策略适合帧间选择,而层自适应的策略适合帧内选择,这为视觉几何Transformer的效率改进提供了实践指导。
核心方法
GoToHunt方法采用两阶段分层token选择框架来加速视觉几何Transformer。整体思路是在全局注意力层之前,在K/V空间中进行token选择,确定每个查询token应该与哪些key/value tokens交互。第一阶段是帧间选择,在帧级别确定应该保留哪些帧,采用基于多样性的策略,通过最远点采样选择能够最大化场景覆盖的帧集。第二阶段是帧内选择,在选中的帧内进一步丢弃冗余token,采用层自适应策略,根据不同层的注意力特性应用不同级别的token剪枝。这种分层设计使得方法在有限的token预算下能够有效地保留关键信息,同时显著降低计算成本。
本文的核心创新是将token选择分解为帧间和帧内两个层次,并针对每个层次设计了专门的选择策略。对于帧间选择,作者发现基于多样性的策略(选择能够最大化视图空间覆盖的帧集)优于基于相似性或激活的直观策略。对于帧内选择,作者发现不同层的注意力模式存在显著差异:早期层的注意力分布比较均匀(稀释),而中后期层的注意力更加集中(出现峰值)。基于这一观察,作者提出了层自适应的帧内选择策略,对早期层应用更激进的剪枝(甚至用局部注意力替代全局注意力),对中后期层采用更保守的策略。这种分层自适应的设计使得方法能够在保持性能的同时实现显著的加速效果。
方法步骤详情
GoToHunt方法的具体实现包括以下步骤。第一步,帧间选择:使用place recognition模型提取每张图像的d维特征,然后计算图像之间的余弦距离 $d(i,j) = 1 - \frac{\langle f_i, f_j \rangle}{\|f_i\|_2 \|f_j\|_2}$。在允许每个查询只与K个帧的token交互的预算约束下,使用最远点采样(FPS)贪婪算法选择帧子集 $S^*$,最小化最大距离到最近选中帧的距离:$S^* = \arg\min_{S \subseteq \{1,...,N\}, |S|=K} \max_{i \in \{1,...,N\}} \min_{j \in S} d(i,j)$。第二步,帧内选择:分析全局注意力层的注意力模式,计算每层的归一化熵和Top-1 token权重。根据分析结果设置两个阈值 $l_{local}$ 和 $l_{sample}$:对于 $l < l_{local}$ 的层,用局部注意力替代全局注意力;对于 $l_{local} \leq l < l_{sample}$ 的层,应用帧内下采样(下采样因子为 $\sigma$);对于 $l \geq l_{sample}$ 的层,保持完整的全局注意力。在实验中,默认设置 $l_{local} = 2$,$l_{sample} = 9$,$K = 25$,$\sigma \in \{2, 3\}$。
技术新颖性
本文的技术新颖性体现在多个方面。首先,这是第一个系统性地将视觉几何Transformer的加速问题形式化为token选择问题的工作,提供了一个简单而通用的问题框架。其次,提出的两阶段分层token选择策略是新颖的,将选择过程分解为帧间和帧内两个层次,每个层次都有专门的设计。第三,通过分析全局注意力模式,发现了早期层的注意力稀释现象和中后期层的注意力峰值现象,这为层自适应的帧内选择策略提供了理论基础。第四,通过大量的消融实验,系统性地探索了不同的token选择策略,提供了实践指导。最后,这是一个训练无关的方法,不需要昂贵的模型重新训练,可以作为一个通用的插件应用于不同的视觉几何Transformer模型。
实验结果
实验结果在多个任务和数据集上验证了GoToHunt方法的有效性。在相机位姿估计任务上,在7-Scenes数据集上,GoToHunt ($\sigma=2$) 的ATE为0.0673,优于VGGT基线的0.0698,也优于FastVGGT的0.0727和LiteVGGT的0.0798。在Neural RGB-D数据集上,GoToHunt的ATE为0.0267,显著优于VGGT的0.0374。在TUM-Dynamics数据集上,GoToHunt的性能与基线相当。在3D点云重建任务上,在7-Scenes数据集上,GoToHunt的Accuracy平均值为0.0152,与VGGT的0.0171相当,Completion为0.0188 vs 0.0184,Normal Consistency为0.5567 vs 0.5568。在Neural RGB-D数据集上,GoToHunt的Accuracy为0.0127,优于VGGT的0.0160。在视频深度估计任务上,在Bonn数据集上,GoToHunt ($\sigma=3$) 的AbsRel为0.0288,优于π3基线的0.0333,也优于Speed3R的0.0314。在推理效率方面,在500帧的输入序列上,GoToHunt的推理时间为41.2秒,相比VGGT的288秒降低了85.7%,相比FastVGGT的84.6秒降低了51.3%,虽然略慢于LiteVGGT的36.5秒(但LiteVGGT需要昂贵的模型重新训练)。更重要的是,GoToHunt实现了接近线性的推理时间缩放,随着输入帧数增加,推理时间增长缓慢。消融实验验证了方法设计的合理性:基于多样性的帧间选择策略优于基于时间邻近性、共可见性或注意力的直观策略;层自适应的帧内选择策略优于均匀下采样;增加帧间选择预算K可以从40到100帧的性能逐步改善;方法对层阈值 $l_{local}$ 和 $l_{sample}$ 的选择比较鲁棒。特别是在超参数敏感性分析中,当K从25增加到60时,ATE从0.0677降到0.0673,RPE-rot从0.4495降到0.4203,推理时间从41.2秒增加到63.7秒,显示了性能与效率的良好平衡。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 相机位姿估计 (7-Scenes) | ATE (m) | 0.0673 ($\sigma=2$) | VGGT 0.0698, FastVGGT 0.0727, LiteVGGT 0.0798 | 比VGGT降低3.6%,比FastVGGT降低7.4%,比LiteVGGT降低15.7% |
| 相机位姿估计 (7-Scenes) | RPE-rot (deg) | 0.4471 ($\sigma=2$) | VGGT 0.4953, FastVGGT 0.4254, LiteVGGT 0.6888 | 与VGGT相当,比LiteVGGT降低35.1% |
| 相机位姿估计 (Neural RGB-D) | ATE (m) | 0.0267 ($\sigma=2$) | VGGT 0.0374, FastVGGT 0.0377 | 比VGGT降低28.6%,比FastVGGT降低29.2% |
| 3D点云重建 (7-Scenes) | Accuracy (mean) | 0.0152 ($\sigma=2$) | VGGT 0.0171 | 比VGGT降低11.1% |
| 3D点云重建 (Neural RGB-D) | Accuracy (mean) | 0.0127 ($\sigma=2$) | VGGT 0.0160 | 比VGGT降低20.6% |
| 视频深度估计 (Bonn) | AbsRel | 0.0288 ($\sigma=3$) | π3 0.0333, Speed3R 0.0314 | 比π3降低13.5%,比Speed3R降低8.3% |
| 推理效率 (500帧) | 推理时间 (秒) | 41.2 | VGGT 288.0, FastVGGT 84.6, LiteVGGT 36.5 | 比VGGT降低85.7%,比FastVGGT降低51.3%,略高于LiteVGGT但无需重新训练 |
| 推理效率 (时间缩放) | 缩放特性 | 接近线性缩放 | VGGT二次增长 | 在大规模输入场景下优势明显 |
局限与改进
本文也存在一些局限性。首先,方法引入了几个超参数,包括帧间选择预算K、下采样因子$\sigma$、层阈值$l_{local}$和$l_{sample}$,虽然作者通过消融实验证明了方法对这些超参数的选择比较鲁棒,但在实际应用中仍需要根据具体场景进行调整。其次,帧间选择策略依赖于place recognition模型提取的特征,这增加了一些额外的计算开销,尽管这个开销相对较小。第三,方法假设输入帧具有一定的多样性,对于极端情况(如所有帧都非常相似),基于多样性的帧间选择策略可能效果不佳。第四,虽然方法在多个任务上验证了有效性,但实验主要基于VGGT和π3两个模型,在其他视觉几何Transformer架构上的泛化能力还需要进一步验证。最后,方法主要针对静态场景,对于动态场景的处理能力还有待探索。作者自己也指出,当帧间选择预算K增加到接近总帧数时,性能并不单调提升,甚至会下降,这表明当前视觉几何Transformer的训练可能还没有达到最优,还有改进空间。
独立分析的弱点
从独立分析的角度,本文存在几个可以改进的弱点。首先,帧间选择策略目前基于静态的图像特征,没有考虑查询的具体特性,这意味着所有查询都使用相同的帧集,可能不是最优的。可以考虑设计查询感知的帧选择策略,让不同的查询根据其特性选择不同的帧集。其次,帧内选择策略虽然考虑了层级的差异,但主要是基于统计特性(熵),没有考虑token的具体内容或语义信息。可以探索基于语义或重要性的token选择策略。第三,方法目前是完全训练无关的,虽然这是优点但也限制了性能上限。可以考虑轻量级的自适应训练,在不重新训练整个模型的情况下,让模型学会更好的token选择策略。第四,实验主要关注准确性指标,但对实际应用中的延迟、内存占用等工程指标的评估还不够充分。最后,方法目前主要针对离线处理场景,对于在线或实时应用场景的适应性还需要进一步研究。
未来方向
基于本文的成果,可以延伸出多个未来研究方向。首先,作者提出可以探索基于学习的路由机制在注意力层中的应用,这与本文的帧间选择思想是一致的,但可以通过学习来实现更智能的选择。其次,可以研究将本文的方法应用于其他类型的Transformer模型,不仅仅是视觉几何Transformer,还包括其他领域的多模态或长序列模型。第三,可以探索自适应的token选择策略,根据输入内容动态调整选择策略,而不是使用固定的超参数。第四,可以研究在动态场景中的应用,扩展方法的时间维度,考虑视频序列中的时序一致性。第五,可以探索与模型压缩技术(如量化、剪枝)的结合,实现进一步的加速。第六,作者发现增加帧间选择预算时性能并不单调提升,这暗示当前视觉几何Transformer的训练可能还不够优化,可以探索如何改进训练策略,让模型更好地利用多视图信息。最后,可以开发更高效的实现,特别是针对大规模部署场景的优化。
复现评估
论文提供了项目网站(https://zsh2000.github.io/good-token-hunting.github.io/),作者承诺会开源代码,这有助于复现。实验使用了公开的数据集(7-Scenes、Neural RGB-D、TUM-Dynamics、Bonn)和公开可用的模型(VGGT、π3),实验设置相对透明。作者详细描述了实现细节,包括超参数设置、硬件配置(单张NVIDIA L40S GPU,48GB CUDA内存)、评估指标等。然而,论文没有明确说明是否已经开源代码,具体的实现细节可能需要查看项目网站或联系作者。此外,论文中使用了一些具体的实现技巧(如FPS算法、熵计算),虽然描述了原理,但具体的代码实现可能存在细微差异。总体来说,论文的复现难度中等,有经验的研发人员应该能够基于论文描述复现主要结果,但可能需要一些调试和实验。
论文图表
表7比较了不同方法在不同输入序列长度下的推理时间。输入序列长度从100到500帧不等。VGGT基线的推理时间从13.6秒增长到288.0秒,显示出二次增长特性。GoToHunt的推理时间从7.8秒增长到41.2秒,显示出接近线性的缩放特性。LiteVGGT是最快的(从4.5秒到36.5秒),但需要昂贵的模型重新训练。GoToHunt在保持训练无关的同时,实现了接近LiteVGGT的效率。
这个表格是证明方法效率优势的核心证据,它量化地展示了85%的时间节省和接近线性的缩放特性,这是本文最重要的贡献之一。
表9分析了层阈值$l_{local}$和$l_{sample}$对性能的影响。实验在π3模型上进行,固定K=25。当$l_{local}$从1变化到4,$l_{sample}$从8变化到10时,ATE在0.0567到0.0578之间变化,RPE-rot在0.3406到0.3446之间变化。性能在不同配置下保持相对稳定,说明方法对这些超参数的选择比较鲁棒。
这个消融实验验证了层自适应策略的鲁棒性,说明方法不需要精细的超参数调优就能获得良好的性能,这增强了方法的实用性。