Geometry-Misalignment in Distributional Learning
Abstract
Distributional learning problems optimize discrepancies between probability measures, including optimal transport or Sinkhorn divergence, yet are typically optimized using Euclidean first-order methods in parameter space. We show this mismatch is structural rather than algorithmic. We introduce geometry-misalignment, a local condition number that measures distortion between Euclidean geometry and the intrinsic geometry induced by a distributional objective. For a broad class of problems, we establish lower bounds demonstrating that Euclidean first-order methods incur an unavoidable convergence slowdown proportional to misalignment, even under intrinsic strong convexity and smoothness. We further prove geometry-aware preconditioned methods attain matching upper bounds independent of misalignment, yielding a sharp separation between Euclidean optimization and geometry-aware optimization. Beyond convergence rates, we show geometry-misalignment induces an optimization-dependent excess risk term under finite budgets, directly linking optimization geometry with statistical efficiency. We develop a geometry-calibrated optimization framework that estimates misalignment and selectively activates geometry-aware updates when necessary. Experiments on distribution matching for domain adaptation validate the theory, with improvements concentrated in high-misalignment regimes and negligible overhead.
Lay Summary
Many machine learning methods compare whole data distributions rather than individual prediction errors. These methods are useful in tasks such as matching data from different sources, adapting models to new domains, and making models more robust. However, they can be difficult to train with standard optimization tools. This paper explains one important reason, i.e., the directions that look equally important in the model parameters may have very different effects on the distributions being compared. We introduce a way to measure this mismatch. When the mismatch is small, standard training methods work well. When it is large, they can become slow or unstable. We prove that this slowdown is unavoidable for standard methods in such cases, and we show that geometry-aware updates can remove the slowdown. We then propose an adaptive training method that uses standard updates when they are sufficient and switches to geometry-aware updates only when needed. Experiments show that the method improves training most in the difficult high-mismatch cases while avoiding unnecessary extra computation in easier cases.