Unbalanced optimal transport

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

Variants

Unbalanced OT loss 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 2103.03606 §2. Card: Unbalanced minibatch Optimal Transport; applications to Domain Adaptation.

Unbalanced Sinkhorn divergence 2103.03606 §2

A debiased version. It adds a mass-difference term and subtracts half of each self-transport term 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 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.

Minibatch estimators 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 2103.03606 §5.2. See the minibatch page for the full estimator.

Label-aware cost for domain adaptation (JUMBOT) 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 2103.03606 §5.2. The paper reports this as its domain-adaptation method. Card: Unbalanced minibatch Optimal Transport; applications to Domain Adaptation.

When to use which

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

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 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 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 2103.03606 app. B.

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

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