3D重建不再逐层全注意力:按敏感度分配稀疏与线性,1024视角端到端提速最高5.6倍
核心概要
Zebra3D把预训练3D重建模型的全局注意力按层敏感度换成线性与瓦片稀疏混合,在约94%平均注意力稀疏度下仍保持接近稠密模型的位姿精度,1024视角端到端提速最高5.6倍。
深度剖析
预训练的前馈式3D重建模型可以不再让每一层全局注意力都做完整的softmax两两比较,而是按层分工:一部分层用线性注意力,另一部分层只对选中的空间瓦片做精确softmax。 此前加速这类模型的做法是把同一种近似算子统一套到整条注意力栈上,例如统一稀疏、统一线性或统一键值下采样,并不区分各层对近似的容忍度。 在VGGT、HunyuanWorld-Mirror和Pi3三个骨干上,按敏感度分配在固定0.5L/0.8S预算下取得最高的RealEstate10K AUC@30,分别为0.730、0.850和0.866,而均匀、反向和随机分配明显更低。
转换后的模型可以直接使用,不需要梯度更新;可选的少量微调再把精度拉回来。 训练无关的转换只做一次闭式拟合,而此前的稀疏或线性替换通常需要重新训练或蒸馏才能恢复精度。 训练无关的0.5L/0.6S配置在约79%平均注意力稀疏度下,RealEstate10K AUC@30与稠密VGGT、HYWM、Pi3的差距分别为0.7、0.4和0.6个百分点;经过3000步微调后,Zebra3D-VGGT在约94%稀疏度下达到80.5% AUC@30,与同数据同训练计划的稠密VGGT微调相差1.1个百分点。
长序列上的收益随视角数增加而放大,因为全局注意力在总运行时间中的占比随视角数上升。 在100视角时提速有限,而稠密模型在1000视角时全局注意力已成为主要开销。 在ScanNet-50的1000帧设置下,Zebra3D-KD(0.7L/0.8S)的Chamfer距离为0.460、提速5.60倍,稠密VGGT为0.466、1.00倍;在1024视角下,三个骨干的端到端提速分别达到5.6倍、4.6倍和5.6倍量级。
启示与展望
这套方法面向已经训练好的前馈式3D重建骨干,把它们的全局注意力层替换为线性与瓦片稀疏的混合,转换本身不需要梯度更新,可选的微调需要一次冻结稠密教师的额外前向。评估覆盖7Scenes、NRGBD、DTU上的重建,RealEstate10K与CO3Dv2上的相机位姿,以及ScanNet-50上100到1000帧的长序列,另在Tanks and Temples与ETH3D上检验宽基线场景。作者把输入自适应分配和更广的硬件支持列为后续工作。
作者指出,收益主要体现在长序列,8到16视角时提速有限,因为图像编码器和预测头会逐渐成为端到端加速的上限。层分配按骨干固定,来自单层边际敏感度,没有建模多层同时替换之间的相互作用,也不随输入自适应。分配在测试过的校准数据集上稳定,但对差异较大的架构和领域是否成立仍待确认。优化内核与基准面向AMD MI355X上的Triton,其他加速器上的可移植性与性能需要进一步评估。可选蒸馏在微调时需要冻结的稠密教师,但不增加推理开销。
