Source-linked AI summary
Wasserstein-2 Generative Networks
Alexander Korotin, Vage Egiazarian, Arip Asadulaev, Alexander Safin, Evgeny Burnaev
TL;DR
The paper addresses how to fit well-structured generative mappings when multiple mappings can transport the same distributions. It trains cyclically monotone Wasserstein-2 maps with an end-to-end non-minimax algorithm, reporting faster convergence than prior methods and potential across practical tasks.
Problem
Multiple mappings can transport the same distributions, while prior GAN and Wasserstein-2 approaches use minimax optimization.
Source: §1 · pp. 1–2Method
The paper trains cyclically monotone quadratic-cost transport maps with an end-to-end non-minimax algorithm using input convex neural networks and cycle-consistency regularization.
Source: §1 · pp. 2–3Results
Takeaways & Limitations
The algorithm's demonstrated scope includes latent-space mass transport, image-to-image color and style transfer, and domain adaptation.
Source: §6 · p. 9A cycle-monotone model may transfer only some desired style attributes when the target transfer is not cycle monotone, such as replacing snow with grass.
Source: §5.3 · p. 9
Abstract
from arXiv · showhide
We propose a novel end-to-end non-minimax algorithm for training optimal transport mappings for the quadratic cost (Wasserstein-2 distance). The algorithm uses input convex neural networks and a cycle-consistency regularization to approximate Wasserstein-2 distance. In contrast to popular entropic and quadratic regularizers, cycle-consistency does not introduce bias and scales well to high dimensions. From the theoretical side, we estimate the properties of the generative mapping fitted by our algorithm. From the practical side, we evaluate our algorithm on a wide range of tasks: image-to-image color transfer, latent space optimal transport, image-to-image style transfer, and domain adaptation.
1 INTRODUCTION
The paper seeks well-structured generative mappings and restricts them to cyclically monotone maps, which are unique for almost every pair of continuous distributions. It develops a non-minimax Wasserstein-2 training method and evaluates it on several practical tasks.
1 INTRODUCTION
- For almost every pair of continuous distributions, a unique cycle-monotone mapping transports one distribution to the other.This restriction narrows the class of candidate generative mappings.Source: §1 · p. 2
- Gradients of convex functions yield cycle-monotone mappings, and input convex neural networks provide a practical class of such functions.The paper uses ICNN gradients to represent generative maps.Source: §1 · p. 2
- GAN training and prior Wasserstein-2 approaches use minimax objectives, motivating a training method that avoids minimax optimization.The paper describes GANs as typically requiring minimax optimization and identifies the same drawback in prior W2 methods.Source: §1 · p. 2
- The contributions include an end-to-end non-minimax algorithm for quadratic-cost optimal transport and theoretical bounds on its fitted mapping's approximation properties.The paper also develops an ICNN class and evaluates on color transfer, latent-space transport, style translation, and domain adaptation.Source: §1 · pp. 2–3
- The proposed method scales to high dimensions and converges up to 10x faster than its predecessors.It extends Makkuva et al. (2019) and replaces minimax optimization with cyclic regularization.Source: §1 · p. 3
2 RELATED WORK
The section reviews GAN and optimal-transport approaches to generative mapping, noting that several W2 methods use minimax optimization or require sequential stages.
2 RELATED WORK
- GANs train a generator against a discriminator, and their minimax objective makes training hard.The discriminator estimates a dissimilarity between generated and real distributions.Source: §2 · p. 3
- Optimal-transport methods seek generative mappings that minimize a specified transport cost; for quadratic cost, the value is the squared Wasserstein-2 distance.The optimal mapping for quadratic cost is cyclically monotone.Source: §2 · p. 3
- Makkuva et al. (2019) removes Taghvaei and Jalali’s costly inner optimization sub-problem but remains minimax; Seguy et al. (2017) uses a non-minimax objective in two sequential steps.Source: §2 · p. 3
3 PRELIMINARIES
The preliminaries connect quadratic-cost optimal transport to convex potentials and their gradients, then describe correlation optimization and extraction of approximate transport maps.
3 PRELIMINARIES
- Under the paper's assumptions of continuous distributions with finite second moments on equal-dimensional spaces, the optimal map exists and is unique on the support of P up to small sets.The map is also invertible.Source: §3 · p. 4
- The W2 dual formulation minimizes over convex potentials, and the value of that minimum is called cyclically monotone correlation.For quadratic cost, finding the optimal map is equivalent to maximizing correlations.Source: §3 · p. 4
- The gradient of the optimal convex potential gives the W2 optimal map, while the conjugate potential's gradient gives its inverse.The paper approximates the primal potential with input convex functions and optimizes correlations.Source: §3 · p. 4
- When the approximate correlation is close to the optimal correlation, the extracted approximate generator pushes P close to Q.Computing the earlier objective requires recovering inverse values for samples from Q, an expensive convex optimization subproblem.Source: §3 · pp. 4–5
4 AN END-TO-END NON-MINIMAX ALGORITHM
The paper proposes a cycle-consistency-regularized, non-minimax objective for fitting cyclically monotone maps and provides approximation guarantees for the resulting generators.
4.1 ALGORITHM
- Cycle-consistency regularization removes the minimax objective by encouraging the two learned gradient maps to be inverses.When the gradients are exact inverses and the forward gradient is injective, the corresponding potentials are convex conjugates up to a constant.Source: §4.1 · p. 5
- The algorithm estimates correlations and cycle-consistency penalties from mini-batches, combines them in a weighted total loss, and takes stochastic gradient steps over both potentials.The regularization coefficient is λ > 0.Source: §4.1 · pp. 5–6
- The approach converges up to 10x faster than minimax alternatives by Makkuva et al. (2019) and Taghvaei and Jalali (2019).The comparison is reported in the experiments section.Source: §4.1 · p. 6
4.2 APPROXIMATION PROPERTIES
- Under the stated convexity, smoothness, and bijective-gradient assumptions, the fitted forward and inverse generators have bounds on how closely they push the distributions to their targets.The guarantees concern approximators obtained by optimizing regularized correlations.Source: §4.2 · p. 6
1. Correlation Upper Bound (regularized correlations dominate over the true ones)
The section bounds regularized correlations by true correlations plus an approximation-order term, then links approximation quality to forward and inverse W2 generative accuracy.
1. Correlation Upper Bound (regularized correlations dominate over the true ones)
- Regularized correlations do not exceed true correlations plus an O(ϵX + ϵY) term.Source: §1 · p. 7
- The fitted potentials push P and Q within O(ϵX + ϵY) of Q and P, respectively.Fully connected ICNNs satisfy the universal approximation property Chen et al. (2018).Source: §1 · p. 7
- The generative guarantees also extend to latent-space mappings and non-continuous distributions.For latent-space mappings, the target-to-generated distribution bound depends on latent fit quality and auto-encoder reconstruction loss.Source: §1 · p. 7
5 EXPERIMENTS
Experiments test the approach in Gaussian transport, latent-space transport, and image-to-image style translation, with results presented as evidence of feasibility rather than state-of-the-art performance.
5 EXPERIMENTS
- The experiments cover Gaussian optimal transport, latent-space mass transport, and style translation; the authors frame them as feasibility tests, not state-of-the-art claims.Additional experiments cover color transfer and domain adaptation.Source: §5 · p. 7
5.1 OPTIMAL TRANSPORT BETWEEN GAUSSIANS
- W2GN, MM-1, and MM-2 achieve L2-UVP below 3% even in R4096, while W2GN converges several times faster and LSOT error grows with dimension.The table compares L2-UVP across dimensions; the convergence comparison is attributed to the minimax methods' inner optimization cycle.Source: Table 1 · p. 8, Figure 2 · p. 8, §5.1 · p. 8
5.2 LATENT SPACE OPTIMAL TRANSPORT
- For CelebA latent-space transport, applying the map to standard-normal noise improves FID, although the generated images do not change significantly.Images are 64 × 64; Table 2 reports FID scores, and Figure 3 compares decoded images before and after transport.Source: Table 2 · p. 8, Figure 3 · p. 8, §5.2 · p. 8
5.3 IMAGE-TO-IMAGE STYLE TRANSLATION
- In winter-to-summer style transfer, the model colors trees green but has difficulty replacing snow with grass.The experiments use randomly cropped 128 × 128 RGB images from Winter2Summer and Photo2Cezanne.Source: §5.3 · p. 9, Figure 4 · p. 9
- Because quadratic-cost OT is permutation invariant and ignores relations among pixels or channels, it may struggle with convolutional architectures designed to preserve local image structure.The authors suggest exploring costs defined on image Gaussian pyramids or perceptual quadratic costs.Source: §5.3 · p. 9
6 CONCLUSION
The conclusion presents a non-minimax method for learning quadratic-cost optimal transport maps, supported by approximation theory and experiments. The appendices prove the main results and establish mathematical properties used in those proofs.
6 CONCLUSION
- The paper develops an end-to-end non-minimax algorithm for quadratic-cost optimal transport, establishes approximation-theoretic justification, and reports experiments across several practical tasks.The conclusion names latent-space mass transport, image-to-image color and style transfer, and domain adaptation.Source: §6 · p. 9
A PROOFS
- The proofs establish the main theorems, justify estimating the quality of a latent-space map combined with its decoder, and extend the method to distributions without a density.The appendix also translates the proofs to a single-potential correlation-optimization case.Source: A PROOFS · p. 12
A.1 PROPERTIES OF WASSERSTEIN-2 METRIC AND RELATION TO L2 SPACES
- The appendix relates Kantorovich transport plans to Monge maps and identifies the square root of the quadratic transport cost as the Wasserstein-2 metric.When an optimal map exists, its induced plan is given by the identity and the map; otherwise, the plan formulation relaxes the map formulation.Source: §A.1 · p. 12
- Appendix A.1 supplies Lipschitz and L2 inequalities for Wasserstein-2 distance that are used in the proofs.The Lipschitz property bounds the distance between mapped distributions by the original distance times the mapping's Lipschitz constant.Source: §A.1 · pp. 12–13
A.2 PROOFS OF THE MAIN THEORETICAL RESULTS
The appendix derives generative guarantees from correlation optimization under smoothness or strong-convexity conditions, then extends the analysis to latent-space generation and distributions without densities. It also points to the ICNN architectures used in the experiments.
A.2 PROOFS OF THE MAIN THEORETICAL RESULTS
- The basic single-potential method without regularization has analogous forward and inverse generative guarantees.The appendix states that the forward and inverse mappings push their source distributions to O(ε)-close target distributions.Source: §A.2 · pp. 15–16
- Better approximation of regularized correlations yields forward and inverse generated distributions that are O(ε)-close in W2 under the theorem’s conditions.The forward map pushes P close to Q, while the inverse map pushes Q close to P.Source: §A.2 · p. 16
A.3 FROM LATENT SPACE TO DATA SPACE
- For latent-space transport, decoded distribution error is bounded by contributions from latent-fit quality and autoencoder reconstruction loss.The decoder is assumed to be L-Lipschitz.Source: §A.3 · pp. 17–18
A.4 EXTENSION TO THE NON-EXISTENT DENSITY CASE
- The main theoretical results require finite second moments and densities on equal-dimensional spaces, while manifold-supported or discrete distributions may lack densities.The appendix proposes Gaussian smoothing and a de-smoothing bound whose bias depends on the noise statistics.Source: §A.4 · p. 18
B NEURAL NETWORK ARCHITECTURES
- Appendix B describes the general ICNN architecture and the DenseICNN and ConvICNN realizations used in experiments.Source: B NEURAL NETWORK ARCHITECTURES · p. 19
B.1 GENERAL INPUT-CONVEX ARCHITECTURE
The appendix builds convex potentials with input convex neural networks, then specifies dense and convolutional realizations for different tasks. It also directs readers to the experimental details and additional results.
B.1 GENERAL INPUT-CONVEX ARCHITECTURE
- The general ICNN combines linear and convexity-preserving blocks so that every neuron, including the output, is convex in the input.Quadratic activations between blocks are used as building blocks for fitting the identity mapping.Source: §B.1 · p. 19
B.2 DENSE INPUT CONVEX NEURAL NETWORK
- DenseICNN uses low-rank positive quadratic forms to reduce quadratic-layer storage from quadratic scaling to O(r · N_in · N_out).The resulting forms have rank at most r; full-rank quadratic layers remain practical when input dimension is low.Source: §B.2 · p. 20
B.3 CONVOLUTIONAL INPUT CONVEX NEURAL NETWORK
- ConvICNN serves image-processing tasks by mapping a 128 × 128 RGB image to a scalar whose input gradient is the generator.Its convolutional layers use 128 channels.Source: §B.3 · pp. 20–21, Figure 7 · p. 21
C EXPERIMENTAL DETAILS AND EXTRA RESULTS
- Appendix C organizes additional experiments on toy problems, latent-space transport, color transfer, domain adaptation, and style transfer.Source: C EXPERIMENTAL DETAILS AND EXTRA RESULTS · p. 21
C.1 GENERAL TRAINING DETAILS
Training initializes the potentials near inverse mappings and uses clipping, Adam, and a memory-saving gradient heuristic. The appendix also reports the cost of evaluating the objective and its gradients.
C.1 GENERAL TRAINING DETAILS
- Each task pretrains one potential toward the identity map, copies its weights to the other potential, then optimizes with mini-batch SGD, weight clipping, and Adam.Source: §C.1 · p. 21
- Disabling back-propagation through a repeated inverse-potential term reportedly speeds convergence and reduces memory and computation, and the authors use this heuristic in all experiments.Source: §C.1 · p. 21
- High-dimensional experiments add a second cycle-consistency regularizer, for which the authors state analogous guarantees but omit it from the proofs for simplicity.The added term is used for latent-space transport, domain adaptation, and style transfer.Source: §C.1 · p. 21
C.2 COMPUTATIONAL COMPLEXITY
- For DenseICNN, objective-gradient computation takes roughly 8–12× the time of evaluating ψθ(x), while evaluating either potential gradient takes roughly 3–4×.The appendix also characterizes objective and gradient evaluation as comparable up to a constant factor to evaluating ψθ(x).Source: §C.2 · p. 22
C.3 TOY EXPERIMENTS
Across these experiments, W2GN converges faster than competing minimax approaches in high-dimensional Gaussian transport, while toy, latent-space, and color-transfer examples demonstrate its behavior across several settings. The results also identify smoothness and cycle-monotonicity as relevant boundaries for some mappings.
C.3 TOY EXPERIMENTS
- For the 8-Gaussian mixture, the forward mapping’s overall structure appears by iteration 200, while the inverse mapping is learned around iteration 2000.The passage relates the delay to the inverse mapping’s larger Lipschitz constant as it unsqueezes dense Gaussian masses.Source: §C.3 · p. 22
- For a 100-Gaussian mixture, the model learns all modes without mode collapse, though the passage does not claim a perfect fit.The theoretical results require smooth P and Q; discontinuous-gradient ICNNs can produce torn mappings, and their theoretical generalization is uncertain without smoothing.Source: §C.3 · p. 22
C.4 GAUSSIAN OPTIMAL TRANSPORT DETAILS AND DISCUSSION
- W2GN, MM-1, and MM-2 achieve L2-UVP<3% even in R4096, while W2GN converges several times faster than the minimax approaches.The comparison attributes the speed difference to the inner optimization cycle in MM-1 and MM-2.Source: §C.4 · p. 24, Figure 12 · p. 25
C.5 LATENT SPACE OPTIMAL TRANSPORT DETAILS
- The fitted latent-space map is explicit and provides a tractable density, using the approximate inverse and the change-of-variables formula.Additional generated images and a PCA visualization of the fitted and autoencoder latent spaces are provided.Source: §C.5 · pp. 25–26, Figure 14 · p. 25, Figure 15 · p. 26
C.6 COLOR TRANSFER
- Mini-batch optimization removes palette-size limitations for color transfer, and results are presented for images of approximately 10 megapixels.The task maps one image’s color palette to another; the cited passage reports results but does not quantify their quality.Source: §C.6 · p. 26
C.7 DOMAIN ADAPTATION
The domain-adaptation experiment maps USPS features toward MNIST features in an unsupervised setting, using no label information. Its reported scores are comparable to Seguy et al.’s, and the mapping applies to new samples.
C.7 DOMAIN ADAPTATION
- The method addresses unsupervised domain adaptation with labels available only in the source domain, training the mapping as a gradient of a convex function that applies to new samples.The experiment transports USPS features toward MNIST features and evaluates a 1-NN classifier on the unmapped and mapped features.Source: §C.7 · pp. 26–28
- The reported domain-adaptation scores are comparable to those reported by Seguy et al.The authors did not reproduce Seguy et al.’s experiment because source code was unavailable; their comparison refers to that paper’s reported scores.Source: §C.7 · p. 28
- The compared approach is not directly applicable to out-of-training-sample examples and is infeasible for large datasets.The passage contrasts this limitation with the proposed method.Source: §C.7 · p. 28
C.8 IMAGE-TO-IMAGE STYLE TRANSFER
On Winter2Summer and Photo2Cezanne image datasets, the networks change colors while preserving image structure. The examples also show limits: some outputs make only a simple green-channel adjustment, and the authors suggest the desired manipulation may not be cycle monotone.
C.8 IMAGE-TO-IMAGE STYLE TRANSFER
- The networks change colors while preserving image structure, but the reported examples do not remove large snow masses or cover green trees with snow in the stated seasonal transformations.The authors suggest, without claiming an explanation, that the desired manipulation may not be cycle monotone.Source: §C.8 · p. 29
- Some poorly performing examples merely increase or decrease the green component, while more advanced image manipulations may be required.The cases are shown for 128 × 128 pixel images in Figure 21.Source: §C.8 · p. 29, Figure 21 · p. 30
- The experiments apply the method directly to images without specific preprocessing or feature extraction; the authors suggest feature-based or image-embedding OT as a possible direction for unexpected results.They identify this as a challenge for further research.Source: §C.8 · p. 29