BA-T: An Iterative Transformer for Two-View Bundle Adjustment¶
会议: NeurIPS 2026
arXiv: 2606.03287
代码: https://github.com/zhangganlin/BA-T
领域: 几何匹配与定位
关键词: 光束法平差, 迭代Transformer, 双目3D重建, 隐式Token残差, 位姿估计
一句话总结¶
受经典光束法平差(BA)信息传递机制启发,BA-T 提出了一种在隐式 Token 空间中循环复用的紧凑迭代 Transformer 层(仅 38M 解码器参数),通过显式建模隐式几何残差并分步更新相对位姿与局部几何,以极小参数量达到并超越深层堆叠前馈解码器的位姿与重建精度。
研究背景与动机¶
从多视角图像中精确估计相机相对位姿并恢复稠密三维几何结构,是机器人导航、具身智能及 AR/VR 等空间智能任务的核心基础。近年来,以 DUSt3R、VGGT、ViSTA 为代表的前馈稠密重建模型取得了显著突破,它们摆脱了传统两阶段特征匹配与三角化的复杂流程,直接在单次前向推理中回归 3D 点云图与相机位姿。为了强制维持多视角之间的一致性,这类方法普遍依赖极为沉重的跨视角注意力解码器堆叠(解码器参数量常在 113M 至 765M 之间)。然而,深层注意力堆叠仅能隐式拟合多视角几何约束,不仅参数开销与显存占用巨大,而且在推理时缺乏自我纠错机制——一旦单次前向预测存在局部对齐偏差,模型无法像优化算法那样针对预测残差进行反馈调整,在低重叠度或弱纹理场景下极其脆弱。
与之形成鲜明对比的是,经典视觉测量学中的光束法平差(Bundle Adjustment, BA)通过非线性最小二乘优化(如高斯-牛顿法与舒尔补分解),在相机位姿与三维空间点之间交替迭代传播误差信息,具备极强的几何约束刚性与自修正能力。尽管学术界曾尝试将可微 BA 模块嵌入深度学习流程(如 DROID-SLAM、BA-Net),但这类方法通常将网络预测解码为显式几何实体(显式特征轨迹或 3D 点云),再调用二阶显式求解器求逆 Hessian 矩阵。这不仅在推理时带来高昂的计算延迟与显存开销,更割裂了端到端 Transformer 隐式高维表征的学习能力,在缺乏共视约束时极易陷入局部极小值。
上述分歧暴露出现有技术的核心矛盾:前馈模型拥有表达能力极强的隐式高维 Token 表征,却缺乏结构化的迭代几何修正能力;经典几何优化具备严密的迭代纠错机制,却受限于显式几何空间与二阶求解器的计算瓶颈。针对这一瓶颈,本文的核心思考在于:能否跳过显式几何方程与 Hessian 求逆,直接在 Transformer 的高维 Token 隐空间中复刻 BA 经典的交替信息流传播结构?核心 idea:将经典光束法平差中「计算残差 \(\rightarrow\) 舒尔补更新相机位姿 \(\rightarrow\) 条件更新场景几何」的依赖逻辑,提炼为一个可循环复用的紧凑隐式 Transformer 单层,在高维 Token 空间内显式计算隐式残差并交替迭代细化相对位姿与稠密几何。
方法详解¶
整体框架¶
BA-T 整体遵循双视角端到端输入输出模式。输入图像对首先经过冻结/微调的预训练 ViT 编码器提取单视角特征,映射为每张图像 \(N\) 个局部几何 Token \(\bm{g}_a, \bm{g}_b \in \mathbb{R}^{N \times D}\)(维度 \(D=768\));同时初始化一个全局可学习的相对相机位姿 Token \(c_{a \rightarrow b} \in \mathbb{R}^{1 \times D}\)(初始解码状态接近单位变换)。随后,这组 Token 送入 BA-T 核心迭代循环。
在每一次迭代中,BA-T 单层解码器严格遵循经典 BA 的三步信息流:首先通过相机条件变换与软注意力匹配计算跨视角隐式几何残差;接着模拟舒尔补的信息聚合,由单个相机 Token 对全部几何残差 Token 进行跨注意力查询并更新相对位姿;最后利用更新后的位姿对另一视角进行空间对准,通过对称的跨视角注意力完成局部几何 Token 的精细化更新。经过 \(K\) 轮循环复用后,最终的隐式 Token 送入轻量回归头,分别输出各视角的公制度量点云图与相对相机外参。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["输入图像对 (I_a, I_b) 与特征编码<br/>初始几何 Token (g_a, g_b) 与可学习相机 Token c_{a→b}"] --> B["相机引导的几何变换与隐式残差计算<br/>AdaLN 条件变换 + 跨视角软匹配 + MLP 残差"]
B --> C["舒尔补启发的相机位姿更新<br/>相机 Token 对全部隐式残差做 1-to-N 交叉注意力聚合"]
C --> D["对称式局部几何更新<br/>新位姿空间对齐 + 双向跨视角注意力 + 视角内自注意力"]
D -->|未达到迭代轮数 k < K| B
D -->|完成迭代 k = K| E["轻量回归预测头<br/>输出各视角稠密 3D 点云图与相对相机位姿"]
关键设计¶
1. 相机引导的几何变换与隐式残差计算:在隐式表征中模拟重投影误差
在经典 BA 中,几何残差的构建依赖于将三维点经相机位姿变换投影至另一视角像素平面,再与测量特征对比。然而在隐式 Token 空间中,两幅图像的几何 Token 分别定义在各自的局部观察坐标系中,直接对比失去几何意义。为此,BA-T 借鉴扩散模型 DiT 中的自适应归一化机制(AdaLN),引入相机条件变换函数 \(\pi_c^{a \rightarrow b}(g_a^p)\)。该模块利用当前估计的相机 Token \(c_{a \rightarrow b}\) 预测尺度缩放与平移调制向量,对视角 \(a\) 的局部几何 Token \(g_a^p\) 进行特征空间坐标对齐: $\(\pi_c^{a \rightarrow b}(g_a^p) := \text{Condition}(g_a^p, c_{a \rightarrow b})\)$ 变换后的几何 Token 需要与视角 \(b\) 中对应的区域进行匹配。由于图像 Patch 投影后往往落在多个 Token 之间或遭遇遮挡,硬离散匹配不可导且鲁棒性差。BA-T 采用软注意力加权机制:以对准后的 Token \(\pi_c^{a \rightarrow b}(g_a^p)\) 为查询,计算其与视角 \(b\) 中所有 Token \(\bm{g}_b\) 的跨视角注意力权重 \(\bm{\alpha}_p\),得到连续融合的对应 Token \(g_b^q = \sum_{i<N} \alpha_p^i g_b^i\)。随后通过线性投影层 \(f(\cdot)\) 提取属性,并经由双层 MLP 计算两者的特征差值,构造隐式几何残差 Token: $\(r_{a \rightarrow b}^p = \operatorname{MLP}\Big(f(g_b^q) - f\big(\pi_c^{a \rightarrow b}(g_a^p)\big)\Big)\)$ 每个视角 \(a\) 的局部 Token 均生成一个对应的隐式残差 Token,全图共生成 \(N\) 个残差 Token \(\bm{r}_{a \rightarrow b}\),在隐式高维空间完整编码了当前相对位姿与几何估计下的投影不一致性。
2. 舒尔补启发的相机位姿更新:残差到相对位姿的单对多交叉注意力聚合
在经典高斯-牛顿 BA 算法中,法方程通过舒尔补消元法优先求解相机位姿增量 \(\Delta \bm{c} = [\bm{B} - \bm{E}\bm{C}^{-1}\bm{E}^T]^{-1}(\bm{v} - \bm{E}\bm{C}^{-1}\bm{w})\)。其物理实质是将空间各观测点产生的耦合残差信息,经由二阶信息加权聚合到低维的相机参数空间中。BA-T 避免了求解海森矩阵求逆的巨大开销,而是通过单对多(One-to-Many)跨视角注意力机制拟合该信息聚合过程。 具体而言,模型以单个相对相机 Token \(c_{a \rightarrow b}^{(k)}\) 作为唯一的 Query,以全局 \(N\) 个隐式残差 Token \(\bm{r}_{a \rightarrow b}^{(k)}\) 作为 Key 与 Value: $\(\Delta c_{a \rightarrow b}^{(k)} = \operatorname{CrossAttn}\big(c_{a \rightarrow b}^{(k)},\; \bm{r}_{a \rightarrow b}^{(k)}\big)\)$ 通过这种注意力汇聚,全场景各处几何对齐的误差信号根据其置信度与相关性,被自适应加权压缩至相对相机 Token 中,实现位姿的加法更新 \(c_{a \rightarrow b}^{(k+1)} = c_{a \rightarrow b}^{(k)} + \Delta c_{a \rightarrow b}^{(k)}\)。为了同时保证双向几何变换的一致性,网络额外引入一个极轻量的 MLP 学习隐式反变换,由 \(c_{a \rightarrow b}\) 直接推导出反向相机 Token \(c_{b \rightarrow a}\),免去双向维护冗余参数。
3. 对称式局部几何更新:更新位姿引导下的跨视角特征交互与几何细化
在相机位姿获得修正之后,经典 BA 紧接着将位姿增量代回几何方程,求解各空间点的坐标调整 \(\Delta \bm{g} = \bm{C}^{-1}(\bm{w} - \bm{E}^T \Delta \bm{c})\)。BA-T 相应地利用最新估计的相机参数 \(c_{a \rightarrow b}^{(k+1)}\) 与 \(c_{b \rightarrow a}^{(k+1)}\),驱动局部几何 Token 的双向互补更新。 当细化视角 \(a\) 的几何结构时,模型首先将视角 \(b\) 的当前几何 Token 通过最新反向相机位姿变换对齐到视角 \(a\) 的坐标系中,即 \(\bm{\pi}_{c^{(k+1)}}^{b \rightarrow a}(\bm{g}_b^{(k)})\);随后,视角 \(a\) 的局部 Token 通过跨视角注意力查询变换后的对应特征,并结合视角内的 3 层自注意力层维持空间结构平滑性与一致性: $\(\Delta \bm{g}_a^{(k)} = \operatorname{CrossAttn}\Big(\bm{g}_a^{(k)},\; \bm{\pi}_{c^{(k+1)}}^{b \rightarrow a}\big(\bm{g}_b^{(k)}\big)\Big)\)$ $\(\Delta \bm{g}_b^{(k)} = \operatorname{CrossAttn}\Big(\bm{g}_b^{(k)},\; \bm{\pi}_{c^{(k+1)}}^{a \rightarrow b}\big(\bm{g}_a^{(k)}\big)\Big)\)$ 这种位姿先导、双向对称的跨视角几何更新,确保了每次几何修正都在当前最优对准空间下进行,消除了传统单前向注意力网络在坐标系不对齐状态下强行混合特征导致的几何畸变。
损失函数 / 训练策略¶
BA-T 采用多任务联合监督目标。单次迭代 \(k\) 输出的几何与位姿预测同时受三项损失约束: $\(L^{(k)} = \lambda_{\text{pmap}} L_{\text{pmap}} + \lambda_{\text{pose}} L_{\text{pose}} + \lambda_{\text{gc}} L_{\text{gc}}\)$ 其中 \(L_{\text{pmap}}\) 为局部点云图的 \(L_1\) 真实距离回归损失;\(L_{\text{pose}}\) 包含相对旋转矩阵与平移矢量的测地距离/角误差损失,并辅以循环一致性正则项;\(L_{\text{gc}}\) 强制视角 \(a\) 预测点云经相对位姿变换后与视角 \(b\) 真实几何重合的跨视角几何一致性损失。
为了训练单一 Transformer 层具备逐步纠错能力,BA-T 在训练阶段展开 \(K=5\) 次迭代并引入时间递增的深度监督策略: $\(L_{\text{total}} = \sum_{k=1}^{K} \alpha_{\lambda}^{K-k} L^{(k)}, \quad (0 < \alpha_{\lambda} < 1)\)$ 超参数设定为 \(\alpha_\lambda = 0.8\)。随着迭代轮数递增,损失权重成指数增加。这种衰减加权机制迫使网络在早期轮次完成粗对齐,而在后序轮次专注于亚像素级别的精细残差修正,有效防止了循环权重在单一固定点上的退化。
实验关键数据¶
主实验¶
在真实室内场景基准 7Scenes 与 BundleFusion 的两视角重叠图像对上,将 BA-T(测试时 \(K=4\) 迭代)与主流多视角/双视角前馈三维重建模型进行全方位对比。评测指标包括相对位姿曲线下面积(Pose AUC @5°/10°/20°,取旋转与平移角误差的最大值)以及局部三维几何指标(Chamfer 距离、点云精度 Acc、完整度 Comp 与相对深度阈值准确率 \(\delta_{1.25}\))。
| 模型类别 | 方法 | 解码器参数量 | 7Scenes 位姿 AUC@5°/10°/20° ↑ | 7Scenes 几何 Chamf.↓ / \(\delta_{1.25}\)↑ | BundleFusion 位姿 AUC@5°/10°/20° ↑ | BundleFusion 几何 Chamf.↓ / \(\delta_{1.25}\)↑ |
|---|---|---|---|---|---|---|
| 多视角训练 | VGGT [39] | 605M | 0.19 / 0.39 / 0.60 | 0.10 / 0.97 | 0.18 / 0.36 / 0.56 | 0.06 / 0.99 |
| MapAnything [17] | 171M | 0.09 / 0.23 / 0.44 | 0.10 / 0.98 | 0.03 / 0.11 / 0.27 | 0.15 / 0.98 | |
| Depth Anything 3 [19] | 765M | 0.22 / 0.43 / 0.64 | 0.08 / 0.98 | 0.20 / 0.37 / 0.57 | 0.06 / 0.99 | |
| 双视角训练 | DUSt3R [42] | 227M | 0.09 / 0.22 / 0.42 | 0.13 / 0.81 | 0.16 / 0.33 / 0.52 | 0.25 / 0.58 |
| MASt3R [18] | 227M | 0.15 / 0.32 / 0.52 | 0.16 / 0.79 | 0.16 / 0.32 / 0.50 | 0.24 / 0.58 | |
| ViSTA [49] | 113M | 0.13 / 0.33 / 0.56 | 0.10 / 0.98 | 0.15 / 0.31 / 0.51 | 0.06 / 0.99 | |
| ViSTA† (等参基线) [49] | 38M | 0.01 / 0.06 / 0.17 | 0.15 / 0.94 | 0.02 / 0.07 / 0.17 | 0.09 / 0.98 | |
| BA-T (本文, K=4) | 38M | 0.16 / 0.37 / 0.59 | 0.11 / 0.98 | 0.16 / 0.34 / 0.54 | 0.07 / 0.99 |
注:ViSTA† 为使用与 BA-T 完全相同编码器、训练数据与 38M 解码器参数预算的传统堆叠解码器基线。
消融实验¶
在 7Scenes 数据集上固定 38M 解码器参数、相同编码器与训练设置,对 BA-T 的各组件进行系统消融阶梯评估(评估位姿旋转 AUC、平移 AUC 与几何对应点误差 Corr.、相对深度误差 Rel.):
| 阶梯设置 | 配置描述 | 旋转 AUC @5°/10°/20° (deg) ↑ | 平移 AUC @0.05/0.10/0.20 (m) ↑ | 对应误差 Corr.↓ | 相对深度 Rel.↓ | 几何 \(\delta_{1.05}\) / \(\delta_{1.25}\) ↑ |
|---|---|---|---|---|---|---|
| (a) 传统堆叠 | ViSTA† (单次前向, 无迭代) | 0.426 / 0.662 / 0.823 | 0.104 / 0.272 / 0.510 | 0.052 | 0.079 | 0.380 / 0.940 |
| (b) 传统迭代 | ViSTA† w/ 循环迭代训练 | 0.491 / 0.712 / 0.851 | 0.151 / 0.346 / 0.580 | 0.045 | 0.060 | 0.516 / 0.968 |
| (h) 极简 BA-T | + 隐式残差 (无几何变换与软匹配) | 0.637 / 0.810 / 0.904 | 0.296 / 0.534 / 0.735 | 0.039 | 0.051 | 0.565 / 0.971 |
| (g) 完整 BA-T | ++ 几何空间对准 & 软对应匹配 | 0.683 / 0.836 / 0.917 | 0.422 / 0.643 / 0.805 | 0.026 | 0.048 | 0.601 / 0.976 |
| (c) 单项剔除 | BA-T w/o 几何空间变换 \(\pi\) | 0.665 / 0.827 / 0.912 | 0.385 / 0.613 / 0.787 | 0.027 | 0.050 | 0.599 / 0.974 |
| (d) 单项剔除 | BA-T w/o 逐迭代递增深度监督 | 0.667 / 0.827 / 0.912 | 0.404 / 0.625 / 0.793 | 0.036 | 0.049 | 0.620 / 0.974 |
| (e) 单项剔除 | BA-T w/o Token 软对应匹配 | 0.677 / 0.833 / 0.916 | 0.412 / 0.635 / 0.802 | 0.026 | 0.048 | 0.591 / 0.976 |
关键发现¶
- 纯粹循环迭代并非性能核心,BA 结构化信息流才是本质:对比 (a) 到 (b),对传统注意力层直接进行 4 轮迭代训练仅使旋转 AUC@5° 提高 0.065,只解释了总增益的 25%;而从 (b) 引入隐式残差结构跃迁到 (h) 即夺回了 76% 的旋转提升与 54% 的平移提升,证明显式构建残差驱动更新是迭代收敛的关键。
- 几何空间对齐是消除跨视角歧义的基石:在 (c) 中移除相机引导的 AdaLN 几何变换后,位姿与对应点误差显著劣化,证明在未对齐的独立视角坐标系下计算残差会严重误导注意力更新。
- 极小参数实现极致能效比:对比经典可微优化求解器 DROID-SLAM,BA-T 推理耗时仅 30.92 ms(提速约 5.2 倍),显存峰值降低 45%(1.35 GB vs 2.43 GB),且免除了内参标定要求与求解奇异 Hessian 的崩溃风险。
亮点与洞察¶
- 将 BA 信息流提炼为单层隐式算子:巧妙地将二阶优化中的「计算残差 \(\rightarrow\) 舒尔补更新位姿 \(\rightarrow\) 反代更新几何」解构为「AdaLN+MLP残差 \(\rightarrow\) 1-to-N 跨注意力位姿更新 \(\rightarrow\) 对称几何跨视角注意力」,在隐式 Token 空间完整保留了优化物理意义。
- 软注意力突破传统像素级对应局限:传统 BA 依赖严格匹配的特征点,而 BA-T 借助 Transformer 全局软关联能力对局部几何 Patch 进行注意力加权,自然克服了重复纹理、视差断裂以及亚 Patch 级几何偏移带来的优化发散问题。
- 结构先验替代暴力参数堆叠:用仅 38M 参数的单一可复用层实现了媲美甚至超越 227M(DUSt3R)与 605M(VGGT)大型解码器的位姿估计性能,为轻量化移动端空间计算与端侧 SLAM 提供了全新范式。
局限与展望¶
- 训练分布受限于室内静态场景:当前训练数据集(ScanNet、BundleFusion、ARKitScenes 等)均局限于室内刚性静态场景与准针孔相机模型,在室外大规模开阔环境、动态运动物体或强畸变鱼眼镜头下的鲁棒性仍有待验证。
- 位姿与几何更新的收益不对称性:实验中位姿指标随迭代提升幅度远大于几何指标,这是因为几何 Token 本身已继承了预训练 ViT 强大的单目先验,而全局相机 Token 仅从接近单位矩阵的随机态起步,导致几何解码在后序迭代中主要体现为微调对齐而非重构。
- 多视角拓扑扩展仍处初级阶段:虽然论文展示了通过图拓扑拓展至 4 视角的潜力,但目前实现仍以密集全连接视点图为主,面向数百帧长序列的大规模稀疏视角图 SLAM 框架仍待深入探索。
相关工作与启发¶
- vs DUSt3R / MASt3R: DUSt3R 采用 227M 庞大解码器在单次前向中回归全局坐标,缺乏推断期自纠错机制且在小重叠度下易产生全局漂移;BA-T 保持每视角局部几何表示,通过 38M 轻量层循环修正相对位姿,计算量与显存均显著降低。
- vs ViSTA / ViSTA-SLAM: ViSTA 采用非对称双视角交叉注意力解码器(113M),BA-T 在其基础上重构了注意力信息流,通过引入显式隐空间残差概念,以三分之一的参数量大幅刷新了相对位姿 AUC 与 3D 一致性。
- vs DROID-SLAM / BA-Net: DROID-SLAM 依赖显式光流与高斯-牛顿二阶求解器,在两视角与弱共视下求解非凸目标极其困难且极易失败;BA-T 完全运行在高维隐式 Token 空间中,融合数据驱动先验,展现出更强的鲁棒性与 5 倍以上的推理加速。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ (首次在隐式 Token 空间完整映射并复现了经典 BA 的舒尔补信息流结构)
- 实验充分度: ⭐⭐⭐⭐⭐ (对比了多视角/双视角最新 SOTA,并严格设计了等参对比消融阶梯)
- 写作质量: ⭐⭐⭐⭐⭐ (数学概念与 Transformer 注意力机制映射清晰,实验分析深入透彻)
- 价值: ⭐⭐⭐⭐⭐ (为端侧与轻量化三维重建系统摆脱超大 Transformer 解码器堆叠指明了可行路径)