Sample complexity of unbalanced entropic OT

Technical Summary: Sample Complexity of Unbalanced Entropic OT

Problem Statement
Optimal Transport (OT) provides a geometric framework for comparing probability measures, but classical balanced OT faces two significant limitations in machine learning applications: (1) the rigidity of hard mass conservation constraints, which is unsuitable for data with missing or created mass, and (2) unfavorable high-dimensional sample complexity (the "curse of dimensionality") where empirical convergence rates deteriorate with ambient dimension.

While entropic regularization and unbalanced relaxations (using $\phi$-divergences) have been developed to address these issues individually, existing statistical guarantees for unbalanced entropic OT are primarily limited to the scalar transport value (the objective function). There is a lack of high-probability finite-sample bounds directly on the optimal coupling (transport plan). This gap is critical because, in applications like generative modeling or trajectory inference, the coupling itself is the object of interest for prediction and interpolation, not just the cost value.

Methodology
The paper develops a statistical analysis of unbalanced entropic OT by focusing on the geometric properties of the dual problem. The core methodological innovation is the introduction of a translation-invariant envelope formulation to handle the loss of gauge invariance inherent in unbalanced settings.

1. Dual Formulation and Envelope Construction:
   Unlike balanced OT, where the dual objective is invariant under additive translations of the Kantorovich potentials ($f \to f+\lambda, g \to g-\lambda$), unbalanced OT breaks this symmetry due to marginal divergence penalties depending on absolute potential levels. The authors define a translation-invariant envelope $T_{\alpha, \beta}(f, g) := \inf_{\lambda \in \mathbb{R}} K_{\alpha, \beta}(f+\lambda, g-\lambda)$. This construction explicitly separates the intrinsic degrees of freedom of the dual from the spurious translation direction.

2. Geometric Analysis (Compactness and Strong Convexity):
   The authors prove that under mild assumptions (compact spaces, Lipschitz cost), this envelope functional admits a compact anchored minimization domain ($P_{\alpha, \beta}^0$) where the potentials are fixed to have zero mean with respect to the source measure. Crucially, on this domain, the envelope is uniformly strongly convex with respect to a weighted $L^2$ norm. This recovers the favorable geometric structure lost in the standard unbalanced dual.

3. Control of Translation Parameter:
   While the envelope removes the need for arbitrary gauge fixing, it introduces an optimization over a translation parameter $\lambda$. The paper establishes that under a mild domain condition on the divergence generators (Assumption A2), the optimal $\lambda$ is confined to a uniformly bounded interval $[-R, R]$. This allows for uniform $L^\infty$ bounds on the dual potentials.

4. Statistical Stability:
   Leveraging the strong convexity of the anchored envelope and the boundedness of the translation parameter, the authors derive stability estimates. They show that the empirical dual potentials converge uniformly to the population potentials. These dual convergence rates are then transferred to the primal transport plan via the exponential scaling relationship ($d\pi/d(\alpha \otimes \beta) \propto \exp((f \oplus g - c)/\eta)$).

Key Contributions

• Translation-Invariant Dual Formulation: The paper introduces a rigorous dual formulation that restores strong convexity for unbalanced entropic OT by quotienting out the translation direction, a property previously utilized only for algorithmic stability in Sinkhorn iterations.
• Plan-Level Sample Complexity: The primary contribution is the derivation of high-probability finite-sample bounds for the optimal coupling (transport plan) in the unbalanced setting, rather than just the objective value.
• General Divergence Applicability: The results apply to a broad class of $\phi$-divergences (including KL, $\chi^2$, Hellinger, Jensen–Shannon, and $\alpha$-divergences) without requiring divergence-specific proofs, relying instead on general convexity and compactness arguments.
• Explicit Constants: The paper provides explicit constants for the bounds on the translation parameter and the resulting sample complexity, clarifying the dependence on the entropy parameter $\eta$, the cost function, and the divergence properties.

Main Results
The paper establishes that for empirical measures $\alpha_n, \beta_n$ derived from $n$ samples, the deviation between the empirical optimal plan $\pi_n$ and the population plan $\pi^\star$ satisfies a high-probability bound of the form:
$$|\langle h, \pi^\star - \pi_n \rangle| \lesssim C \cdot e^{C(M + \|c\|_\infty)/\eta} \|h\|_\infty \sqrt{\frac{m_\alpha m_\beta (\log n + t)}{n}}$$
where $M$ is a uniform bound on the dual potentials, $h$ is a test function, and the constants depend on the problem geometry and divergence parameters. Notably, if the sampling distribution matches the product of the marginals, the $\log n$ factor can be removed.

Significance and Claims
The paper claims that these results clarify the practical necessity of regularization in machine learning applications of OT. Specifically:

• Mitigating Dimensionality: Regularization softens the curse of dimensionality, reducing the number of samples required for stable transport estimation.
• Algorithmic Compatibility: The resulting estimators remain compatible with scalable Sinkhorn-type solvers, bridging the gap between statistical theory and computational practice.
• Stability for Inverse Problems: The forward stability estimates provided are essential for companion inverse-problem analyses (e.g., recovering costs from samples), as noted in related work [1].

The authors emphasize that their analysis relies on "very mild and largely divergence-agnostic assumptions," making the theoretical framework robust and applicable to the standard models used in practice. They do not claim to solve the curse of dimensionality entirely but demonstrate how entropic unbalanced OT provides a statistically tractable alternative to unregularized Wasserstein distances in high-dimensional settings.