Add OT-CFM: mini-batch optimal transport coupling for flow matching

[isNeil]

Summary

Implements mini-batch optimal transport (OT) coupling for the flow matching training path as an opt-in feature (use_ot_coupling=False by default, preserving existing FM behaviour).

OT-CFM

Standard conditional flow matching pairs each clean sample x_0 with an independently drawn noise sample ε. These random pairings produce straight-line paths that cross each other in latent space, forcing the model to learn an averaged, curved vector field that requires more Euler steps to integrate accurately at inference.

OT-CFM (Tong et al., NeurIPS 2023) instead finds the assignment that minimises total squared distance across the mini-batch:

σ* = argmin_σ  Σᵢ ‖x_0ⁱ − ε^σ(i)‖²

Non-crossing paths yield a straighter learned vector field that requires fewer inference steps — which is the primary motivation.

Implementation: _ot_permute() in flow_matching.py solves the linear assignment problem via scipy.optimize.linear_sum_assignment (Hungarian algorithm, O(B³), negligible at B ≤ 64; scipy is already a project dependency). One call is inserted in training_losses() before q_sample. Inference is completely unchanged.

Enabling OT:

python main.py ++experiment.diffusion.use_ot_coupling=true

Non-flow experiments are completely unaffected.

Changes

• src/diffusion/flow_matching.py — _ot_permute() helper, use_ot_coupling flag on FlowMatching.__init__ (default False), called in training_losses()
• src/diffusion/__init__.py — use_ot_coupling kwarg forwarded to FlowMatching (default False)
• src/experiments/train_experiment.py — reads use_ot_coupling from diffusion config (default False)

Results

Training quality — RT-1 Fractal (30k steps)

NanoWM-S/2, 12,515 episodes, RTX 4090, bf16-mixed AMP. * = winner at that step. Bold = that model's best across all checkpoints.

PSNR ↑

Step	FM+OT	FM-noOT
5000	19.52 *	19.03
10000	21.18 *	20.84
15000	20.67	20.70 *
20000	21.51	21.87 *
25000	22.22 *	22.04
30000	21.75 *	21.70

SSIM ↑

Step	FM+OT	FM-noOT
5000	0.7139 *	0.6906
10000	0.7570 *	0.7439
15000	0.7385 *	0.7333
20000	0.7665	0.7670 *
25000	0.7750 *	0.7678
30000	0.7633	0.7702 *

LPIPS ↓

Step	FM+OT	FM-noOT
5000	0.3387 *	0.3510
10000	0.2953 *	0.3106
15000	0.2942 *	0.3054
20000	0.2822	0.2786 *
25000	0.2782 *	0.2855
30000	0.2869	0.2851 *

MSE ↓

Step	FM+OT	FM-noOT
5000	0.0508 *	0.0555
10000	0.0369 *	0.0383
15000	0.0406	0.0399 *
20000	0.0344	0.0333 *
25000	0.0284 *	0.0303
30000	0.0303 *	0.0311

OT-CFM leads clearly at 5k–10k (faster early convergence), noOT catches up and leads around 15k–20k, then results are mixed at 25k–30k. Final quality is essentially tied — all differences are within eval noise at 32 val samples.

---

Inference steps — RT-1 30k (mean ± std, 3 seeds × 32 samples)

Both models were trained with diffusion_steps=1000; inference uses Euler integration with varying num_sampling_steps. The primary claim of OT-CFM is that straighter paths require fewer steps to integrate accurately. Bold * = that model's best across all step counts.

PSNR ↑

Steps	FM+OT	FM-noOT
1	21.06±0.01	20.61±0.12
2	21.72±0.05 *	21.31±0.18
5	21.60±0.11	21.35±0.14
10	21.61±0.10	21.55±0.04
20	21.54±0.10	21.58±0.06 *
50	21.44±0.11	21.56±0.07
100	21.39±0.12	21.54±0.07
250	21.37±0.12	21.52±0.07

SSIM ↑

Steps	FM+OT	FM-noOT
1	0.7102±0.0017	0.7049±0.0017
2	0.7701±0.0003 *	0.7674±0.0028 *
5	0.7598±0.0014	0.7628±0.0024
10	0.7500±0.0012	0.7579±0.0000
20	0.7416±0.0010	0.7510±0.0013
50	0.7340±0.0015	0.7451±0.0014
100	0.7315±0.0017	0.7432±0.0012
250	0.7302±0.0018	0.7419±0.0013

LPIPS ↓

Steps	FM+OT	FM-noOT
1	0.3389±0.0063	0.3658±0.0082
2	0.2217±0.0027	0.2370±0.0035
5	0.1865±0.0014	0.1966±0.0022
10	0.1820±0.0009 *	0.1872±0.0008
20	0.1838±0.0013	0.1859±0.0011 *
50	0.1877±0.0017	0.1875±0.0011
100	0.1893±0.0018	0.1886±0.0009
250	0.1902±0.0019	0.1894±0.0009

Takeaway: FM+OT peaks at 10 steps (LPIPS 0.1820); FM-noOT peaks at 20 steps (LPIPS 0.1859) and never reaches OT's best. On PSNR/SSIM both models plateau by step 2–5. In practice OT-CFM lets you run at 10 steps instead of 20+ — a 2× inference speedup at strictly better perceptual quality, with larger gains at very low step budgets (1–5 steps).

Test plan

• Smoke test: python main.py experiment=dino_wm_pusht_flow ++training.max_steps=100
• use_ot_coupling=false (default) reproduces standard FM — no behaviour change
• use_ot_coupling=true OT-CFM run on RT-1 30k; inference step ablation above

🤖 Generated with Claude Code