BA-T: An Iterative Transformer for Two-View Bundle Adjustment

BA-T matches or exceeds the performance of much larger models while using only 16% of their decoder parameters. It achieves this by replacing heavy, deep attention stacks with a single, lightweight, iterative Transformer layer. This layer acts like a smart correction loop.

In 3D reconstruction, the goal is to estimate camera positions (poses) and 3D structure (geometry) from 2D images. Traditionally, this used Bundle Adjustment (BA). BA is a mathematical optimization process. It iteratively tweaks poses and points to minimize reprojection error (the distance between a projected 3D point and its observed 2D location). Modern feed-forward models use deep neural networks. They attempt to regress these values in a single pass. They use massive "decoder" stacks of cross-view attention (a mechanism where one part of a model queries another to exchange information) to force consistency between images.

The Problem

Current feed-forward models suffer from an architectural tension. To ensure a 3D point looks the same from two angles, they rely on deep attention layers. These layers try to brute-force geometric consistency through sheer parameter count.

Because these models typically perform a single forward pass, they lack a formal mechanism to self-correct. If the initial camera pose is off, the model cannot re-think that decision during inference. As shown in, this leads to misalignment.

Local geometry might look correct in isolation but fails to stitch together globally. These models trade efficiency for depth. They attempt to approximate geometric constraints that classical optimization handles with fewer moving parts.

How It Works

The authors propose BA-T. It treats the iterative nature of classical BA as a design template for a Transformer. Instead of a deep stack, they use a single, repeatable layer. This layer performs updates directly in a latent (hidden) token space.

The mechanism follows three functional stages in each iteration :

1. Latent Residual Computation: The model does not work with raw pixels or explicit XYZ coordinates. It uses a "Camera-Conditioned Geometry Transform" to move geometry tokens from one view's coordinate frame into another. It then performs "Token-level Correspondence Matching." This is an attention-based search to find matching tokens between views .

The discrepancy between these tokens is captured as a "latent residual." This is a learned analogue to the reprojection error. 2. Camera Update: The model uses a one-to-many cross-attention mechanism. The camera token acts as the query. It aggregates latent residuals to predict a delta ($\Delta c$) for the camera pose. This mimics the Schur complement step in classical math. In that step, camera updates are driven by the collective error of all observed points. 3. Local Geometry Update: Once the camera pose is refined, the model updates the geometry tokens. It uses symmetric cross-view attention. This ensures geometry in View A and View B respects the new camera relationship.

Repeating this layer $K$ times shrinks the error. This is visible in the decreasing magnitude of residuals in .

Numbers

The most striking takeaway is the parameter efficiency. BA-T achieves competitive accuracy to models like DUSt3R and VGGT. However, it uses only 38M decoder parameters. This is roughly 16% of the parameter count of some larger counterparts.

On the 7Scenes dataset, BA-T (with $K=4$ iterations) reaches a Pose AUC @5° of 0.683. This outperforms the much larger ViSTA model. Regarding geometry, the authors measure a significant reduction in 3D correspondence error. In a comparison with a baseline using iterative training, BA-T reduced correspondence error by 44%.

The standard stacked decoder approach only saw a 27% reduction.

Efficiency extends to the hardware footprint. In multi-view testing on BundleFusion, the authors report a running time of 50.62 ms. Peak GPU memory usage was only 1.35 GB. This is significantly leaner than MapAnything (3.42 GB) or VGGT (5.11 GB). This makes it a viable candidate for edge devices.

What's Missing

There are gaps a production engineer should note. First, the paper focuses heavily on two-view reconstruction. They demonstrate a multi-view extension in and . However, they admit broader multi-view configurations remain largely unexplored. We do not know how complexity scales in massive, complex loop-closure scenarios.

Second, geometry refinement is described as "indirect." The geometry is updated after the camera. Consequently, geometry metric improvements are less dramatic than pose improvements. The model is excellent at fixing camera position. It is slightly less efficient at fixing object shape.

Finally, the paper lacks a discussion on iterative loop stability. In many recurrent architectures, extreme initial guesses can cause divergence. The authors show convergence in . However, they do not stress-test the limits of this stability.

Should You Prototype This

Yes, if your bottleneck is memory or latency on edge hardware. Swapping a deep decoder for a single 38M-parameter layer is a massive win. The efficiency gains are significant. You gain nearly 5x speed compared to some solver-based baselines.

Code is reportedly available at https://github.com/zhangganlin/BA-T. If you are building real-time SLAM (Simultaneous Localization and Mapping) or 3D scanning for robotics, this architecture is worth a prototype. Test if latent refinement holds up under your specific sensor noise profiles.

Figures from the paper