CyclOT: Learning Quadratic Optimal Transport Maps via Synchronized Forward-Backward Interpolants
Authors: Shizhou Xu, Jiachen Liu, Shih-Hsin Wang, Stefan Broecker, Yuhao Huang, Bao Wang, Thomas Strohmer
Organizations: SLAC National Accelerator Laboratory, Stanford University, Menlo Park, CA 94025, USA · Department of Mathematics, University of California, Davis, Davis, CA 95616, USA · Department of Computer Science and Information Engineering, National Taiwan University, Taipei 10617, Taiwan · Department of Mathematics and Scientific Computing and Imaging Institute, University of Utah, Salt Lake City, UT 84112, USA · Department of Computer Science, University of California, Davis, Davis, CA 95616, USA
We study the recovery of forward and reverse quadratic optimal-transport maps from unpaired samples in high dimensions. We introduce a bidirectional neural framework in which the learned maps induce forward and backward displacement interpolants, while the training objective combines bidirectional quadratic action, discriminator-restricted Jensen-Shannon endpoint objectives, and two-sided cycle consistency. The construction requires neither precomputed sample pairings nor an explicit convex-potential parameterization. For absolutely continuous probability measures supported on a compact convex set, and under the stated generator-approximation, discriminator-richness, and minimizer-attainment conditions, we prove a population recovery theorem: for every prescribed accuracy, the sum of the corresponding L2 errors between any global minimizer and the forward and reverse quadratic Brenier maps is below that accuracy, provided the discriminator level is sufficiently large and the annealing action weight becomes sufficiently small. Moreover, the cycle loss is bounded above by λW22(μ0,μ1). Complementary results quantify approximate invertibility and show that exact endpoint Jensen-Shannon divergence and cycle consistency control missing target mass and many-to-one map collapse, respectively. Experiments on Swiss roll, MNIST, CelebA, single-cell perturbation data, and chest X-ray images evaluate endpoint fidelity, transport cost, inverse consistency, and the geometry of the induced interpolations.
Figures & tables
Figure 1: Overview of our proposed target-free OT learning via synchronized forward–backward interpolants. The forward map fθ transports μ0 from left to right, while the backward map gϕ transports μ1 from right to left; Both fθ and gϕ are neural network-approximated maps. The forward and backward induced marginals are matched at the two endpoints, providing marginal supervision for the transport maps.
{x0i}i=1B∼μ0,{x1i}i=1B∼μ1.
Algorithm 1 Forward-Backward Neural Optimal Transport
Figure 2: Qualitative visualization of the Gaussian-to-Swiss-roll experiment. The left panel shows pointwise map-induced interpolation paths from source samples to transported endpoints. The right panel shows forward marginal matching, reverse marginal matching, and target-domain cycle reconstruction.
Method
Global map?
Cost ↓
TDR ↑
TPS ↑
FID ↓
Class-TV ↓
Train Time (min) ↓
Sinkhorn
No
245.201
0.992
0.984
93.773
0.072
–
CyclOT
Yes
212.43
0.957
0.954
92.269
0.039
7
Neural OT
Yes
133.68
0.956
0.952
120.45
0.130
198
ICNN OT
Yes
210.82
0.805
0.749
1565.65
0.162
43
Rectified flow
Yes
246.18
0.565
0.567
1282.97
0.171
29
OT-CFM
Yes
100.87
0.390
0.393
841.36
0.241
32
Table 1: MNIST 0–4→5–9 transport. We report transport cost, target-domain rate (TDR), target probability score (TPS), MNIST-feature FID, and total variation distance between predicted and real target-class distributions over digits 5,…,9 . Sinkhorn hard matching is a finite-sample transductive baseline and does not define a learned global map. Among learned global maps, CyclOT achieves the best MNIST-feature FID and target class-distribution matching, while maintaining high target-domain validity.
Figure 3: MNIST 0–4→5–9 transport.
Figure 4: Training dynamics on the CelebA female-to-male translation task for Neural OT and CyclOT.
Method
Cost ↓
FID ↓
TDR ↑
TPS ↑
∣Enrich\mbox−k100−0.5∣↓
FIE ↓
MIE ↓
ICNN OT
2266.7034
305.2227
46.0%
0.4851±0.25
0.3759
3.2M
35.2M
Neural OT
687.9927
17.8291
79.9%
0.7447±0.27
0.0228
17.6M
1.6M
CyclOT
1069.7815
25.4478
92.0%
0.8602 ± 0.20
0.0357
2.73M
2.73M
Table 2: Quantitative comparison between Neural OT and CyclOT on the CelebA female-to-male translation task.
Input set
TPS
Predicted male
Predicted female
Real source female
0.1278 ± 0.23
479 (9.6%)
4521 (90.4%)
Real target male
0.8893 ± 0.21
4629 (92.6%)
371 (7.4%)
Table 3: FairFace classifier baseline results on real CelebA source-female and target-male images. Each evaluation set contains 5,000 images. The classifier correctly predicts 90.4% of the source-female images as female and 92.6% of the real target-male images as male.
Figure 5: Cycle visualizations on CelebA. The first block shows source female samples x , forward mapped samples fθ(x) , and cycle reconstructions gϕ(fθ(x)) . The second block shows source male samples y , backward mapped samples gϕ(y) , and cycle reconstructions fθ(gϕ(y)) .
Figure 6: Comparison results for female-to-male face translation on CelebA. All three methods are trained for 4 hours under the same experimental setting.
Figure 7: Target Domain Rate (TDR) versus training time for Neural OT and CyclOT.
Table 5: Quantitative comparison between ICCN OT, Neural OT, and CyclOT on the chest x-ray healthy-to-pneumonia transition task.
Input set
TPS
Real source healthy
0.1876 ± 0.1382
Real target pneumonia
0.6254 ± 0.1902
Table 6: Pneumonia classifier baseline results on real source (healthy) and target (pneumonia) images.
Figure 8: Chest x-ray transport results for CyclOT. The first row shows source (healthy) radiographs x , the second row shows the transported samples fθ(x) , and the third row shows the cycle reconstructions gϕ(fθ(x)) . The learned transport changes the target-domain radiographic appearance while preserving the coarse anatomical layout.
Figure 9: Chest x-ray transport results for other tested methods. The first row shows source (healthy) radiographs x and the second row shows the transported samples fθ(x) .
Appendix figures & tables5 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 10: Additional qualitative results on the MNIST Digit Transfer (from digit 0-4 to digit 5-9) translation task.
Figure 11: Qualitative results using the first 50 female images from the CelebA dataset as fixed source samples. Odd-numbered rows show the source female images, and the corresponding even-numbered rows show the mapped male images generated by CyclOT. The full image grid is displayed without post-hoc cropping.
Figure 12: Qualitative results using the first 50 non-smiling images from the CelebA dataset as fixed source samples. Odd-numbered rows show the source non-smiling images, and the corresponding even-numbered rows show the mapped smiling images generated by CyclOT. The full image grid is displayed without post-hoc cropping.
Method
Direction
Cost ↓
∣Enrich\mbox−k100−0.5∣↓
TPS ↑
ICNN OT
Pos → Neg
73142.1836
0.4986
0.1517 ± 0.0260
Neg → Pos
65095.5671
0.4984
0.1513 ± 0.0409
Neural OT
Pos → Neg
1623.5243
0.0507
0.4180 ± 0.1811
Neg → Pos
1984.1612
0.3246
0.4762 ± 0.1860
Ours
Pos → Neg
1472.9777
0.0133
0.4484 ± 0.1856
Neg → Pos
1717.2522
0.3169
0.4098 ± 0.1873
Appendix
Table 7: Quantitative comparison between ICNN OT, Neural OT, and our method on the chest x-ray healthy to pneumonia transition task, broken down by transport direction.
Figure 13: Ablation study on the cycle-consistency loss. Removing the cycle loss leads to mode collapse transport (right) compared to the full model (middle).
Optimal Transport (OT) provides a principled framework for learning transformations between probability distributions from unpaired samples. In many applications, however, a single transformation must map several source distributions to a common target distribution. For example, image restoration might require handling different types of degradation without knowing the degradation of each input at inference time. Simple approaches of pooling the source distributions only encourage alignment with the target at the aggregate level and may leave individual sources misaligned. In our paper, we consider the simultaneous OT problem which formalizes the task of learning a shared transport map that minimizes the average transport cost while aligning each source distribution with a prescribed target. We propose a neural method for solving the simultaneous OT problem by learning a shared transport map that minimizes the average transport cost while aligning each source distribution with a prescribed target. We derive a max-min formulation for learning this map. We illustrate its application to image restoration, where a single model handles multiple degradation types using a common collection of clean target images.
Milena Gazdieva, Kirill Sokolov, Jiawei Chen +2
Applied AI Institute AXXX Moscow, Russia · Lomonosov Moscow State University Moscow, Russia · MIRIAI Moscow, Russia
We propose an implicit neural formulation of optimal transport that eliminates adversarial min--max optimization and multi-network architectures commonly used in existing approaches. Our key idea is to parameterize a single potential in the Kantorovich dual and reformulate the associated c-transform as a proximal fixed-point problem. This yields a stable single-network framework in which dual feasibility is enforced exactly through proximal optimality conditions rather than adversarial training. Despite the inner fixed-point computation, gradients can be computed without differentiating through the fixed-point iterations, enabling efficient training without requiring implicit differentiation. We further establish convergence of stochastic gradient descent. The resulting framework is efficient, scalable, and broadly applicable: it simultaneously recovers forward and backward transport maps and naturally extends to class-conditional settings. Experiments on high-dimensional Gaussian benchmarks, physical datasets, and image translation tasks demonstrate strong transport accuracy together with improved training stability and favorable computational and memory efficiency.
Yesom Park, Eric Gelphman, Stanley Osher +1
Department of Mathematics, University of California, Los Angeles · Department of Applied Mathematics and Statistics, Colorado School of Mines
The optimal transport (OT) map provides a geometric transformation for aligning probability distributions and has become a useful tool in machine learning. However, existing estimators of the OT map still exhibit a gap between sharp statistical guarantees and practical parametric estimation based on stable training objectives. Theoretical estimators achieve minimax optimal convergence rates, but they are typically nonparametric and can incur demanding implementation design or inference costs. Practical estimators are parametric and scalable, but their statistical guarantees remain underexplored, and their min-max, adversarial-like training objectives can be sensitive to optimization algorithms. We propose BROT (Barycentric Regression for OT), a simple two-step method that first computes the unregularized OT plan and then fits a deep neural network (DNN) to the induced barycentric targets by least-squares regression. Under standard regularity conditions, we prove that the DNN estimator of BROT attains the minimax convergence rate, when the ground-truth OT map is Lipschitz. Numerical studies on synthetic datasets and an image dataset show that BROT provides accurate map estimates, strong target distribution matching, and competitive transport costs, compared to existing estimation methods. Experiments on two downstream tasks, single-cell perturbation prediction and unsupervised domain adaptation, further suggest that the accurate estimation of BROT can translate into stronger task performance.