Optimal transport mapping via input convex neural networks

In plain words

Moving data from one distribution to another can be done in many ways. The cheapest way, in terms of total squared distance moved, is the optimal transport map. Finding it from samples is hard in high dimensions, because the usual mathematical formulation has a huge number of constraints that cannot be checked directly.

This paper's idea is that the best way to move points is the gradient of a convex function, a bowl-shaped function with no dips. The authors represent two such functions with neural networks that are built to stay convex, and train them against each other in a minimax game. After training, the gradient of one network tells every point where to go. Because a gradient can jump, the map can split one region into several separate targets without spreading points into the gaps in between, which standard generators tend to do.

The paper reports results on two-dimensional toy data and on Gaussian pairs in 784 dimensions, and it gives a small MNIST example in a 16-dimensional learned embedding.

TL;DR

The paper learns the quadratic-cost optimal transport map from samples without regularising the dual problem. It parametrises two convex functions with input convex neural networks (ICNNs) and trains them in a minimax game (Algorithm 1). The transport map is the gradient of the second function §3. Author-reported results show a small error in the transported mean on 784-dimensional Gaussian pairs tab. 1. Qualitatively, the method reproduces discontinuous maps that continuous generators do not §4.3.

Problem

Monge's problem asks for a map that transports one distribution to another with minimal expected quadratic cost §2. The Kantorovich dual is a supremum over pairs of functions that must satisfy an infinite family of pointwise constraints §2, eq. 3. Estimating it by sampling is hard to do without biasing the result §1, §3. Common fixes add a regulariser, which changes the problem and biases the map §2. GAN-style generators are continuous functions, and this produces spurious mass between disjoint target supports §4.3.

The method assumes the source distribution Q has a density on R^d, so that the optimal map is unique and given by a gradient §3, Theorem 3.1.

Method

The transport map is written as T(y) = grad g(y), where g is convex. The quadratic cost is one half of the squared Euclidean distance §2.

1. Semi-dual form. The dual of the transport cost can be written with a single convex function f and its convex conjugate f* §3, eq. 4. Direct estimation of f* is awkward because it needs a conjugate per sample [§3.1, Remark on Taghvaei et al.].
2. Minimax reformulation. The conjugate f* is replaced by a second convex function g. The paper shows that W_2^2 equals a supremum over convex f of an infimum over convex g, plus a constant, and that the optimal g has gradient equal to the optimal transport map §3, Theorem 3.2, eq. 5. The relaxation g in L^1 instead of convex g does not change the optimal value §3, Remark.
3. ICNN parametrisation. Both f and g are ICNNs §3.1, eq. 6. An ICNN layer is z_{l+1} = sigma_l(W_l z_l + A_l x + b_l), with W_0 = 0, and the network is convex in x if the W_l are non-negative, sigma_0 is convex, and every later sigma_l is convex and non-decreasing §3.1. The paper cites Chen et al. for the claim that ICNNs can approximate any convex function on a compact set §3.1.
4. Training. Algorithm 1 alternates K Adam steps on g with one Adam step on f, followed by clipping f's hidden-to-hidden weights at zero alg. 1. The g-objective adds a penalty for negative weights of g, relaxing the convexity constraint on g §4, eq. 8.

Pseudocode for Algorithm 1 in the notation of the paper, with the objective of eq. 9 written out:

# Minimax training of the potentials f and g  [alg. 1] [eq. 9]
# Input:  source Q, target P, batch size M, inner steps K, rounds T
# Output: f and g; the transport map is the gradient of g  [§3]

def train(Q, P, M, K, T):
    # Repeat for T rounds  [alg. 1]
    for t in range(T):
        # One batch of target samples X_1..X_M from P per round  [alg. 1]
        targets = sample(P, M)                      # (M, d)
        # K updates of g, each with a fresh source batch  [alg. 1]
        for k in range(K):
            # Source batch Y_1..Y_M from Q  [alg. 1]
            sources = sample(Q, M)                  # (M, d)
            # Transport map of the sources: gradient of g  [§3.1]
            maps = grad_g(g, sources)               # (M, d)
            # Objective of eq. 9 without the f(X) term  [eq. 9]
            loss_g = mean(f(maps)) - mean(sum(sources * maps, axis=1))
            # Penalty on negative weights of g  [eq. 8]
            loss_g = loss_g + penalty(g, lam)
            # Adam step on the parameters of g only  [alg. 1]
            g = adam_step(g, loss_g)
        # Maximise J over f, i.e. minimise mean f(X) - mean f(map)  [eq. 9]
        maps = grad_g(g, sources)                   # (M, d)
        loss_f = mean(f(targets)) - mean(f(maps))
        # Adam step on f, using the last source batch (not stated)  [alg. 1]
        f = adam_step(f, loss_f)
        # Clip the hidden-to-hidden weights of f at zero  [alg. 1] [§3.1]
        f = clip_hidden_weights(f, minimum=0)

Forward pass of one ICNN, with the convention of §3.1 (L affine layers, the last one producing a scalar):

# Forward pass of an ICNN  [§3.1] [app. C.1]
# Input:  x (M, d); A_l, b_l unconstrained; W_l >= 0 for l >= 1
# Output: scalar value per sample (M, 1)

def icnn(x):
    # Layer 0 has no W term (W_0 = 0); square of leaky ReLU  [app. C.1]
    hidden = square(leaky_relu(x @ A_0 + b_0, beta))    # (M, m)
    for l in range(1, L):
        # Non-negative path from the previous hidden layer  [§3.1]
        hidden = leaky_relu(x @ A_l + b_l + hidden @ W_l, beta)
    # The scalar output also passes through its activation  [§3.1]
    return hidden                                       # (M, 1)

Key equations, written in the paper's notation:

\max_{\theta_f:\, W_\ell \ge 0} \min_{\theta_g} \; J(\theta_f,\theta_g) + R(\theta_g) §4, eq. 9

R(\theta_g) = \lambda \sum_{W_\ell \in \theta_g} \lVert \max(-W_\ell, 0) \rVert_F^2 §4, eq. 8

J(\theta_f,\theta_g) = \frac{1}{M}\sum_{i=1}^{M} \big[ f(\nabla g(Y_i)) - \langle Y_i, \nabla g(Y_i)\rangle - f(X_i) \big] §4

Here X_i are samples from the target P and Y_i from the source Q §4, app. C.1.

Hyperparameters

The paper gives full hyperparameters only for the experiments in the appendix. Where the paper and the code differ, the table shows the paper value and the code default, and the Paper vs code section explains the gap.

Results

Author-reported. Setup: Gaussian source N(0, I_d) and target N(mu, I_d) with mu = alpha (1, ..., 1), d = 784, ICNNs trained with Algorithm 1 [§4.4, app. C.2]. The baseline is the closed-form optimal map T*(x) = x + mu, with exact squared W_2 distance alpha^2 d / 2 §4.4. Table 1 reports the error between the mean of the transported samples and the target mean §4.4, tab. 1.

The ± statistic is not defined in the paper. The repository README says the relative errors were computed over 4 trials README.md, but the paper does not state the number of trials.

Other results are qualitative. The paper gives no numeric score for the 2-D experiments. Figures compare the learned maps with Barycentric-OT, W1-LP and W2GAN on checkerboard and 8-Gaussian data, and show that the proposed map does not leave spurious points between modes §4.1, fig. 3. A second figure shows W1-LP and W2GAN maps changing with random initialisation, while the proposed map is said to be independent of initialisation §4.2, fig. 4. On high-dimensional tasks, Figure 5 shows the estimated W_2 distance converging to the exact value for alpha in {1, 5, 10}; the converged values are not given in the text §4.4, fig. 5a. On the MNIST latent task, the paper reports example couplings such as 1 -> 9, 2 -> 8 and 4 -> 9 §4.4, fig. 5c-d. The Gaussian-to-MNIST samples are shown but not scored §4.4, fig. 6.

Paper vs code

Code: the official-looking repository AmirTag/OT-ICNN (https://github.com/AmirTag/OT-ICNN) at 6caa9b982596a101b90a8a947d10f35f18c7de4e. The README calls it the code and dataset for the article and links the arXiv abstract README.md. The owner is a user account that appears to match author Amirhossein Taghvaei; this is inferred from the handle and not confirmed, and the paper does not link the repository. The repository has no LICENSE file, so code_license is recorded as "None" and no code is quoted here. Searches run: links in the paper and its arXiv comments (none); the arXiv abstract page (none); GitHub repository searches for "optimal transport input convex neural networks makkuva" (no results), "OT-ICNN optimal transport" (found the repository above and a third-party minimal PyTorch reimplementation, which was not used) and "input convex neural networks optimal transport minimax" (no results). The paper's only code link is to a baseline repository, which is not the authors' code. The high-dimensional scripts import optimal_transport_modules, which is not in the repository, so they cannot be run from the repository alone Gaussian_to_Gaussian.py.

Differences found:

1. Depth convention. The paper writes f as L layers with z_L = f, and sets L = 4 for the 2-D tasks §3.1, app. C.1. The 2-D code uses NUM_LAYERS = 4 hidden layers of width 64 and then appends a width-1 output, so five affine maps in total W2-minimax-tf.py and W2-minimax-tf.py.
2. Penalty scale. Eq. 8 uses the squared Frobenius norm §4, eq. 8. The 2-D code uses tf.nn.l2_loss, which is half the sum of squares, so the effective lambda is halved relative to the paper W2-minimax-tf.py. The high-dimensional scripts use the squared norm without the half Gaussian_to_Gaussian.py.
3. Resampling of Y. Algorithm 1 draws a fresh Y batch in each of the K inner steps alg. 1. The 2-D code does the same [code: 2_dim_experiments/W2-minimax-tf.py#L186-L191]. The high-dimensional code draws one Y batch per round and reuses it for all K inner steps Gaussian_to_Gaussian.py; the GMM script does the same GMM_to_GMM.py. The 2-D code also redraws X in every inner step, though the paper draws X once per round alg. 1 W2-minimax-tf.py.
4. Adam weight decay. The high-dimensional optimisers use weight decay 1e-5 Gaussian_to_Gaussian.py and vaeMNIST_to_vaeMNIST.py. The paper does not mention weight decay. The 2-D code does not use it W2-minimax-tf.py.
5. Initialisation. The 2-D code initialises A and W uniformly on [0, 0.1) and b to zero W2-minimax-tf.py. The high-dimensional code draws all parameters from a truncated normal with scale 1/sqrt(fan-in), which can give negative W entries before the first clipping step Gaussian_to_Gaussian.py. The paper does not state an initialisation.
6. Centre scale in 2-D. The paper's checkerboard and 8-Gaussian centres have unit scale app. C.1. The 2-D code multiplies centres by SCALE, whose default is 5.0 W2-minimax-tf.py and W2-minimax-tf.py. The README commands do not override it README.md.
7. Noise scale notation. The paper writes Z ~ N(0, 0.5 I) for the 8 Gaussians and Sigma = diag(0.2, 0.2) for the GMM app. C.1, app. C.2. The code multiplies standard normal samples by VARIANCE, so 0.5 and 0.2 are used as standard deviations W2-minimax-tf.py and GMM_to_GMM.py. If the paper meant covariances, the code uses a different spread.
8. GMM dimension and length. The paper uses d = 784 and 30 epochs for the GMM task app. C.2. The GMM script defaults to INPUT_DIM = 728 and 50 epochs GMM_to_GMM.py and GMM_to_GMM.py. The README runs the script with defaults README.md. Table 1's caption also says 728 while the text says 784 tab. 1 §4.4.
9. Code docstring. The 2-D class docstring describes ReLU layers, but the forward pass uses leaky ReLU with slope 0.2 W2-minimax-tf.py.
10. Command syntax. The README 2-D commands put spaces around =, as in --DATASET_X = 'checker_board_four' README.md. Written this way, argparse probably rejects the arguments, so use --DATASET_X=checker_board_four.

No discrepancy was found for the 2-D learning rate, batch sizes, K, T, the Adam betas, or the MNIST settings app. C.1, app. C.2.

Pitfalls

• Which batch the f-step uses is not stated. Algorithm 1 updates f without naming its batches alg. 1. The 2-D code draws fresh batches for the f-step W2-minimax-tf.py. The high-dimensional code reuses the last source batch Gaussian_to_Gaussian.py.
• Second-order autodiff is needed. The g-objective depends on grad g(Y), so training differentiates through a gradient. The code does this with tf.gradients and torch.autograd.grad(..., create_graph=True) W2-minimax-tf.py and Gaussian_to_Gaussian.py.
• Positivity is enforced in two different ways. f's hidden-to-hidden weights are clipped at zero after each update alg. 1. g's are not clipped in the 2-D code. They receive a penalty instead §4, eq. 8. In the high-dimensional code, g is clipped only when lambda = 0 [code: High_dim_experiments/Gaussian_to_Gaussian.py#L394-L396].
• The scalar output has an activation. The paper's layer equation applies sigma to the last layer too §3.1. The 2-D code does the same W2-minimax-tf.py. Implementers often leave the output unactivated, which changes the function class.
• The quantity is W_2^2 with the 1/2 factor, not W_2 §2, eq. 2. The 2-D code's W2 estimate is the mean of f(grad g(Y)) - f(X) - <Y, grad g(Y)> plus 0.5 E|X|^2 + 0.5 E|Y|^2 W2-minimax-tf.py.
• Exact check. The repository's exact solver uses the POT library with ot.emd on uniform samples W2-minimax-tf.py. Use it on small batches to validate an implementation.
• Stability bound needs strong convexity. Theorem 3.3 bounds the map error only if f is alpha-strongly convex §3.2. The paper does not give alpha for its experiments and does not describe how to enforce it.
• Cost. Each round runs K inner steps, each with a forward pass, a gradient and a second backward pass, on batches of 60 to 1024 app. C.1, app. C.2. The abstract says that scaling to large datasets is challenging [abstract]. The experiments are 2-D, Gaussian pairs in 784 dimensions, and a 16-dimensional VAE latent §4.
• Hidden activation slope in high dimensions. The paper gives beta = 0.2 for the 2-D case app. C.1 but not for the high-dimensional leaky ReLU app. C.2. The ICNN module that defines it is missing from the repository Gaussian_to_Gaussian.py.

Lineage

• Builds on: the semi-dual ICNN approach of the paper's reference taghvaei20192 §2, §3.1, remark, and input convex neural networks of amos2016input §1, §3.1.
• Compared against: Barycentric-OT (seguy2017large), W1-LP (petzka2017regularization) and W2GAN (leygonie2019adversarial) §4.1, app. C.1. The paper also discusses the conjugate-based method of taghvaei20192, which needs an inner convex solve per sample §3.1, remark.
• Superseded by: not known from content already in this wiki.