ITSPACE: Monotone Gaussian Optimal Transport Updates
Woojoo Na, Jennifer Dy
Abstract
Covariance matrices serve as compact descriptors of feature distributions in many machine-learning pipelines, including domain adaptation and Gaussian embeddings. Under a centered Gaussian approximation, the unregularized Wasserstein-2 optimal transport (OT) discrepancy admits a closed form on covariances given by the Bures--Wasserstein (BW) objective on the symmetric positive definite (SPD) cone. We propose ITSPACE (Iterative Transport for Stable Proximal Alignment of Covariance Embeddings), a proximal majorization-minimization method that directly optimizes this exact BW objective through closed-form updates in a square-root factorization. In exact arithmetic, each iteration satisfies a sufficient-decrease inequality for the BW objective; under inexact polar computations, we provide an explicit certificate-gap bound controlling deviations from exact descent. The resulting iterations preserve PSD structure by construction and naturally support rank-restricted factors, making ITSPACE well-suited as a lightweight inner-loop primitive in settings where adaptation must be performed from unlabeled target batches under strict step and compute budgets. Across real-world covariance-alignment benchmarks, ITSPACE reaches low-BW-gap solutions substantially faster than BW-gradient descent, methods based on other covariance geometries, and entropically regularized sample-OT baselines.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Builds on7
- WILDS: A Benchmark of in-the-Wild Distribution ShiftsPang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie et al.ICML 2021 · 1,773 citations
- Variational inference via Wasserstein gradient flowsMarc Lambert, Sinho Chewi, Francis R. Bach, Silvère Bonnabel et al.NeurIPS 2022 · 123 citations
- On Riemannian Optimization over Positive Definite Matrices with the Bures-Wasserstein GeometryAndi Han, Bamdev Mishra, Pratik Kumar Jawanpuria, Junbin GaoNeurIPS 2021 · 55 citations
- Forward-Backward Gaussian Variational Inference via JKO in the Bures-Wasserstein SpaceMichael Ziyang Diao, Krishna Balasubramanian, Sinho Chewi, Adil SalimICML 2023 · 47 citations
- On the Convergence of Projected Bures-Wasserstein Gradient Descent under Euclidean Strong ConvexityJunyi Fan, Yuxuan Han, Zijian Liu, Jian-Feng Cai et al.ICML 2024 · 2 citations
Related papers
- Critical Points and Convergence Analysis of Generative Deep Linear Networks Trained with Bures-Wasserstein LossPierre Bréchet, Katerina Papagiannouli, Jing An, Guido MontúfarICML 2023 · 7 citations
- The Unbalanced Gromov Wasserstein Distance: Conic Formulation and RelaxationThibault Séjourné, François-Xavier Vialard, Gabriel PeyréNeurIPS 2021 · 106 citations
- Entropic Optimal Transport between Unbalanced Gaussian Measures has a Closed FormHicham Janati, Boris Muzellec, Gabriel Peyré, Marco CuturiNeurIPS 2020 · 109 citations
- Linear-Time Gromov Wasserstein Distances using Low Rank Couplings and CostsMeyer Scetbon, Gabriel Peyré, Marco CuturiICML 2022 · 73 citations
- Joint Metric Space Embedding by Unbalanced Optimal Transport with Gromov-Wasserstein Marginal PenalizationFlorian Beier, Moritz Piening, Robert Beinert, Gabriele SteidlICML 2025
