---
kind: concept
slug: minibatch-optimal-transport
title: Minibatch optimal transport
summary: >-
  Estimates an optimal transport loss on random minibatch pairs so it can train a model. Balanced minibatch OT forces every sample to move, which pairs unrelated samples; unbalanced variants let mass go unmatched, with unbiasedness and concentration proved. Gains in domain adaptation are author-reported; one paper so far.
tags: [optimal-transport, minibatch, domain-adaptation, unbalanced-optimal-transport]
papers: ["2103.03606"]
related: [optimal-transport, domain-adaptation]
status: draft
updated_at: 2026-10-19
---

## 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 [arxiv: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 [arxiv: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
[arxiv:2103.03606 §3]. The unbalanced version lets each pair leave some mass
unmatched and charges a penalty for it [arxiv: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 [arxiv: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 [arxiv: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 [arxiv:2103.03606 §2].
Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

### 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 ε [arxiv:2103.03606 §2, eq. 2]. It is solved with a
generalised Sinkhorn iteration, whose complexity the paper quotes as
Õ(n²/ε) for the KL case [arxiv:2103.03606 §2]. The paper proves the minibatch
estimator is unbiased [arxiv:2103.03606 §4.1]. Its concentration bound depends
on the sample size n and minibatch size m, and not on the dimension d
[arxiv:2103.03606 §4.2]. Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

### 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 [arxiv:2103.03606 §2, eq. 3]. The paper
states it is convex, symmetric and positive definite when e^(−C/ε) is a
positive definite kernel [arxiv:2103.03606 §2]. The minibatch losses are still
not definite for 1 < m < n [arxiv:2103.03606 §4.1].
Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

### 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 [arxiv:2103.03606 §4.1, eq. 8]. Both are
unbiased. The incomplete estimator's deviation term shrinks as k grows, with
the rate given in [arxiv:2103.03606 §4.2, Thm. 1].
Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

## When to use which

| Situation | Use | Why |
|---|---|---|
| Masses are believed equal and every sample should be matched | Balanced OT, or UOT with large τ | Large τ recovers the balanced problem [arxiv:2103.03606 §2]. The paper argues balanced loss is fragile under outliers and imbalanced minibatches [arxiv:2103.03606 §3] |
| Class imbalance, partial overlap or outliers between the two sets | UOT with a KL penalty | The paper's bound on the unbalanced loss stays bounded under an outlier [arxiv:2103.03606 §3, eq. 4]. Author-reported partial DA average is 75.5 with JUMBOT against 60.9 for DeepJDOT [arxiv:2103.03606 tab. 1] |
| The loss must be zero when comparing a set with itself | Sinkhorn divergence | The plain entropic loss is not zero on identical sets [arxiv:2103.03606 §2] |
| Per-step compute is tight | k = 1 incomplete estimator | The paper reports state-of-the-art results with one minibatch pair per step [arxiv:2103.03606 §5.2]. The concentration bound loosens as k falls [arxiv:2103.03606 §4.2] |
| Batch size is small | UOT | Author-reported: JUMBOT accuracy stays about constant across batch sizes, while DeepJDOT drops by about 4 to 6 points at small batches [arxiv:2103.03606 app. C] |
| A full transport plan between the two sets is needed | Not the averaged minibatch plan | The averaged plan does not respect the marginals and is not an OT plan in general [arxiv:2103.03606 app. A] |
| The loss is trained by gradient descent | ε > 0 | With ε > 0 the optimal plan is unique and the loss is differentiable in the cost. With ε = 0 the paper needs Clarke derivatives [arxiv:2103.03606 app. B] |

## 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).

```pseudocode
# 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
  [arxiv:2103.03606 §4.1] [code: Domain_Adaptation/train_jumbot.py#L270].
- 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
  [code: digits/jumbot.py#L93-L108] [arxiv: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 [code: Domain_Adaptation/train_jumbot.py#L270-L273].
- The paper's domain-adaptation loss adds a label term to the feature cost
  and weights the transfer term by η3 [arxiv: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
[arxiv:2103.03606 tab. 1] and 70.1 with the best score over epochs
[arxiv: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 [arxiv:2103.03606 §5.2]. Compare numbers
only within one protocol.

## Papers

- [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606): defines minibatch UOT estimators with unbiasedness, concentration and gradient results, and applies them to domain adaptation as JUMBOT.