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

Ganlin Zhang^1,2 · Weirong Chen^1,2 · Daniel Cremers^1,2 · Xi Wang^1,2,3

¹ TU Munich, ² MCML, ³ ETH Zurich

NeurIPS 2026

arXiv | Project Page | Code

TL;DR: BA-T is a repeatable, compact Transformer layer that refines camera pose and local geometry in latent space, inspired by bundle adjustment.

---

Model description

Given an image pair (A, B), BA-T predicts:

• a dense 3D point map (with confidence) for each image, in that image's own camera frame;
• the relative camera pose between the two images, in both directions (each with a confidence).

Both images are encoded by a shared ViT-L encoder (patch size 16, RoPE, 224x224 input). A single shared BA-T layer then refines the image tokens and a relative camera token recursively, K times; K is chosen at test time (default 8, useful range roughly 4 to 24). The refinement is causal, so the prediction after iteration k equals a run with K=k.

The same weights also refine more than two views jointly on a view graph: every view cross-attends to its neighbours and every edge has its own relative camera token. The inference code builds the graph automatically from encoder-feature similarity and recovers global camera poses with confidence-weighted pose graph optimization.

Checkpoints

file	description
bat.pth	Default. The paper model, further fine-tuned with more refinement iterations and on multi-view graphs (2 to 8 views).
bat_paper.pth	The checkpoint used for the results in the paper (two-view).

Both are plain PyTorch state_dicts (about 1.4 GB each) with the same architecture, and work with all scripts in the code repository.

Usage

Install the code following the GitHub README, then download a checkpoint into pretrains/:

git clone https://github.com/zhangganlin/BA-T.git && cd BA-T
# install torch, then: pip install -r requirements.txt
hf download zhangganlin/BA-T bat.pth --local-dir pretrains

Two views:

python demo.py --images path/to/A.jpg path/to/B.jpg --output outputs/demo

from ba_t.data import load_views
from ba_t.inference import load_model, run_pairs

model = load_model('pretrains/bat.pth', device='cuda')
view_a, view_b = load_views(['A.jpg', 'B.jpg'])
out = run_pairs(model, view_a, [view_b], num_refine_iter=8)

T_b_a = out['a'][0]['relative_pose']   # (1,4,4) camera A -> camera B
pts_a = out['a'][0]['pts3d']           # (1,224,224,3) in the camera frame of A
pts_b = out['b'][0]['pts3d']           # (1,224,224,3) in the camera frame of B

Multiple views:

python demo_multiview.py --images path/to/folder --output outputs/demo_mv

from ba_t.inference import run_multiview

out = run_multiview(model, load_views('path/to/folder'))
out['cam_to_world']   # (V,4,4) camera poses, view 0's camera is the world frame
out['pts3d']          # (V,224,224,3) point maps, each in its own camera frame

Conventions

• Input: images are centre-cropped and resized to 224x224; only the central square is seen by the model.
• Cameras follow OpenCV (x right, y down, z forward).
• relative_pose of A is T_B<-A = inv(c2w_B) @ c2w_A, mapping camera-A points into camera B.
• Scale is arbitrary (not metric); within a pair or a multi-view run, points and translations share one scale.
• Point confidence is >= 1 (higher is better); pose confidences are in (0, 1).

Performance

With the optional curope kernel and xformers on an RTX 4070 at K=8: one pair takes about 50 ms and 1.3 GB of GPU memory; 8 views take 0.4 s and 1.6 GB, 64 views 3 s and 3.3 GB.

Acknowledgement

The code builds on ViSTA-SLAM and DUSt3R.

Citation

@inproceedings{zhang2026bat,
  title     = {{BA-T}: An Iterative Transformer for Two-View Bundle Adjustment},
  author    = {Zhang, Ganlin and Chen, Weirong and Cremers, Daniel and Wang, Xi},
  booktitle = {Advances in Neural Information Processing Systems (NeurIPS)},
  year      = {2026}
}