CV君 2025-11-19 15:57 江苏
Transformer太慢？CMU用“信心”给它装上涡轮！
论文标题: Co-Me: Confidence-Guided Token Merging for Visual Geometric Transformers
作者: Yutian Chen, Yuheng Qiu, Ruogu Li, Ali Agha, Shayegan Omidshafiei, Jay Patrikar, Sebastian Scherer
机构: 卡内基梅隆大学 (Carnegie Mellon University), Field AI
项目主页（代码还未放出）: https://github.com/co-me-tokens/CoMe
最近，来自卡内基梅隆大学（CMU）和Field AI的研究团队推出了一项名为Co-Me的技术，全称是置信度引导的Token合并（Confidence-Guided Token Merging）。这个方法可以说是给计算量巨大的视觉几何Transformer模型装上了一个“涡轮增压器”，能在不重新训练、不改变模型主干的情况下，实现惊人的推理加速。
对于熟悉CV领域的朋友来说，Transformer的大名早已如雷贯耳。尤其在三维视觉任务中，像VGGT、MapAnything这类视觉几何Transformer模型，凭借其强大的能力，在三维重建、场景理解上取得了非常好的效果。但“能力越大，计算量越大”，Transformer自注意力机制带来的二次方复杂度，让它们在实时应用和边缘设备部署上显得力不从心。如何为这些强大的模型“瘦身减负”，同时又不牺牲性能，成了一个亟待解决的难题。
研究背景
为了解决Transformer的效率问题，学术界和工业界提出了不少方案。一类是“剪枝”（Token Pruning），比如DynamicViT，它会动态地丢掉一些不那么重要的Token。这种方法在分类等任务上效果不错，但在三维重建这类需要密集预测的任务中，直接丢掉Token会损失空间信息，导致重建细节崩坏。
另一类是“合并”（Token Merging），比如ToMe，它会把相似的Token合并起来，保留了空间覆盖，但效果往往有限。尤其在复杂的几何任务中，简单地基于相似度合并，可能会错误地融合掉一些包含关键几何信息的区域。
而今天的主角Co-Me，则另辟蹊径。它的核心思想非常直观：模型自己最清楚哪些区域的预测是可靠的。既然如此，我们何不利用模型本身的“自信心”来决定计算资源的分配呢？
Co-Me: 基于置信度引导的Token合并
Co-Me巧妙地利用了视觉几何Transformer模型在预测时会同步生成“置信度图”（Confidence Map）这一特性。置信度高的区域，通常是纹理丰富、几何结构稳定的地方；而置信度低的区域，则可能是天空、反光表面或者遮挡区域。Co-Me认为，在这些低置信度区域，我们不需要投入那么多计算，可以把它们的Token合并起来，从而节省计算量。
整个流程如上图所示，主要包含几个关键步骤：
置信度预测器
Co-Me并不会等到整个模型推理完才去拿最终的置信度图，那就太晚了。它设计了一个非常轻量的“置信度预测器”，这个预测器通过“蒸馏”的方式，从Transformer的中间层特征中学习预测最终的置信度。这个过程完全是自监督的，不需要任何额外的标注数据，训练起来又快又方便。论文发现，在VGGT的第15层进行蒸馏，效果最好，既能准确预测置信度，又能给后续层留出足够的加速空间。
Token合并与分割
有了置信度预测，接下来的操作就顺理成章了。
掩码生成 (Mask Generation): 首先，根据预设的合并比例（比如p=0.5），Co-Me会筛选出置信度最低的那些Token，生成一个合并掩码。
合并 (Merge): 对于标记为“合并”的Token组，Co-Me会将它们简单地取平均，融合成一个Token。而高置信度的Token则保持不变。这样一来，Token序列的长度就大大缩短了。
分割 (Split): 在经过Attention或MLP模块的计算后，需要将序列恢复到原始长度，以便下游的预测头使用。Co-Me会把之前合并的那个Token复制多份，填充回原来的位置。
整个过程通过一个高效的自定义CUDA算子实现，最大程度地减少了操作本身带来的额外开销。
注意力偏差校正
这里有一个非常关键的技术细节。当多个Token被合并成一个后，它在Attention计算中的“分量”也应该相应变大，否则会破坏原始Attention权重的分布。
如上图所示，直接合并（左）会扭曲注意力分布。为了解决这个问题，作者们引入了一个巧妙的注意力偏差校正（Attention Bias Correction）。具体来说，如果一个Token是由n个原始Token合并而来的，那么在计算softmax之前，就给它的logit值加上一个偏置项 。这个简单的操作，等效于将其在softmax后的权重放大了n倍，从而完美地补偿了因合并造成的权重损失，保证了模型性能的稳定。
实验证明，这个小小的改动效果显著，能将误差降低多达4倍。
实验结果：快，且稳！
Co-Me的效果可以说非常惊艳。
在VGGT和MapAnything这两个SOTA模型上，Co-Me都取得了显著的加速效果。
对于VGGT，在处理长序列（512帧）时，Co-Me实现了高达 11.3倍 的加速！对于更高的合并率（p=0.9），加速比甚至可以达到惊人的 26.65倍。
对于MapAnything，也实现了最高 7.2倍 的加速。
更重要的是，如此高的加速比，并没有以牺牲大量性能为代价。在单目深度估计、多视图深度估计、位姿估计和点云估计等多个任务上，Co-Me加速后的模型性能与原始模型相比，只有非常轻微的下降，甚至在某些场景下（如ETH3D数据集），由于Co-Me去除了低置信度区域的噪声，性能反而还有所提升！
作者还对比了其他合并策略，如基于相似度合并（Merge by Sim）、随机选择一个（Pick-one）或直接丢弃（Drop-all）。结果显示，Co-Me的置信度引导策略在性能-速度的权衡上全面胜出。
研究团队还在NVIDIA Jetson Thor这样的边缘计算设备上进行了实测。Co-Me加速的MapAnything模型能够以3.5 FPS的帧率运行，实现了1.5倍的端到端加速，将复杂的3D重建任务带入了近实时的领域。
当然，Co-Me也并非完美。在处理一些非常细长的物体时（比如灯柱、雕像的手臂），如果这些区域被判定为低置信度并被合并，可能会导致重建细节的丢失。
总结
总的来说，Co-Me是一项非常实用且巧妙的工作。它抓住了视觉几何Transformer自身就能提供“哪里重要”这一信息的特点，设计了一套即插即用、无需重训的加速方案。通过蒸馏轻量级置信度预测器和注意力偏差校正等一系列创新，实现了性能和效率的极佳平衡。这项技术让原本笨重的视觉几何Transformer变得更加轻快，为它们在自动驾驶、机器人、AR等领域的实时应用扫清了一大障碍。
你觉得这种动态调整计算量的方法，未来还能用在哪些地方？欢迎在评论区留下你的看法！