Minibatch optimal transport

Overview

Optimal transport (OT) compares two sets of samples by the cheapest way to move one onto the other. The exact problem is too costly to run on full datasets inside a training loop, so the loss is estimated on small random minibatches drawn from each set and averaged 2103.03606 §4.1. The invariant the paper proves for its unbalanced estimators is that they are unbiased for their full-data expectation and concentrate around it as the number of minibatch pairs grows 2103.03606 §4.1–§4.2, Thm. 1.

The weak point of the balanced version is that each minibatch pair must send all of its mass. Small batches then match samples that do not correspond 2103.03606 §3. The unbalanced version lets each pair leave some mass unmatched and charges a penalty for it 2103.03606 §2.

Variants

Balanced minibatch OT

Each minibatch pair is solved as a balanced OT problem and the results are averaged. The paper's toy example shows samples from one source cluster paired with two different target clusters 2103.03606 §1, fig. 1. Balanced OT is also sensitive to a single outlier: its loss grows as the outlier moves away from the support of the other distribution 2103.03606 §3, eq. 5. The balanced problem is the limit of the unbalanced one as the marginal penalty τ goes to infinity, with equal masses 2103.03606 §2. Card: Unbalanced minibatch Optimal Transport; applications to Domain Adaptation.

Unbalanced minibatch OT with a KL marginal penalty

The per-pair problem adds a KL penalty of weight τ on each marginal and an entropic term of weight ε 2103.03606 §2, eq. 2. It is solved with a generalised Sinkhorn iteration, whose complexity the paper quotes as Õ(n²/ε) for the KL case 2103.03606 §2. The paper proves the minibatch estimator is unbiased 2103.03606 §4.1. Its concentration bound depends on the sample size n and minibatch size m, and not on the dimension d 2103.03606 §4.2. Card: Unbalanced minibatch Optimal Transport; applications to Domain Adaptation.

Sinkhorn divergence

The entropic term means OT of a set with itself is not zero, so the loss is not a proper distance. The Sinkhorn divergence subtracts the self-transport terms and adds a mass-difference term 2103.03606 §2, eq. 3. The paper states it is convex, symmetric and positive definite when e^(−C/ε) is a positive definite kernel 2103.03606 §2. The minibatch losses are still not definite for 1 < m < n 2103.03606 §4.1. Card: Unbalanced minibatch Optimal Transport; applications to Domain Adaptation.

Complete versus incomplete estimator

The complete estimator averages over every pair of minibatches [arxiv:2103.03606 §4.1, eq. 6]. It is combinatorially expensive, so the incomplete estimator averages over k randomly drawn pairs 2103.03606 §4.1, eq. 8. Both are unbiased. The incomplete estimator's deviation term shrinks as k grows, with the rate given in 2103.03606 §4.2, Thm. 1. Card: Unbalanced minibatch Optimal Transport; applications to Domain Adaptation.

When to use which

Implementation sketch

The common core is the incomplete estimator with uniform weights. The variants differ in the per-pair loss (balanced, unbalanced, or Sinkhorn divergence).

# Minibatch UOT loss from k random minibatch pairs [arxiv:2103.03606 §4.1]
# Inputs:  X (n, d) source samples, Y (n, d) target samples
# Output:  scalar estimate of the expected minibatch loss

def minibatch_uot(X, Y, m, k, tau, eps):
    total = 0
    for _ in range(k):
        # Draw one index tuple per side, without replacement  [§4.1, eq. 8]
        I = sample_without_replacement(n, m)
        J = sample_without_replacement(n, m)
        # Ground cost between the two minibatches  [§4.1, eq. 1]
        C = pairwise_cost(X[I], Y[J])                 # (m, m)
        # Per-pair unbalanced loss with uniform weights  [§2] [§4.1]
        total += uot_loss(u, u, C, tau, eps)          # scalar
    return total / k

Notes on the implementation:

• The weights are m_α/m per sample, so they carry the mass of the source distribution, not 1/m. Unit-mass weights match only when the mass is 1 2103.03606 §4.1 train_jumbot.py.
• The plan is computed from a detached cost, and the loss is the plan-weighted sum of the cost. The plan is treated as a constant in the gradient jumbot.py 2103.03606 app. B.
• The paper does not state a Sinkhorn iteration cap, tolerance or log-domain stabilisation. With small ε and float32 costs, underflow of exp(−C/ε) should be checked train_jumbot.py.
• The paper's domain-adaptation loss adds a label term to the feature cost and weights the transfer term by η3 2103.03606 §5.2, eq. 9–10. The label term is absent from the generic estimator above.

Contradictions

Only one paper so far, so there are no cross-paper disagreements. Within the paper, Office-Home JUMBOT is 70.0 with the end-of-training score 2103.03606 tab. 1 and 70.1 with the best score over epochs 2103.03606 app. C, Table 4. The authors state the two protocols are not comparable. The authors also say their ROT numbers differ from the original ROT paper because the protocol differs 2103.03606 §5.2. Compare numbers only within one protocol.

Papers

• Unbalanced minibatch Optimal Transport; applications to Domain Adaptation: defines minibatch UOT estimators with unbiasedness, concentration and gradient results, and applies them to domain adaptation as JUMBOT.