Sample complexity of unbalanced entropic OT

Optimal transport (OT) has evolved as a powerful metric for distribution comparison, with the Wasserstein distance being a prominent example. Early work focused on balanced OT, enforcing strict marginal constraints, which limited scalability and robustness in real-world noisy data. The advent of entropic regularization, notably Sinkhorn algorithms, revolutionized computational efficiency, enabling large-scale applications in image processing, generative modeling, and domain adaptation. However, classical OT suffers from the curse of dimensionality, with sample complexity deteriorating exponentially as dimension grows. To address data with missing or created mass, unbalanced OT (UOT) formulations emerged, relaxing marginal constraints via divergence penalties like KL or χ². Despite computational advances, theoretical understanding of the statistical properties, especially at the plan level, remained limited. Existing results mainly focused on scalar costs or objective value stability, leaving a gap in understanding the sample complexity of the actual coupling in high-dimensional, noisy settings.

The core challenge is to establish finite-sample guarantees for the optimal coupling in unbalanced entropic OT, particularly in high dimensions where data noise and mass variation complicate geometric stability. Traditional methods lack geometric invariance and do not provide plan-level bounds, which are essential for applications like stochastic process inference and generative modeling. The absence of a unified geometric framework hampers understanding of how regularization influences sample efficiency and stability. Consequently, practitioners face uncertainty about the number of samples needed for reliable estimation, limiting the method’s practical deployment in large-scale problems.

This work introduces a translation-invariant dual envelope that restores geometric symmetry lost in unbalanced formulations, enabling a unified analysis of stability and statistical guarantees. Key innovations include:

• �� Establishing compactness and strong convexity of the dual variables within a carefully constructed anchored set, independent of divergence specifics;
• �� Deriving explicit bounds on the translation parameter, ensuring uniform control over dual potentials;
• �� Translating geometric estimates into high-probability bounds on the empirical transport plan, valid across broad φ-divergences.

These contributions collectively provide a rigorous foundation for understanding the sample complexity of unbalanced entropic OT, bridging geometric, statistical, and computational perspectives.

• �� Define a translation-invariant envelope Tα,β by minimizing the dual objective over scalar shifts, removing gauge freedom.
• �� Prove that Tα,β is strongly convex on a compact anchored set, ensuring stability of dual potentials.
• �� Impose mild Assumptions A1 and A2 to bound the dual potentials and restrict the translation parameter within a finite interval.
• �� Use geometric estimates and concentration inequalities to relate empirical and population dual solutions.
• �� Derive explicit high-probability bounds for the deviation of the empirical transport plan from the true plan, based on sample size n, divergence parameters, and geometric constants.
• �� Implement Sinkhorn-type algorithms with normalization strategies exploiting geometric symmetry for scalable computation.

Synthetic high-dimensional datasets (up to 500 dimensions) were generated with controlled noise levels to validate theoretical bounds. The experiments compared regularized unbalanced OT against classical Wasserstein, measuring sample complexity and estimation error. Hyperparameters such as regularization strength η and divergence parameters were tuned to optimize stability. Ablation studies examined the impact of geometric normalization and divergence choice. Results consistently showed that the proposed geometric framework reduced sample requirements by over 50%, with errors decreasing by 20-30%, confirming the theoretical predictions. Additional tests across different divergence types demonstrated the universality of the geometric properties.

The experiments confirmed that, under mild assumptions, the sample complexity for estimating the optimal plan scales as O(1/√n), with explicit constants derived from geometric bounds. In high dimensions, this translates to a significant reduction in data needed compared to unregularized OT. The geometric analysis ensures robustness against noise, with the regularized method maintaining stability even when data is highly corrupted. Across divergence types, the bounds hold uniformly, validating the broad applicability of the theoretical framework. These results demonstrate that entropy regularization not only accelerates computation but also enhances statistical efficiency in complex, high-dimensional settings.

The framework applies directly to large-scale distribution matching tasks such as generative modeling, domain adaptation, and stochastic process inference, especially in high-dimensional data environments like images, text, and sensor data. It provides practitioners with quantifiable data requirements, enabling efficient data collection and model training. The robustness and scalability make it suitable for industry applications requiring reliable distribution alignment under noise and mass variation, including healthcare, finance, and autonomous systems.

The analysis assumes compactness and Lipschitz continuity of the cost function, which may not hold in all real-world scenarios. Computational costs, while improved, remain significant for extremely large datasets, necessitating further approximation techniques. The bounds depend on divergence-specific assumptions, potentially limiting effectiveness under extreme noise or non-standard divergences. Extending the theory to non-compact or non-smooth settings remains an open challenge, requiring additional geometric and statistical tools.