Optimal transport

Overview

Given samples from a source distribution and a target distribution, optimal transport asks for the map that moves each source point to a target point at the lowest total cost. With squared Euclidean cost, the problem is Monge's formulation, and its dual is a supremum over pairs of functions that must satisfy an infinite family of pointwise constraints [arxiv:1908.10962 §2, eq. 3]. The shared difficulty is estimating this from samples: the constraints are hard to enforce, and common regularisers change the problem and bias the map 1908.10962 §1, §2. Training a model on the transport cost adds a second difficulty: the exact problem is too expensive on a full dataset, so it is usually estimated on random minibatch pairs, and balanced minibatch estimates force every sample to move 2103.03606 §1, §3.

The variants below answer different questions. One returns a map from the source to the target. Another is a loss, estimated on minibatches, that a model can be trained with. The third uses a transport plan as a soft assignment between the positions of two images, with no map and no training loss. The fourth applies the one-dimensional case as a distance between two spectra, to score predictions; it is neither a map nor a loss. The last two use OT as a reward for reinforcement learning: one matches keys to fingertips at each step, the other compares an agent's frames with expert frames. A seventh variant uses the same entropic solver for clustering rather than matching: it assigns training pixels to class prototypes under an equal-size constraint. An eighth variant relaxes the exact problem by dropping one marginal constraint at a time, so that it is cheap enough to evaluate at every step of an image optimisation.

Variants

Minimax potentials with input convex networks

The invariant that this variant relies on is that, for a source with a density, the optimal map under quadratic cost is the gradient of a convex function 1908.10962 §3, Theorem 3.1. Searching over convex functions is therefore enough.

The dual is rewritten in a semi-dual form with one convex function f and its conjugate 1908.10962 §3, eq. 4. The conjugate is replaced by a second convex function g, and the optimal g has a gradient equal to the transport map 1908.10962 §3, Theorem 3.2, eq. 5. Both functions are parametrised as input convex neural networks and trained by alternating updates 1908.10962 §3.1, alg. 1. The authors report that the learned map can split a region into separate targets without placing mass in the gaps, which continuous generators do not do 1908.10962 §4.3. Card: Optimal transport mapping via input convex neural networks.

Unbalanced minibatch OT

Balanced minibatch OT averages the balanced problem over small random subsets. Each subset must send all of its mass, so it pairs samples that do not correspond 2103.03606 §1, fig. 1. The variant replaces the balanced problem in each minibatch with unbalanced OT: the marginals get a KL penalty of weight τ, and an entropic term of weight ε is added [arxiv:2103.03606 §2, eq. 2]. It is solved with a generalised Sinkhorn iteration that is proved to converge 2103.03606 §2.

The authors prove that the averaged estimates are unbiased for the expected minibatch value, both over all minibatch pairs and over k random pairs 2103.03606 §4.1, eq. 6–8. They also give a concentration bound whose rate depends on the batch and sample sizes, not on the dimension 2103.03606 §4.2, Theorem 1. Stochastic gradient descent on the loss converges to Clarke critical points, with ε > 0 for a unique, differentiable plan 2103.03606 §4.3, Theorem 2; app. B. For outliers, the paper bounds the unbalanced loss under an added outlier where the balanced loss grows with the outlier's distance 2103.03606 §3, eq. 4–5, fig. 2.

The application is JUMBOT, a domain-adaptation loss that matches source and target minibatches by features and by the source label, with one minibatch pair per step 2103.03606 §5.2, eq. 9–10. Card: Unbalanced minibatch optimal transport; applications to Domain Adaptation. The minibatch-loss side has its own page: Minibatch optimal transport. The unbalanced problem itself, with its divergence and Sinkhorn form, is on Unbalanced optimal transport.

Sinkhorn assignment in a dense matcher

This variant is neither a map nor a training loss. A Sinkhorn solver runs on the pairwise feature scores of two images, with a learned "dustbin" score for positions that have no partner. Its confidence matrix selects the coarse matches 2104.00680 §3.3, §3.6. The paper's other option is dual softmax, the product of two softmaxes, which needs no iterative solver 2104.00680 §3.3. The OT option runs three Sinkhorn iterations 2104.00680 §3.6.

The two options are close in the paper's tables. On ScanNet pose AUC at 10°, the OT model scores 40.39 against 40.8 for dual softmax. On MegaDepth at 10°, it scores 67.14 against 69.19 2104.00680 tab. 2, tab. 3. The authors time the OT option at 130 ms against 116 ms for dual softmax at 640×480 2104.00680 §3.6. Card: LoFTR: Detector-Free Local Feature Matching with Transformers.

Sinkhorn equipartition clustering for prototypes

Here the Sinkhorn form is used inside a training loss for segmentation, not to match two sets of points. For the labelled pixels of one class in a batch, the paper seeks an assignment to that class's K prototypes. Each pixel gets one prototype, and each prototype receives about N/K pixels (equipartition). The constraints are relaxed with an entropy term weighted by κ, and the solution is a row and column scaling of an exponentiated similarity matrix, computed by a few Sinkhorn–Knopp iterations 2203.15102 §3, eq. 8–10. The paper's default is κ = 0.05 with three iterations 2203.15102 §5.1, alg. 1. The prototypes are then updated by momentum toward the normalised mean of their assigned pixels and receive no gradient 2203.15102 §3, eq. 14.

Two details differ between paper and code. The paper takes a soft Sinkhorn solution and then a hard argmax, while the released code draws a hard Gumbel-softmax sample for the prototype update. Equipartition is also enforced within each batch, so a class with few pixels in a batch is still spread over K prototypes; the paper does not discuss this [arxiv:2203.15102 Paper vs code; §3]. On ADE20K val with four reproduced baselines, the authors report +0.6 to +1.2 mIoU, for example 42.0 to 43.0 for HRNetV2-W48 and 50.9 to 51.7 for SegFormer MiT-B4 2203.15102 tab. 1. On ADE20K-847 the paper reports 19.4 mIoU against 17.7 for a parametric head 2203.15102 tab. 4. Only a Cityscapes HRNet-W48 configuration is released 2203.15102 Paper vs code. Card: Rethinking Semantic Segmentation: A Prototype View. The segmentation side is on Semantic segmentation, and the prototype side on Prototype learning.

Relaxed EMD as a style loss

This variant does not compute the exact earth mover's distance. Between two sets of VGG16 features, the exact value is too costly to evaluate at every gradient step, so the paper relaxes it. Each step keeps one marginal constraint and drops the other, and the relaxed value is the larger of the two one-sided averages of nearest-neighbour costs 1904.12785 §2.2, eq. 5–8. The ground metric is cosine distance. The authors report that a Euclidean ground metric gave significantly worse results 1904.12785 §2.2, eq. 9. Cosine distance ignores feature magnitude, so the style term also matches the mean and covariance of the two feature sets 1904.12785 §2.2, eq. 10.

The relaxed term is one part of a per-image optimisation. The output image is refined against fixed VGG16 features, with a self-similarity content term and a coarse-to-fine schedule 1904.12785 §2, §2.3, §2.5, eq. 1. The relaxation is a lower bound on the exact value. Over 900 pairs its mean ratio to the exact EMD is 0.60, with a standard deviation of 0.04 1904.12785 §4.3. The authors say a tighter approximation may give better results, and leave it to future work 1904.12785 §5.

The authors report a 662-worker study, with three regimes, in which their method gets higher style scores than four prior methods at a given content score. The text gives no numeric table 1904.12785 §4.1, fig. 8. The runtime is 95 s for a 512 px image, against 33 s for Gatys et al. on the authors' hardware 1904.12785 tab. 1. Card: Style Transfer by Relaxed Optimal Transport and Self-Similarity. The style side is on Neural style transfer.

One-dimensional transport distance as a metric

This variant is not a map, a plan or a loss. The NPSS score for human motion compares the power spectrum of each true joint trajectory with that of the prediction, and it is used to score predictions, not to train them 1809.03036 §3.5. The paper's eq. 8 is an L1 norm of the two normalised spectra, and the paper gives it the symbol emd. The released code instead takes the L1 norm of their cumulative sums, which in one dimension is the earth mover's distance between the two normalised spectra. The code also weights each joint by its unnormalised power, where the paper's weight is one by construction 1809.03036 eq. 8–9; Paper vs code. Card: A Neural Temporal Model for Human Motion Prediction. The motion side is on Motion prediction.

Per-step finger assignment as a reward

This variant uses a hard matching, not an entropic plan. At each control step, every key to press gets one fingertip, each fingertip takes at most one key, and the total Euclidean distance is minimised 2408.11048 §4.1, eq. 1. The paper uses this matching in place of human fingering labels. The reward term is one exponential of the matched total, above a threshold of 0.01; the constant in the exponent is not stated 2408.11048 §4.1, eq. 2. The authors report that this reward matches a human-fingering RL baseline on two songs and beats a no-fingering baseline by a wide margin [arxiv:2408.11048 §4.2, fig. 2]. These results come from simulated Shadow hands with one RL agent per song, and the paper gives no per-song table. Card: RP1M: A Large-Scale Motion Dataset for Piano Playing with Bi-Manual Dexterous Robot Hands. The card's Paper vs code section says the released code computes the reward in a different form, so use the paper's equation only with that in mind. The imitation side is on Imitation learning.

Temporal OT reward from demonstrations

Here the transport plan is between the agent's frames and one or two expert videos, and each agent step is rewarded with its negative cost under the plan 2410.21795 §3.1, eq. 5. The paper adds two order-aware changes: the cost of a frame pair is averaged over a short window of consecutive frames 2410.21795 §3.2.2, eq. 6, and the plan is restricted to a band around the diagonal §3.2.3, eq. 9. The reward is then used in a standard off-policy RL update 2410.21795 alg. 1. The authors report an average success of 61.1 over nine Meta-World tasks with two demos, against 35.9 for ADS and 20.8 for the sparse task reward. The standard deviations are large on several tasks 2410.21795 tab. 1. The method assumes the expert moves at about the agent's speed, and the authors report a drop when the expert is three or four times faster 2410.21795 §3.2.3, app. A.4, tab. 3. Card: Robot Policy Learning with Temporal Optimal Transport Reward. The reward-design side is on Reward design, and the demonstration side on Learning from demonstration. The card's code has no license file, so only the paper is cited here.

When to use which

Choose by what the task needs, and only the claims the cards support.

Implementation sketch

The minimax loop of the ICNN variant is in its card 1908.10962 alg. 1. The core of the minibatch variant is the estimator below. Learning rates, batch sizes and the other settings are in the cards and in the minibatch concept page.

# Minibatch UOT estimate from k random minibatch pairs  [arxiv:2103.03606 eq. 8]
# Inputs: X (n, d) source samples, Y (n, d) target samples, batch size m
#         k pairs to draw, tau marginal penalty, eps > 0 entropic weight
# Output: scalar estimate of the expected minibatch value

def minibatch_uot(X, Y, m, k, tau, eps):
    total = 0
    for _ in range(k):
        # Two index sets of size m, drawn without replacement  [§4.1]
        I = sample_indices(len(X), m)
        J = sample_indices(len(Y), m)
        # Ground cost between the two minibatches  [eq. 1]
        C = cost(X[I], Y[J])                        # (m, m)
        # Unbalanced OT value with uniform weights  [§2] [eq. 2]
        total = total + uot_value(C, eps, tau)
    return total / k

Contradictions

None known between the cards: they report different tasks and share no numbers. Within the cards, five points need care:

• The minibatch paper gives Office-Home JUMBOT as 70.0 with end-of-training scores, and 70.1 with the best score over epochs in its appendix table. The authors say the protocols are not comparable [arxiv:2103.03606 §5.2, app. C, Table 4].
• The style-transfer paper's α schedule halves four times from 16.0, which ends at 2, but the text says the last scale uses 1.0. The code ends at 2.0 1904.12785 §2.5; Paper vs code. Use the card's pitfalls section before choosing a final α.
• The ICNN paper's Table 1 caption gives the dimension as 728, while the text gives 784 1908.10962 tab. 1, §4.4.
• The LoFTR ablation's full model (tab. 6: 20.06 / 40.8 / 57.62 at 5°, 10°, 20°) does not match its OT row in tab. 2 (21.51 / 40.39 / 57.96), although the text says the ablation uses OT 2104.00680 tab. 2, tab. 6, §4.4. Compare the ablation rows only with each other.
• The NPSS paper's eq. 8 is an L1 norm of normalised spectra, but the released code computes an L1 norm of cumulative sums, a different quantity 1809.03036 eq. 8; Paper vs code. Scores from one cannot be compared with scores from the other.
• The TemporalOT paper's Table 3 gives Button-press at 1× as 89.0 (3.8). That matches the ADS entry in its Table 1, not TemporalOT's 92.4 (3.6) 2410.21795 tab. 1, tab. 3. Use Table 1 for TemporalOT and treat the Table 3 value as uncertain.

Differences between a paper and its code are in the cards' "Paper vs code" sections, not here.

Papers

• Optimal transport mapping via input convex neural networks: learns the quadratic-cost transport map from samples with two input convex networks trained in a minimax game.
• Unbalanced minibatch Optimal Transport; applications to Domain Adaptation: replaces balanced minibatch OT with unbalanced OT, proves unbiasedness and concentration, and applies it to domain adaptation (JUMBOT).
• LoFTR: Detector-Free Local Feature Matching with Transformers: uses a Sinkhorn OT solver with a dustbin score as one of two confidence options for a detector-free matcher; dual softmax is the other.
• Rethinking Semantic Segmentation: A Prototype View: replaces learned class weights with prototypes found by Sinkhorn equipartition clustering of training-pixel embeddings.
• Style Transfer by Relaxed Optimal Transport and Self-Similarity: optimisation-based style transfer whose style term is a relaxed EMD over VGG16 features, with a self-similarity content term.
• A Neural Temporal Model for Human Motion Prediction: proposes NPSS, a power-spectrum score for predicted motion whose distance is a one-dimensional transport distance in the released code.
• RP1M: A Large-Scale Motion Dataset for Piano Playing with Bi-Manual Dexterous Robot Hands: a simulated piano dataset whose fingering reward is an optimal-transport matching between keys and fingertips, in place of human fingering labels.
• Robot Policy Learning with Temporal Optimal Transport Reward: an order-aware OT reward from one or two expert videos, with a context cost and a banded masked Sinkhorn plan.