---
kind: concept
slug: unbalanced-optimal-transport
title: Unbalanced optimal transport
summary: >-
  Optimal transport in which mass may stay unmatched: each marginal is penalised (KL in the paper) instead of forced to equal the input mass, and an entropic term makes the plan solvable by Sinkhorn-type iterations. Gives a bounded response to outliers where balanced OT grows, and a debiased Sinkhorn divergence. One paper so far; domain-adaptation gains are author-reported.
tags: [optimal-transport, sinkhorn, minibatch, domain-adaptation]
papers: ["2103.03606"]
related: [optimal-transport, minibatch-optimal-transport, domain-adaptation]
status: draft
updated_at: 2026-10-19
---

## Overview

Balanced optimal transport moves all of the source mass onto the target mass
at the lowest cost, so every source sample must be sent somewhere
[arxiv:2103.03606 §2]. Unbalanced optimal transport (UOT) drops that hard
constraint. The plan's two marginals are compared with the input measures
through a penalty of weight τ, using a Csiszár divergence D_φ. An entropic
term with weight ε keeps the problem smooth and solvable [arxiv:2103.03606 §2]:

$$ \operatorname{OT}_\phi^{\tau,\varepsilon}(\alpha,\beta,c) = \min_{\pi} \int c\,d\pi + \varepsilon\,\mathrm{KL}(\pi\,|\,\alpha\otimes\beta) + \tau\big(D_\phi(\pi_1\|\alpha) + D_\phi(\pi_2\|\beta)\big) $$

The invariant the variants share is that the penalty is finite. As τ grows
with equal input masses, the problem approaches balanced OT
[arxiv:2103.03606 §2]. The paper takes φ convex and lower semicontinuous with
φ(1) = 0 and a positive recession constant; the KL case is
φ(x) = x log x − x + 1 [arxiv:2103.03606 §2, app. A].

Unbalanced OT matters in training because its minibatch estimates stay
unbiased and well-behaved. The minibatch estimator is covered on
[minibatch-optimal-transport](/concepts/minibatch-optimal-transport).

## Variants

### Unbalanced OT loss [arxiv:2103.03606 §2]

The raw problem above, with ε > 0. With ε > 0 the optimal plan is unique and
the loss is differentiable in the cost. With ε = 0 it is not, and the paper
uses Clarke derivatives for its gradient results [arxiv:2103.03606 §4.3,
app. B]. The loss is not a metric: with the entropic term, OT(β, β) is not
zero [arxiv:2103.03606 §2]. Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

### Unbalanced Sinkhorn divergence [arxiv:2103.03606 §2]

A debiased version. It adds a mass-difference term and subtracts half of each
self-transport term [arxiv:2103.03606 §2]:

$$ S_\phi^{\tau,\varepsilon}(\alpha,\beta,c) = \operatorname{OT}_\phi^{\tau,\varepsilon}(\alpha,\beta,c) + \frac{\varepsilon}{2}(m_\alpha-m_\beta)^2 - \frac12\operatorname{OT}_\phi^{\tau,\varepsilon}(\alpha,\alpha,c) - \frac12\operatorname{OT}_\phi^{\tau,\varepsilon}(\beta,\beta,c) $$

The paper states that S is convex, symmetric and positive definite when
e^(−C/ε) is a positive definite kernel [arxiv:2103.03606 §2]. Use it when you
need a divergence that is zero for identical inputs. Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

### Minibatch estimators [arxiv:2103.03606 §4.1–4.2]

Both forms can be averaged over minibatch pairs. The complete average over all
pairs is unbiased. An incomplete average over k random pairs is cheaper and
also unbiased, with concentration bounds that depend on k [arxiv:2103.03606
§4.1–4.2]. JUMBOT uses k = 1 [arxiv:2103.03606 §5.2]. See the
[minibatch page](/concepts/minibatch-optimal-transport) for the full estimator.

### Label-aware cost for domain adaptation (JUMBOT) [arxiv:2103.03606 §5.2]

The cost between a labelled source sample and an unlabelled target sample is a
weighted sum of their squared feature distance and the cross-entropy between
the source label and the target prediction. The UOT plan on that cost is
combined with a source classification loss [arxiv:2103.03606 §5.2]. The paper
reports this as its domain-adaptation method. Card: [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606).

## When to use which

| Situation | Use | Why |
|---|---|---|
| Source and target have different class proportions, or some samples are outliers | UOT loss with finite τ | Unmatched mass is penalised, not forced across. The paper shows that balanced OT's bound on an outlier's effect grows with its displacement, while the UOT bound stays bounded [arxiv:2103.03606 §3, eq. 4–5]. The class-imbalance evidence is a qualitative 2D gradient-flow experiment [arxiv:2103.03606 §5.1] |
| You train a model on minibatch estimates and need unbiasedness | UOT loss with ε > 0 | Gives unique plans and a differentiable loss, so SGD is justified [arxiv:2103.03606 §4.3] |
| You need a divergence that is zero for identical inputs | Unbalanced Sinkhorn divergence | The plain loss is not zero at identical inputs [arxiv:2103.03606 §2] |
| Equal masses, and an exact transport plan is the goal | Balanced OT | Recovered as τ grows [arxiv:2103.03606 §2]. The paper does not argue for UOT here |
| Choosing ε | Smallest ε your solver handles stably, and ε > 0 | The paper's quoted Sinkhorn complexity is Õ(n²/ε) [arxiv:2103.03606 §2]. At ε = 0 the plan is not unique and the loss is not differentiable [arxiv:2103.03606 app. B] |
| Choosing τ | Start near the paper's settings below, then sweep | The appendix text reports that very small τ causes negative transfer, that accuracy peaks near τ = 1 on digits, and that large τ enforces the marginals at a small accuracy cost [arxiv:2103.03606 app. C]. These are text, not tables |

Paper's settings for reference. They are author-chosen, not general defaults
[arxiv:2103.03606 app. C]:

| Setting | τ | ε | Minibatch m |
|---|---|---|---|
| Digits | 1 | 0.1 | 500 |
| Office-Home DA | 0.5 | 0.01 | 65 |
| VisDA | 0.3 | 0.01 | 72 |
| Office-Home partial DA | 0.06 | not stated | 65 |

The paper does not state a Sinkhorn iteration cap, a tolerance, or a
log-domain stabilisation. Check for underflow of exp(−C/ε) yourself when ε is
small and costs are large [arxiv:2103.03606 §2, app. C].

## Implementation sketch

Solving the plan is a black box here. The paper uses a generalised Sinkhorn
iteration that it proves converges [arxiv:2103.03606 §2]. The step that
matters is that the plan is held fixed when differentiating: the derivative of
the loss with respect to the cost is the plan [arxiv:2103.03606 app. B].

```pseudocode
# UOT loss on one minibatch pair  [arxiv:2103.03606 §2] [§4.1]
# Inputs:  xs (m, ...) source batch, xt (m, ...) target batch
# Output:  scalar loss, differentiable in the cost
def uot_loss(xs, xt, tau, eps):
    # Pairwise ground cost between the two batches  [§4.1, eq. 1]
    cost = pairwise_cost(xs, xt)                    # (m, m)
    # Mass per sample. The paper's vectors carry m_α/n each  [§4.1]
    a = ones(m) / m                                 # (m,)
    b = ones(m) / m                                 # (m,)
    # Entropic plan with penalised marginals  [§2, eq. 2]
    plan = unbalanced_sinkhorn(a, b, cost.detach(), eps, tau)
    # Plan held fixed; the gradient is the plan  [app. B]
    return sum(plan * cost)
```

```pseudocode
# Unbalanced Sinkhorn divergence on the same batches  [§2, eq. 3]
# Cross term minus half of each self term; the mass term is zero here
# because both batches carry the same mass
def uot_divergence(xs, xt, tau, eps):
    ot_st = uot_loss(xs, xt, tau, eps)
    ot_ss = uot_loss(xs, xs, tau, eps)
    ot_tt = uot_loss(xt, xt, tau, eps)
    return ot_st - 0.5 * ot_ss - 0.5 * ot_tt
```

Notes on the sketch. The sketch uses unit total mass, which matches the
paper's form only when the total mass is 1 [arxiv:2103.03606 §4.1]. With
unequal masses, add the mass term (ε/2)(m_α − m_β)² from [§2, eq. 3].

## Contradictions

None across papers; this page has one source. Within the paper, the main
tables report end-of-training scores, while appendix Table 4 reports
best-over-epochs scores for Office-Home. JUMBOT averages 70.0 in the main
table and 70.1 in the appendix table [arxiv:2103.03606 §5.2, app. C]. Do not
compare the two. The authors say their ROT numbers differ from the original
ROT paper because of protocol [arxiv:2103.03606 §5.2]. That original paper is
not in the wiki, so this cannot be checked here.

## Papers

- [Unbalanced minibatch Optimal Transport; applications to Domain Adaptation](/papers/2103.03606): defines unbalanced OT with KL marginal penalties and its minibatch estimators, proves unbiasedness, concentration and gradient properties, and applies it to domain adaptation as JUMBOT.