PaperlayerView exact revision on arXiv

Source-linked AI summary

Simultaneous Neural Optimal Transport

Milena Gazdieva, Kirill Sokolov, Jiawei Chen, Evgeny Burnaev, Alexander Korotin

arXiv:2609.37424v1cs.LGcs.AI

TL;DR

Learning one transformation for several source distributions is difficult because matching their pooled mixture need not align each source individually. SimNOT learns a shared neural transport map with source-specific objectives; experiments include synthetic transport and image restoration, with results under degradation shifts.

  • Problem

    A pooled source-to-target map need not align each source individually, while separate maps do not provide one shared transformation.

    Source: §1 · p. 2

  • Method

    SimNOT learns a shared neural transport map with a separate distribution-matching objective for each source, using unpaired samples.

    Source: §1 · p. 2

  • Results

    SimNOT preserves perceptual similarity more consistently under degradation-parameter shifts, although Conditional UOT performs best on training degradations.

    Source: §5.2 · p. 9

  • Takeaways & Limitations

    Experiments show shared transformations can be learned from unpaired samples and handle multiple image degradations without degradation labels at inference.

    Source: §6 · p. 10

    The neural approximation results in the stated subsection assume the target domain is a compact box.

    Source: §4.3 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

The paper addresses how to learn one transport map that aligns multiple source distributions with a common target. It proposes SimNOT, a neural method with a shared map and separate source objectives.

1 INTRODUCTION

  • Pooling sources can match the target in aggregate while leaving individual sources misaligned, whereas separate maps require selecting a source-specific model.The paper motivates a shared map that explicitly accounts for each source.Source: §1 · p. 2
  • Simultaneous OT formalizes learning a common map that aligns several sources with a prescribed target while minimizing average transport cost.The shared map couples source-specific transport tasks, retaining alignment requirements lost when sources are treated only as a mixture.Source: §1 · p. 2
  • SimNOT uses an unbalanced formulation with divergence penalties and learns one neural transport map alongside source-specific potentials.The source identity selects a potential during training, but is not required by the shared map at inference.Source: §1 · p. 2
  • The paper establishes an exact semi-dual representation and neural approximation guarantees, then demonstrates SimNOT on synthetic transport and unpaired image restoration.The restoration evaluation includes unseen degradation strengths.Source: §1 · p. 2

2 RELATED WORK

The related work distinguishes pairwise neural OT and unbalanced OT from simultaneous transport, where one transformation must serve multiple sources.

2 RELATED WORK

  • Continuous neural OT solvers learn transport maps or plans from samples using duality-based objectives, while the cited methods address pairs of distributions.Partial transport can relax distribution matching while controlling input-output similarity.Source: §2 · p. 3
  • Wang & Zhang (2025) formulate simultaneous OT with a common map or kernel; this paper develops a neural method for learning shared transformations from samples.The paper's divergence-based marginal relaxation differs from their inequality-constrained unbalanced formulation.Source: §2 · p. 3
  • Prior continuous unbalanced OT solvers learn maps between two measures; this work combines marginal relaxation with a shared transformation and source-specific potentials.The paper positions UOTM as particularly relevant to its optimization principles.Source: §2 · p. 3

3 BACKGROUND

The background moves from pairwise OT formulations to simultaneous transport with a shared conditional rule, then sets out the sample-based learning problem.

3 BACKGROUND

  • The section reviews OT formulations relevant to the paper and specifies simultaneous OT with a common target.It also introduces the marginal relaxation used in the paper.Source: §3 · p. 3

3.1 OPTIMAL TRANSPORT

  • Monge OT seeks a cost-minimizing map, while the Kantorovich relaxation permits splitting input mass and has an optimal plan under the paper's assumptions.Unbalanced OT instead penalizes deviations of plan marginals from prescribed measures; indicator penalties recover the balanced constraints.Source: §3.1 · pp. 3–4

3.2 SIMULTANEOUS OPTIMAL TRANSPORT

  • Balanced simultaneous OT uses the same conditional output distribution for each source and minimizes average transport cost; for overlapping sources, independent maps may disagree.For pairwise disjoint supports, simultaneous OT reduces to independent pairwise OT when those maps exist.Source: §3.2 · pp. 4–5
  • The simultaneous unbalanced formulation retains a shared conditional distribution while divergence penalties relax marginal agreement; the paper identifies this divergence-based formulation as previously unst studied.The plans may have non-unit mass, with penalties for deviations from each source and the common target.Source: §3.2 · p. 5

3.3 COMPUTATIONAL SETUP

  • The computational setup assumes unknown source and target distributions observed only through finite unpaired samples and seeks a shared map that generalizes to new inputs.This sample-based continuous setting differs from discrete transport between empirical measures.Source: §3.3 · p. 5

4 SIMULTANEOUS NEURAL OT METHOD

The section derives a semi-dual objective for simultaneous unbalanced transport and parameterizes it with a shared neural map and source-specific potentials. It also establishes approximation and consistency results for the neural formulation, with a deterministic-optimum result under specific cost, penalty, and source-density conditions.

4.1 DERIVING THE OPTIMIZATION OBJECTIVE

  • The section derives a simultaneous unbalanced-OT objective and semi-dual whose distribution dependence is through expectations, making it estimable from samples.The semi-dual result is stated under atomlessness and other assumptions; the attainment remark distinguishes optimal-value equality from existence of a primal minimizer.Source: §4 · p. 5, §4.1 · pp. 5–6

4.2 PARAMETRIZATION AND TRAINING

  • A shared stochastic map realizes each conditional transport kernel; neural potentials are source-specific, training alternates potential ascent with map descent, and inference needs neither source index nor potentials.The transport cost is scaled by τ; smaller τ favors balanced simultaneous OT, while larger τ allows greater marginal deviations.Source: §4.2 · p. 6

4.3 THEORETICAL PROPERTIES

  • A shared neural map can approximate any shared stochastic transport mechanism simultaneously across sources with a common accuracy bound.The approximation result concerns a map with a noise input and assumes a compact-box target for the stated approximation setting.Source: §4.3 · pp. 6–7
  • With sufficiently expressive map and potential classes, the restricted population value converges to the exact value J*, without requiring convergence of network parameters.The corollary specifies increasing ReLU-map classes and potential classes dense in C(Y)^K.Source: §4.3 · p. 7
  • For quadratic cost and KL marginal penalties, an attained optimum is deterministic under the stated regularity condition on source densities.The theorem assumes strictly positive C2 source densities on a common open set and yields a Borel map.Source: §4.3 · pp. 7–8

5 EXPERIMENTS

Experiments span a five-source Gaussian-to-Swiss-roll example and unpaired CelebA restoration across five degradation types. The restoration results compare source-specific distribution matching and reconstruction quality, including tests on unseen bilinear factors.

5.1 TOY EXPERIMENT

  • For five Gaussian sources, the shared map recovers the Swiss-roll spiral for all, though outputs vary in spread and some samples fall between its branches.The map receives input coordinates without a source index.Source: §5.1 · p. 8

5.2 UNPAIRED IMAGE RESTORATION WITH MULTIPLE DEGRADATIONS

  • On held-out training degradations, SimNOT outperforms Pooled UOT in mean and maximum source-specific FID and PSNR, while Conditional UOT achieves the best results.Mean and maximum FID aggregate the five source-specific scores; PSNR averages across images and degradation conditions.Source: Table 1 · p. 9, §5.2 · p. 9
  • At unseen bilinear factors ×3 and ×5, SimNOT achieves the lowest FID and LPIPS and the highest PSNR among restoration methods.Models trained at factor ×4 are evaluated without retraining.Source: Table 2 · p. 9, §5.2 · p. 9

6 CONCLUSION

The conclusion presents SimNOT as a continuous, sample-based neural solver for simultaneous OT, with one shared map and a separate distribution-matching objective for each source. Appendix A establishes the semi-dual formulation and related properties under the theorem assumptions.

6 CONCLUSION

  • SimNOT extends neural OT to simultaneous transport, combining a shared map, separate source objectives, an unbalanced formulation, an exact semi-dual representation, and neural approximation guarantees.The conclusion reports experiments on synthetic data and image restoration, with no source labels required at inference.Source: §6 · p. 10

A.1 PROOF OF THE SEMI-DUAL FORMULATION

  • Under the theorem assumptions, the divergences are jointly lower semicontinuous for narrow convergence and satisfy the data-processing inequality.These properties are used in the semi-dual proof.Source: §A.1 · p. 12
  • The proof establishes that allowing the kernel to depend on both x and y does not change the infimum under Theorem 1 assumptions.This no-relaxation-gap result supports the proof’s reduction to the relaxed formulation.Source: §A.1 · pp. 12, 13, 16

A.2 PROOF OF THE APPROXIMATION RESULTS

The approximation proof relates measurable transport maps to finite ReLU networks and establishes convergence for the potential approximation. It also connects map approximation to the resulting output laws under quadratic cost.

A.2 PROOF OF THE APPROXIMATION RESULTS

  • Every measurable map in the stated class can be approximated arbitrarily well in L2 by finite ReLU networks whose outputs lie in Y.The proof uses Lusin’s theorem, universal approximation on a compact set, and coordinate clipping.Source: §A.2 · p. 16
  • Approximating maps yield convergence of output expectations, and sampling the input gives a coupling of the output laws with the same quadratic cost.The passage identifies these steps with equations 13 and 14.Source: §A.2 · pp. 16–17
  • The proof obtains uniform convergence on compact potential sets, then uses their density in C(Y)^K to identify the limit with J* when Y is compact.The compact-set convergence follows from Dini’s theorem.Source: §A.2 · p. 17

A.3 PROOF OF THE DETERMINISTIC STRUCTURE THEOREM

The proof shows that the optimal shared conditional distribution is almost everywhere deterministic, while the related-work and restoration results distinguish SimNOT from barycenter methods and report its performance across degradation settings.

A.3 PROOF OF THE DETERMINISTIC STRUCTURE THEOREM

  • The optimal conditional measure is a Dirac mass for σ-almost every input, yielding a deterministic shared map after modification on a null set.The proof establishes optimality by contradiction: a cheaper coupling would yield an admissible relaxed-problem solution below J∗.Source: §A.3 · pp. 18–19

B EXTENDED DISCUSSION OF RELATED WORK

  • Unlike OT barycenters, SimNOT uses an accessible prescribed target and a shared transformation; unlike BaryIR, it uses unpaired samples and prescribed-target transport.Standard barycenter problems optimize an unknown common distribution and allow separate transport plans.Source: B EXTENDED DISCUSSION OF RELATED WORK · p. 19

C.1 RESULTS FOR INDIVIDUAL DEGRADATIONS

  • SimNOT has lower FID and higher PSNR than Pooled UOT for JPEG, blur, and noise, but Pooled UOT performs slightly better on downsampling degradations.SimNOT has the best FID and PSNR for Gaussian noise and the best PSNR for JPEG; Conditional UOT has the best FID for the remaining four degradations.Source: §C.1 · p. 19, Table 3 · p. 20

C.2 GENERALIZATION TO UNSEEN DEGRADATION STRENGTHS

  • LPIPS is lower for SimNOT in 9/10 conditions versus classifier-routed Conditional UOT and 7/10 versus Pooled UOT; it has the highest mean PSNR among blind methods, while Pooled UOT has slightly lower mean FID.The known-family reference has better mean scores in all three metrics.Source: §C.2 · p. 19, Table 4 · p. 20

C.3 DEGRADATION CLASSIFICATION RESULTS

The classifier is fully accurate on the original degradation settings, but downsampling parameter shifts disrupt its predictions; Table 6 reports accuracy under parameter shifts.

C.3 DEGRADATION CLASSIFICATION RESULTS

  • The classifier achieves 100% accuracy on the test set with the original degradation parameters.The parameter-shift accuracy is measured against the known degradation family.Source: §C.3 · p. 20, Table 6 · p. 21
  • At bilinear factors ×3 and ×5, no test image is classified as bilinear; approximately 80% and 85%, respectively, are classified as JPEG.For bicubic ×3 most inputs are classified as blur, and at ×5 all are classified as blur.Source: §C.3 · p. 21, Table 6 · p. 21

D.1 GAUSSIAN-TO-SWISS-ROLL EXPERIMENT

The experiment transports five Gaussian sources toward a standardized noisy Swiss-roll target using a shared neural map, then visualizes outputs on fixed evaluation batches.

D.1 GAUSSIAN-TO-SWISS-ROLL EXPERIMENT

  • Five 2D Gaussian sources centered at (0, 0), (3, 0), (0, 3), (−3, 0), and (0, −3) are transported toward a standardized noisy Swiss-roll target.The target is generated from a noisy spiral and standardized using fixed statistics from a 20,000-point calibration sample.Source: §D.1 · p. 22
  • The shared map and five source-specific potentials are fully connected networks with four hidden layers of width 256 and ReLU activations.The map is deterministic and has no residual connection.Source: §D.1 · p. 22
  • For training, the specified cost is the scaled quadratic transport cost c(x, y) = τ∥x −y∥2, with no numerical value for τ given in the passage.Source: §D.1 · p. 22
  • Training alternates potential and map updates for 100,000 iterations, averages source contributions equally, and applies an R1 penalty with coefficient γ = 5.Each potential update samples 512 points from each source and 512 common target points.Source: §D.1 · p. 22
  • Figure 3 displays final-map outputs on fixed batches of 2,048 points per source, applied without rejection sampling, alongside the same target sample in every panel.Source: §D.1 · p. 22

D.2 CELEBA RESTORATION EXPERIMENT

The CelebA experiment trains a shared restoration map from unpaired degraded sources and clean targets, then evaluates image quality and distribution matching on held-out images.

D.2 CELEBA RESTORATION EXPERIMENT

  • The split uses 202,592 CelebA images at 64 × 64, with disjoint source, clean-target, and test sets; source-image clean counterparts are not training supervision.The source, target, and test sets contain 91,166, 91,167, and 20,259 images, respectively.Source: §D.2 · p. 22
  • Five source distributions use bicubic or bilinear resizing from 64 × 64 to 16 × 16 and back, JPEG quality 25, Gaussian blur (standard deviation 1.5), or Gaussian noise (standard deviation 0.05).The blur uses a 13 × 13 kernel; noise is added to pixel values in [0, 1] and clipped to [0, 1].Source: §D.2 · pp. 22–23
  • The shared stochastic map uses independent 100-dimensional Gaussian noise, while five independently parameterized potentials score the sources; the map receives no source index.The map uses an NCSN++ generator, and each potential uses a Discriminator_large network.Source: §D.2 · p. 23
  • Training runs for 100,000 iterations, alternating potential and map updates with 16 images per source and a common batch of 16 clean targets.Reported restoration results use the final exponential moving average map checkpoint.Source: §D.2 · p. 23
  • Pooled UOT shares a map and potential across the source mixture, whereas Conditional UOT conditions both on degradation labels supplied by a separately trained five-class classifier for blind inference.The conditional model is also evaluated with the known degradation family.Source: §D.2 · pp. 23–24
  • Evaluation reports source-specific mean and maximum FID, PSNR against clean references, and FID, PSNR, and LPIPS under fixed-model parameter shifts.For bilinear downsampling, factors ×3 and ×5 correspond to intermediate resolutions 21 × 21 and 12 × 12.Source: §D.2 · pp. 23–24

E LIMITATIONS AND FUTURE RESEARCH DIRECTIONS

SimNOT’s per-source potentials increase training memory and computation as the number of sources grows, while training also requires samples grouped by source. The shared map remains unchanged and inference does not require source labels.

Scaling to many sources

  • As source count grows, per-source potentials increase training memory and computational requirements while the transport network remains unchanged.Scaling to many sources could be improved by sharing potential parameters or sampling a subset of sources per training step.Source: E LIMITATIONS AND FUTURE RESEARCH… · p. 24
  • Sharing potential parameters or sampling fewer sources per training step are proposed ways to improve scalability, affecting training only.Inference continues to use a single shared map regardless of source count.Source: E LIMITATIONS AND FUTURE RESEARCH… · p. 24

Source grouping during training

  • Training requires samples grouped by source distribution so each sample can be associated with its corresponding potential; such grouping may be incomplete or unavailable in other applications.The image-restoration setup has this information because the degradation process is known during training, while inference does not require source specification.Source: E LIMITATIONS AND FUTURE RESEARCH… · p. 24
Loading 2609.37424v1…