PaperlayerView exact revision on arXiv

Source-linked AI summary

Unbalanced minibatch Optimal Transport; applications to Domain Adaptation

Kilian Fatras, Thibault Séjourné, Nicolas Courty, Rémi Flamary

arXiv:2103.03606v1cs.LGmath.STstat.ML

TL;DR

Large-scale OT is computationally expensive, and minibatch approximations can introduce undesirable couplings through marginal constraints and subsampling. The paper replaces balanced minibatch OT with unbalanced OT, establishing theoretical properties and evaluating it for domain adaptation. Its experiments report state-of-the-art performance across challenging domain adaptation tasks.

  • Problem

    OT is costly on large datasets, while minibatch OT can amplify undesirable coupling effects caused by marginal constraints and subsampling.

  • Method

    The paper applies unbalanced OT with relaxed marginal constraints at the minibatch level and studies its statistical and optimization properties.

  • Results

    JUMBOT is state of the art on 9 of 12 domain adaptation tasks and averages 2% above competitors.

  • Takeaways & Limitations

    Minibatch unbalanced OT provides a more efficient distributional comparison and reaches state-of-the-art performance on challenging domain adaptation problems.

  • Takeaways & Limitations

    The analysis assumes compactly supported distributions, which bounds the ground cost.

Abstract

from arXiv · show

Optimal transport distances have found many applications in machine learning for their capacity to compare non-parametric probability distributions. Yet their algorithmic complexity generally prevents their direct use on large scale datasets. Among the possible strategies to alleviate this issue, practitioners can rely on computing estimates of these distances over subsets of data, {\em i.e.} minibatches. While computationally appealing, we highlight in this paper some limits of this strategy, arguing it can lead to undesirable smoothing effects. As an alternative, we suggest that the same minibatch strategy coupled with unbalanced optimal transport can yield more robust behavior. We discuss the associated theoretical properties, such as unbiased estimators, existence of gradients and concentration bounds. Our experimental study shows that in challenging problems associated to domain adaptation, the use of unbalanced optimal transport leads to significantly better results, competing with or surpassing recent baselines.

1. Introduction

OT provides a useful distributional loss but is expensive at scale, motivating minibatch approximations. The paper proposes unbalanced minibatch OT to reduce sampling-induced undesirable couplings while retaining scalability and theoretical properties.

  • Motivation: OT losses compare empirical distributions but have cubic computational cost, while entropic regularization reduces complexity to approximately quadratic.The stated complexities are O(n^3log(n)) for OT and Õ(n^2) for Sinkhorn.
  • Minibatch Optimal Transport: Minibatch OT computes OT solutions on data subsets and averages them as a scalable proxy for the full problem.Its minimizers converge to the true minimizer as minibatch size increases, according to prior work.
  • Minibatch Optimal Transport: Minibatch sampling can create couplings between clusters that would not be coupled by exact OT, producing undesirable smoothing effects.Figure 1 illustrates this behavior in a two-dimensional toy example.
  • Proposed approach: The paper uses unbalanced OT with relaxed marginal constraints at minibatch level to obtain more robust transport plans.The formulation is intended to reduce undesired couplings and scale with minibatch size.
  • Assumptions: The analysis assumes source and target distributions have compact support, ensuring a bounded ground cost.The bound is denoted by a strictly positive constant M.

2. Related work and background

The paper situates unbalanced OT among scalable and robust transport formulations. It uses soft marginal penalties, generalized Sinkhorn computation, and related divergence-based approaches to address computational and robustness concerns.

  • Scope: The section introduces the use of UOT and related formulations for robust machine learning.It also discusses empirical estimation difficulties caused by the curse of dimensionality.
  • Unbalanced Optimal Transport: Unbalanced OT relaxes mass conservation, allowing transport plans to create or destroy mass through soft marginal penalties.The considered formulation uses Csiszár divergences and generalized Sinkhorn iterations.
  • Unbalanced Optimal Transport: For KL marginal penalties, the generalized Sinkhorn algorithm has complexity Õ(n^2/ϵ).Balanced OT is recovered as the marginal penalty parameter τ tends to infinity for equal-mass inputs.
  • Related formulations: Entropy regularization can lose metric properties because the regularized loss may be nonzero when comparing a distribution with itself.The paper motivates an unbalanced Sinkhorn divergence to mitigate this issue while preserving computational benefits.
  • Robustness: UOT is described as more robust to outliers because it does not require exact marginal matching.Other robustness approaches include partial OT, distributionally robust models, and marginal relaxations with divergence penalties.

3. Minibatch OT and robustness to sampling

Balanced minibatch OT inherits OT’s sensitivity to outliers and subsampling, which can force inappropriate matches. The paper analyzes this effect and shows that unbalanced minibatch OT limits outlier influence and reduces variance in a toy setting.

  • Sampling sensitivity: Balanced OT must transport outliers because of its marginal constraints, and minibatching can make ordinary samples appear as minibatch-level outliers.This can force matches that differ from those in the full transport plan.
  • Outlier analysis: As an outlier moves farther from the distributions’ supports, OT loss increases, whereas the UOT upper bound saturates.After saturation, the remaining UOT cost corresponds to removing the outlier with cost proportional to its mass.
  • Toy experiment: UMBOT has lower variance than MBOT for k = 30 and k = 500, while increasing k lowers the variance of subsample MBOT.Each computation is repeated five times.
  • Class imbalance: A two-dimensional example compares unbalanced and balanced minibatch OT when source and target class proportions differ.The source is balanced across classes, whereas the target contains 3 red and 7 green samples.

4. Unbalanced Minibatch Optimal Transport

The paper defines unbalanced minibatch OT estimators by averaging UOT over sampled minibatches, then establishes their statistical and optimization properties. These estimators are unbiased, admit concentration guarantees and generalized-gradient results, while enabling scalable computation.

  • Estimator construction: Unbalanced minibatch OT replaces full-distribution UOT with an expectation over UOT problems computed on sampled minibatches.The formulation uses minibatches of size m and can be approximated by subsampling k pairs of minibatches.
  • Estimator construction: The complete and incomplete minibatch UOT losses are unbiased estimators of the expected minibatch UOT objective.They are respectively complete and incomplete U-statistics, although the losses are not definite because self-comparisons can remain positive.
  • Statistical properties: Under bounded-cost assumptions, UOT is finite and its set of optimal transport plans is compact, supporting the concentration analysis.The ground cost is bounded when the source and target distributions have compact support.
  • Statistical properties: Increasing the data size n and batch count k brings the subsampled estimator closer to its expectation, with a convergence rate in n that is nearly optimal.The deviation bound uses an upper bound M for UOT and, unlike original unbalanced OT, does not depend on the data dimension.
  • Optimization properties: The expected minibatch UOT objective permits exchanging expectation and generalized differentiation, including unregularized UOT where optimal plans need not be unique.The negative expected objective is Clarke regular, providing a nonsmooth optimization framework.
  • Optimization properties: With decreasing step sizes, stochastic gradient descent converges almost surely to critical points of the Clarke generalized derivative for minibatch UOT losses.The paper uses this result to justify SGD optimization of the parametric objective.

5. Experiments

The experiments evaluate JUMBOT’s unbalanced minibatch OT in gradient-flow, domain-adaptation, and partial-domain-adaptation settings. Across these settings, the method is reported as more robust to imbalanced classes, minibatch effects, and incorrect couplings than balanced minibatch OT baselines.

  • Unbalanced minibatch OT gradient flow: In gradient-flow experiments, minibatch OT is not robust to imbalanced classes, whereas minibatch UOT better respects target-distribution shape.The experiments use imbalanced 2D clusters and compare minibatch sizes of 64 and 128.
  • Method: JUMBOT combines a neural embedding-space cost with a label-space loss to learn a joint source–target distribution map for domain adaptation.Its transportation cost uses representation distance and cross-entropy, with an additional source-domain cross-entropy term.
  • Domain adaptation: JUMBOT achieved the best digit-dataset performance with a margin of at least 1.5 points and a 4% improvement over DEEPJDOT.Digit experiments covered USPS-to-MNIST, MNIST-to-MNIST-M, and SVHN-to-MNIST, with three runs per experiment.
  • Domain adaptation: JUMBOT achieved the best average accuracy on Office-Home and the best accuracy on 11 of 12 scenarios, while also leading VisDA by 4% and 2%, respectively.The reported results use ResNet-50 and benchmark comparisons against DANN, CDAN-E, DEEPJDOT, ALDA, and ROT.
  • Domain adaptation: 7% of DEEPJDOT connections linked samples with different labels on average, compared with 0.7% for JUMBOT during USPS-to-MNIST training.The analysis associates these incorrect connections with transferring wrong labels to the target and reduced overall accuracy.
  • Partial domain adaptation: On partial Office-Home, JUMBOT was state of the art on 9 of 12 tasks, averaged 2% above competitors, and was 15% higher than DEEPJDOT.The partial-domain setting uses target labels as a subset of source labels and evaluates 25 selected target categories.

6. Conclusion

The paper argues that minibatch sampling amplifies balanced OT’s undesirable coupling effects and proposes minibatch UOT to mitigate them. It reports preserved theoretical properties, dampened negative couplings, and state-of-the-art performance on challenging domain-adaptation problems.

  • Conclusion: Minibatch sampling amplifies balanced OT’s sensitivity to marginal constraints and subsampling, which can harm learning performance.The paper motivates relaxing marginal constraints at the minibatch level to address these effects.
  • Conclusion: Minibatch UOT dampens negative coupling effects and provides a more efficient distribution-comparison measure while preserving the paper’s theoretical properties.The conclusion specifically highlights its use as a scalable alternative for machine-learning applications.
  • Conclusion: The method reaches state-of-the-art performance on challenging domain-adaptation problems.The conclusion presents this as evidence supporting minibatch unbalanced OT in machine-learning applications.

Domain Adaptation

The paper formulates minibatch unbalanced optimal transport and develops its basic properties. It emphasizes relaxed marginal constraints, averaging over minibatches, and limitations of the resulting transport plans and losses.

  • Minibatch UOT formulation: Minibatch UOT applies unbalanced optimal transport with relaxed marginal constraints to minibatch pairs.The formulation computes UOT at the minibatch level rather than between full distributions.
  • Minibatch UOT formulation: The averaged minibatch transport matrix aggregates optimal plans over all possible minibatch couples using lifted plans.Each minibatch plan is embedded into the full sample space before averaging.
  • Limitations: The averaged minibatch UOT plan generally is not an OT plan because its relaxed marginals need not match the original marginals.For Sinkhorn divergence, three terms prevent defining an associated averaged minibatch transport matrix.
  • Basic properties: The minibatch UOT losses are positive and symmetric but not definite: for nontrivial X and 1 < m < n, ¯h_m(X, X) > 0.The nonzero self-loss follows because some minibatch pairs contain different samples.
  • Averaging construction: Averaging uniform minibatch weights recovers the original uniform weights through a combinatorial counting argument.The proof counts how often each index appears across m-tuples without repeated elements.

B.1.1. ROBUSTNESS

This section contrasts UOT and OT under an outlier perturbation. It presents UOT as robust to an outlier, whereas the corresponding OT result is not robust.

  • Robustness comparison: Lemma 1 separates robustness into a UOT result showing tolerance to an outlier and an OT result showing lack of robustness.The comparison uses a measure perturbed by adding a Dirac outlier outside the other measure’s support.
  • Outlier model: The perturbation is modeled as ˜μ = ζμ + (1 − ζ)δ_z, with z outside ν’s support.The measures are assumed to have compact support.
  • UOT robustness proof: The UOT proof constructs a feasible perturbed plan using optimal potentials and optimizes a mass parameter κ.The argument uses KL joint convexity and then minimizes the resulting upper bound with respect to κ.

B.1.2. UOT PROPERTIES

The paper establishes finiteness, boundedness, compactness, convexity, and Lipschitz properties for UOT under stated cost and mass assumptions. These properties support later statistical and optimization analysis.

  • Finiteness and bounds: Under finite ⟨abᵀ, C⟩, the UOT cost is finite and admits an explicit upper bound.For uniformly bounded costs, the proof uses the product plan abᵀ as a feasible competitor.
  • Plan compactness: Optimal UOT plans lie in a compact set when the ground cost is locally bounded, and the bound can be global for 0 ≤ C ≤ M.The compactness argument bounds the total plan mass and extracts convergent subsequences.
  • Plan geometry: The set of optimal transport plans is convex for regularized and unregularized UOT.This follows from convexity of the optimization objective and feasible set.
  • Optimization properties: The UOT loss is locally Lipschitz in the ground cost and globally Lipschitz when costs are uniformly bounded.The proof combines optimality inequalities with uniform bounds on optimal plan norms.
  • Assumptions: The analysis assumes compactly supported distributions, a positive symmetric definite C1 ground cost, and finite strictly positive input masses.Squared Euclidean cost is used without loss of generality in the stated setup.

B.2.1. PROOF OF THEOREM 1

The proof develops concentration results for complete and incomplete minibatch estimators and a distance-to-marginals result. It relies on bounded UOT losses and plans together with Hoeffding-type arguments.

  • Estimator concentration: The incomplete estimator is analyzed through independent minibatch draws, yielding a deviation bound controlled by the number of batches.Conditioned on the data, the sampled variables are independent, centered, and bounded.
  • Estimator concentration: The complete minibatch UOT estimator is a two-sample U-statistic whose concentration around its expectation follows from boundedness and Hoeffding’s inequality.Compact support and the UOT upper bound ensure the loss is bounded.
  • Theorem 1.1: Theorem 1.1 combines deviations between the incomplete estimator, the complete estimator, and the expectation over minibatches.The resulting maximal deviation depends on empirical sample size and batch count.
  • Theorem 1.2: Theorem 1.2 gives a high-probability bound on the distance between the averaged minibatch UOT plan and the target marginals.The proof uses bounded minibatch plans and Hoeffding’s inequality over sampled minibatch pairs.

B.2.2. PROOF OF THEOREM 2

The proof establishes regularity properties for unbalanced optimal transport costs and justifies exchanging minibatch expectations with generalized gradients under compact-support and integrability assumptions.

  • Regularity and generalized gradients: Clarke differential handles unregularized UOT because the optimal transport plan need not be unique.When ε > 0, the optimal plan is unique, so the convex-hull notation in the generalized gradient is unnecessary.
  • Regularity and generalized gradients: For a C1 parametrization θ ↦ Yθ and C1 cost C, the negative UOT cost is Clarke regular.The proof uses Clarke regularity of the cost composition and a chain-rule argument.
  • Regularity and generalized gradients: The UOT subderivative is the closed convex hull of derivatives induced by optimal transport plans.Danskin’s theorem identifies the subderivatives with optimal plans, while the chain rule propagates them through the parametrized cost.
  • Theorem 2 assumptions: Compact support and a locally integrable bound permit the theorem’s expectation and generalized-gradient operations to be controlled.The assumptions include compact support for X and {Yθ}, an open neighborhood around each parameter, and finite expected bounds.
  • Theorem 2 conclusion: Under these assumptions, the Clarke subdifferential of the expected UOT cost equals the expectation of the pointwise Clarke subdifferential.The resulting expected negative UOT objective is also Clarke regular.

C. Domain adaptation and partial domain adaptation experiments

The domain-adaptation experiments evaluate JUMBOT under stratified minibatch sampling against established baselines across digit, Office-Home, VisDA, and partial-adaptation settings.

  • Experimental setup: JUMBOT uses stratified source minibatches, assigning each class the same number of samples because source labels are available.This sampling procedure follows the setup used for DEEPJDOT.
  • Experimental setup: Digit experiments use a 9-layer CNN, one dense classification layer, 10 source-pretraining epochs, Adam, learning rate 2e−4, and minibatches of 500.Competitors use official implementations with the considered architecture and training procedure.
  • Experimental setup: Office-Home and VisDA experiments use an ImageNet-pretrained ResNet-50 generator and compare JUMBOT with DANN, CDAN-E, ALDA, DEEPJDOT, and ROT.The classifier and discriminator are trained from scratch with learning rates 10 times the generator’s rate.
  • Partial domain adaptation: Partial domain adaptation uses a related architecture and training procedure, with 5000 iterations, batch size 65, and hyperparameters τ = 0.06, η1 = 0.003, η2 = 0.75, and η3 = 10.The setup follows the setting associated with Jian et al. (2020).
  • Partial domain adaptation: JUMBOT’s reported partial-adaptation scores exceed ENT and PADA by at least 5% on average.The authors note uncertainty about whether the compared scores were evaluated at the end of optimization or during training.

C.3. Sensitivity analysis

Sensitivity analyses examine hyperparameters, minibatch size, and training behavior for JUMBOT and DEEPJDOT, showing greater stability for JUMBOT under several perturbations.

  • Sensitivity-study design: Figure 6 reports classification accuracy across τ, ε, and batch-size variations on USPS 7→MNIST and SVHN 7→MNIST.The analysis compares DEEPJDOT and JUMBOT under these hyperparameter changes.
  • Marginal penalization: JUMBOT accuracy increases with τ and peaks around τ = 1, while excessive τ causes a slight decrease.The decrease is attributed to respecting marginal distributions together with OT constraints and minibatch sampling.
  • Marginal penalization: When τ is too small, entropic regularization produces negative transfer in JUMBOT.This identifies a low-τ regime where the marginal penalty does not support the best accuracy.
  • Entropic regularization: Entropy slightly improves JUMBOT and entropic DEEPJDOT results, but excessive entropic regularization reduces accuracy.The authors conjecture that smoothing source labels regularizes the neural network and affects penultimate-layer class clusters.
  • Batch-size robustness: JUMBOT maintains constant accuracy across small batch sizes, whereas DEEPJDOT loses 4% on SVHN 7→MNIST and 6% on USPS 7→MNIST.The comparison supports JUMBOT’s robustness to small minibatches and limited computation budgets.
  • Training behavior: On MNIST 7→M-MNIST, DEEPJDOT begins overfitting after epoch 30 across classes, while JUMBOT remains stable and achieves higher accuracy.The magnitude of DEEPJDOT’s class-accuracy reduction varies across classes.
Loading 2103.03606v1…