面向硬件的校准聚类注意力加速视觉几何Transformer

Visual Geometry Grounded Transformer (VGGT) 标志着 3D 场景重建领域的重大飞跃，因为它是第一个一次性直接联合推断所有关键 3D 属性（相机姿态、深度和密集几何）的模型。然而，这种联合推理机制需要具有极长序列的全局注意力层，这会导致严重的延迟瓶颈。在本文中，我们提出了块式集群注意力（BC注意力）来加速VGGT中的全局注意力层。通过将聚类限制在硬件友好的邻域块内，BC 注意力减少了查询聚类的计算开销以及片内和片外内存之间昂贵的数据移动。这使得 BC 注意力能够扩展到长序列，并在 GPU 上实现实际的延迟改进。此外，我们引入了哈希超平面校准方法和基于阈值的误差补偿方法，以有效减少聚类误差，这是当前聚类注意力机制的瓶颈。总体而言，我们在 GPU 上的实验表明，经过校准的 BC 注意力将全局注意力层加速了 2.10-2.63$\times$，将整个骨干网加速了 1.77-2.35$\times$，对于大场景，损失可以忽略不计 (1%)。在性能损失较小（< 5%）的情况下，校准后的 BC 注意力进一步在全局注意力层上实现了 2.26-2.87$\times$ 延迟改进，在主干网络上实现了 1.90-2.55$\times$ 延迟改进。