PaperlayerView exact revision on arXiv

Source-linked AI summary

Neural Optimal Transport

Alexander Korotin, Daniil Selikhanovych, Evgeny Burnaev

arXiv:2201.12220v3cs.LG

TL;DR

The paper addresses computation of optimal transport maps and plans for strong and weak costs, including cases where deterministic maps may not exist. It introduces a stochastic-map neural algorithm, proves neural universal approximation under stated assumptions, and demonstrates image-translation applications while noting task-specific cost design as a boundary.

  • Problem

    Existing scalable methods focus mainly on strong OT and deterministic maps, although optimal plans can be stochastic and deterministic maps may not exist.

  • Method

    The paper represents transport plans with stochastic neural maps, solves a maximin dual problem, and parameterizes the map and potential with neural networks.

  • Results

    The algorithm computes deterministic and stochastic OT plans for strong and weak costs, with applications to unpaired image-to-image translation.

  • Takeaways & Limitations

    The approach provides a generic tool for aligning distributions with interpretable deterministic or stochastic transport maps and controllable output diversity.

  • Takeaways & Limitations

    The learned OT maps may not suit all unpaired tasks, and continuous neural approximations may be irregular and hard to learn in practice.

Abstract

from arXiv · show

We present a novel neural-networks-based algorithm to compute optimal transport maps and plans for strong and weak transport costs. To justify the usage of neural networks, we prove that they are universal approximators of transport plans between probability distributions. We evaluate the performance of our optimal transport algorithm on toy examples and on the unpaired image-to-image translation.

1 INTRODUCTION

The paper targets scalable computation of deterministic and stochastic optimal transport plans with neural networks, covering both strong and weak costs. It contributes a neural algorithm, a universal-approximation result, and applications to large-scale computer vision.

  • Existing work often computes OT costs for generative modeling, while this paper focuses on computing the OT plan itself.
  • The proposed algorithm computes deterministic and stochastic optimal transport plans with deep neural networks for both weak and strong transport costs.It generalizes previously known scalable approaches.
  • The paper proves that neural networks are universal approximators of transport plans between probability distributions.
  • The algorithm is evaluated on large-scale computer vision tasks, including unpaired image-to-image translation.

2 PRELIMINARIES

The preliminaries distinguish deterministic Monge maps from transport plans that permit mass splitting, and introduce weak costs that depend on conditional output distributions. They also state existence and duality conditions for the costs considered.

  • Monge’s formulation minimizes transport cost over measurable maps T that push P to Q, but it is asymmetric and may not permit mass splitting.
  • Kantorovich’s relaxation minimizes over transport plans with marginals P and Q, allowing deterministic or stochastic optimal plans.
  • Weak OT costs take a point x and a conditional distribution of y as input, generalizing strong OT formulations.
  • Under lower-boundedness, convexity, and lower-semicontinuity assumptions, an optimal plan exists and the costs admit a dual formulation.
  • For strong costs, the dual formulation reduces to the familiar Kantorovich duality, although optimal transport plans need not be unique.

3 RELATED WORK

Prior scalable methods mainly compute strong-cost OT and often seek deterministic maps, whereas deterministic maps may not exist and existing approaches have practical or statistical limitations.

  • Existing scalable methods primarily target strong OT and commonly seek deterministic maps, even though such maps may not exist.
  • Primal approaches use generative models and adversarial regularizers, making optimization difficult to configure and sensitive to hyperparameters.
  • ICNN-based methods have theoretical justification but insufficient performance on practical large-scale problems.
  • Entropy-regularized methods recover plans biased from the true plan and make sampling or density computation difficult.
  • The strongest-performing maximin approach recovers deterministic maps fairly well, but its deterministic-map target remains a key limitation.

4 ALGORITHM FOR LEARNING OT PLANS

The algorithm represents transport plans with stochastic neural maps and derives a maximin reformulation of the OT dual problem. Neural networks parameterize the map and potential, enabling optimization for weak and strong costs, while theory establishes approximation under stated assumptions.

  • 4.1 REFORMULATION OF THE DUAL PROBLEM: The dual problem is reformulated as a saddle-point or maximin optimization in which stochastic transport maps solve the inner problem for an optimal potential.
  • 4.1 REFORMULATION OF THE DUAL PROBLEM: Any transport plan can be represented implicitly by a stochastic map T(x,z), with deterministic maps as the special case independent of z.
  • 4.1 REFORMULATION OF THE DUAL PROBLEM: For the γ-weak quadratic cost, a dual maximizer exists, allowing the saddle-point solution to yield an optimal stochastic transport map.
  • 4.2 NEURAL OPTIMAL TRANSPORT: The practical method parameterizes T and the dual potential with neural networks and trains them using stochastic gradient ascent-descent on sampled batches.
  • 4.3 RELATION TO PRIOR WORKS: The method recovers stochastic plans for weak costs and subsumes deterministic-map approaches for strong costs, including the maximin method for quadratic cost.
  • 4.4 UNIVERSAL APPROXIMATION WITH NEURAL NETWORKS: Neural networks can approximate stochastic transport maps in L2 under compactness and finite-moment assumptions, although the resulting continuous networks may be irregular and hard to learn.

5 EVALUATION

The evaluation applies neural OT to deterministic and stochastic unpaired image translation using strong and γ-weak quadratic costs. Deterministic maps preserve input attributes, while stochastic maps preserve attributes while producing multiple outputs.

  • 5 EVALUATION: The algorithm is evaluated on unpaired image-to-image translation using strong and γ-weak quadratic transport costs, with results reported on unseen test images.The datasets range from 50K to 500K images, with 90% used for training and 10% held out for testing.
  • 5.1 PRELIMINARY EVALUATION: For strong costs, the learned stochastic parameterization becomes independent of noise, producing a deterministic map through conditional collapse.The map T(x,z) reduces to T(x) because dependence on z becomes unnecessary.
  • 5.1 PRELIMINARY EVALUATION: For γ-weak costs, the additional objective term stimulates stochastic plans with high conditional variance rather than conditional collapse.This distinction motivates separate deterministic experiments for γ=0 and stochastic experiments for γ>0.
  • 5.2 ONE-TO-ONE TRANSLATION: Deterministic OT maps preserve image content while translating domains, retaining colors, textures, hairstyles, head forms, backgrounds, or facial structure depending on the dataset pair.Examples include shoes↔handbags, CelebA↔anime, CelebA male→female, and anime→shoes translations.
  • 5.3 ONE-TO-MANY TRANSLATION WITH OPTIMAL PLANS: Stochastic OT maps preserve input attributes and generate multiple outputs for each input when trained with the γ-weak quadratic cost.The experiments use γ values including 2/3 and 1, with results shown in Figures 1b and 6.
  • 5.3 ONE-TO-MANY TRANSLATION WITH OPTIMAL PLANS: The method uses only two neural networks and a single γ parameter to control variability, avoiding the more complicated objectives used by several one-to-many translation methods.The comparison is framed against methods such as AugCycleGAN and M-UNIT.

6 DISCUSSION

The method offers interpretable deterministic and stochastic transport maps with controllable diversity, but its learned saddle-point solutions are not always optimal stochastic transport maps. The γ parameter exposes a variance–similarity trade-off in generated samples.

  • 6 DISCUSSION: The method provides interpretable transport maps and controls generated-sample diversity through γ, while remaining potentially applicable beyond unpaired translation.The paper mentions image restoration and domain adaptation as possible additional tasks.
  • 6 DISCUSSION: The method may not suit every unpaired task because appropriate task-specific transport costs remain an open design consideration.The authors identify designing such costs as a promising research direction.
  • 6 DISCUSSION: The saddle-point formulation can yield stochastic functions that are not optimal stochastic transport maps, and this issue produces conditional collapse for strong costs.The paper identifies studying saddle points and arg inf sets as an important challenge for future work.
  • 6 DISCUSSION: The authors release source code for the experiments and model checkpoints to support reproducibility.Implementation details are provided in the official repository README.
  • A VARIANCE-SIMILARITY TRADE-OFF: Increasing γ raises sample variety but can reduce preservation of the input style, making γ a trade-off between variance and similarity.At γ=0 there is no variety; at γ=1 variety is very high but many samples do not preserve input style.

B TOY 2D EXPERIMENTS

The method learns stochastic transport maps for weak quadratic costs and recovers target distributions on toy 2D examples. Its learned average maps closely match discrete OT references and convex-gradient ground truth.

  • B TOY 2D EXPERIMENTS: The experiments use weak quadratic transport because the method becomes deterministic under the strong quadratic cost and prior benchmarks already evaluate that setting.The weak-cost analysis is limited by nonuniqueness of optimal transport plans and the lack of nontrivial ground-truth pairs.
  • B TOY 2D EXPERIMENTS: The learned stochastic maps match the target distributions for Gaussian-to-eight-Gaussians and Gaussian-to-Swiss-Roll transport.The pushforward distribution matches Q in both experiments.
  • B TOY 2D EXPERIMENTS: The learned average maps nearly match convex-gradient OT maps computed by a discrete weak OT solver.The reference maps are obtained from POT using sufficiently sampled 2D distributions.
  • B TOY 2D EXPERIMENTS: L2-UVP errors are below 1% for Swiss Roll and approximately 3% for eight Gaussians.These values quantify similarity between the learned average map and the ground-truth convex gradient.
  • B TOY 2D EXPERIMENTS: The learned transport segments do not intersect for different inputs, consistent with cycle monotonicity of gradients of convex functions.This visual property supports the expected structure of the recovered maps.

C TOY 1D EXPERIMENTS

The 1D experiments evaluate stochastic plans under the 1-weak quadratic cost on Gaussian-to-mixture distributions. Neural OT produces average maps and plans close to discrete weak OT results across restarts.

  • C TOY 1D EXPERIMENTS: In one dimension, transport plans become 2D distributions that can be conveniently visualized.The figures show stochastic plans learned by Neural OT alongside discrete OT results.
  • C TOY 1D EXPERIMENTS: The experiments use 1-weak quadratic cost, for which every optimal plan has a monotone increasing 1-Lipschitz conditional-mean map.The learned stochastic map is checked against this necessary condition.
  • C TOY 1D EXPERIMENTS: The learned average maps nearly match discrete weak OT solutions for Gaussian-to-two-Gaussians and Gaussian-to-three-Gaussians transport.Discrete weak OT is used as the comparison reference, with four restarts and differing seeds.
  • C TOY 1D EXPERIMENTS: The recovered plans are close to discrete OT and have transport costs that nearly equal the optimal cost.The conclusion uses the observed agreement between the learned average map and the reference map.

D COMPARISON WITH PRINCIPAL UNPAIRED TRANSLATION METHODS

Neural OT is compared with principal one-to-one and one-to-many unpaired translation methods on three image-domain pairs. FID scores are comparable on similar domains, while Neural OT outperforms most alternatives on outdoor-to-church translation.

  • D COMPARISON WITH PRINCIPAL UNPAIRED TRANSLATION METHODS: The comparison includes CycleGAN, DiscoGAN, AugCycleGAN, and MUNIT alongside Neural OT.These methods represent principal one-to-one and one-to-many translation approaches.
  • D COMPARISON WITH PRINCIPAL UNPAIRED TRANSLATION METHODS: Existing one-to-one methods preserve translation style comparably to Neural OT, whereas alternative one-to-many methods do not preserve style.The qualitative discussion distinguishes one-to-one and one-to-many baselines.
  • D COMPARISON WITH PRINCIPAL UNPAIRED TRANSLATION METHODS: On outdoor-to-church translation, most competing methods are outperformed by Neural OT in FID.The comparison is reported in Figure 14 and the last row of Table 1.
  • D COMPARISON WITH PRINCIPAL UNPAIRED TRANSLATION METHODS: Neural OT optimizes only two neural networks, unlike the other compared methods.The network-count comparison is summarized in Table 2.

E EXPERIMENTAL DETAILS

The experiments use normalized and resized image data, UNet transport maps with Gaussian noise, and Adam optimization. Training usually converges in about 40K iterations, but weak-cost optimization can be unstable and stochastic training requires multiple GPUs.

  • E EXPERIMENTAL DETAILS: For image translation, inputs are rescaled to [-1, 1] and resized to 64 × 64 or 128 × 128 without augmentation.Anime faces are additionally cropped after rescaling to 512 × 512.
  • E EXPERIMENTAL DETAILS: The stochastic transport map uses a UNet with Gaussian noise as an additional RGBZ input channel.The noise has image-sized dimensions and axis-wise σ = 0.1.
  • E EXPERIMENTAL DETAILS: Training uses Adam with lr = 1·10^-4, batch size 64, 10 inner transport-map iterations, and four noise samples per image for weak costs.Most unpaired translation experiments converge in approximately 40K iterations.
  • E EXPERIMENTAL DETAILS: The dynamic weak-cost schedule increases γ linearly from 0 to the target value during the first 25K potential-update iterations.The target γ values include 2/3 and 1.
  • E EXPERIMENTAL DETAILS: Weak-cost optimization can become unstable when the input distribution is less dispersed than the target or γ is high.Improving stability and convergence is identified as a research direction.
  • E EXPERIMENTAL DETAILS: Stochastic maps require four Tesla V100 GPUs because multiple noise samples are drawn for each input.Deterministic-map training uses comparable generative-model complexity and typically takes one to three days on one V100.

F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS

For strictly convex weak costs, the paper shows that every solution of the maximin objective is a stochastic optimal transport map. The proof establishes both optimal cost attainment and correct generation of the target distribution.

  • F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS: The argument relies on strict convexity to rule out non-transport-map solutions in the relevant arg inf sets.The preceding lemma identifies optimal maps among minimizers, while strict convexity excludes additional solutions.
  • F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS: Strict convexity guarantees that all solutions of the maximin problem are stochastic optimal transport maps.The result applies when an optimal dual potential exists.
  • F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS: The proof shows that a candidate solution attains the optimal transport cost.This follows from the equality in the maximin objective and the optimality of the reference map.
  • F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS: The candidate map generates the target distribution, completing its characterization as a stochastic optimal transport map.Its conditional output distributions are compared with those of a true stochastic optimal map.
  • F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS: In experiments, strong and weak quadratic costs are not strictly convex, so the theorem’s guarantee does not directly cover those settings.The paper reports promising downstream performance nevertheless and identifies strictly convex costs as future work.
  • F OPTIMALITY OF SOLUTIONS FOR STRICTLY CONVEX COSTS: The framework also covers unbalanced optimal transport only by contrast: the paper’s problem enforces π ∈Π(P, Q), unlike the unconstrained unbalanced formulation.This distinguishes the proposed strong and weak OT setting from related stochastic-generator methods.

H ADDITIONAL EXPERIMENTAL RESULTS

Additional experiments examine stochastic translations across several image-domain pairs and test whether synchronized noise induces shared effects across inputs. The authors report no close relation between outputs produced with the same noise vectors for different inputs.

  • H ADDITIONAL EXPERIMENTAL RESULTS: Figures 15 and 16 provide additional OT-map translation examples.The supplied captions identify both figures as examples using W2 OT maps.
  • H ADDITIONAL EXPERIMENTAL RESULTS: Additional examples cover CelebA-female→anime and outdoor→church translations at 128 × 128 resolution.The figures label these examples as stochastic translations.
  • H ADDITIONAL EXPERIMENTAL RESULTS: Further stochastic examples include handbags→shoes, anime→CelebA-female, anime→shoes, and shoes→handbags translations at 64 × 64 or 128 × 128.These figures provide additional examples across several source and target domains.
  • I EXAMPLES WITH THE SYNCHRONIZED NOISE: The synchronized-noise experiment plots an N × K matrix of generated images Tθ(x_n, z_k) across input images and noise samples.It tests whether shared noise produces common generated structures such as shoe or church forms.
  • I EXAMPLES WITH THE SYNCHRONIZED NOISE: No close relation is found between images generated for different inputs with the same noise vectors.This conclusion is based on synchronized-noise examples for handbags→shoes and outdoor→church translations.
  • I EXAMPLES WITH THE SYNCHRONIZED NOISE: Figures 23 and 24 visualize synchronized-noise inputs and translated outputs for the experiment.The figures compare outputs produced from shared noise inputs across different input images.
Loading 2201.12220v3…