ITSPACE: Monotone Gaussian Optimal Transport Updates
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.
Lay Summary
Many machine-learning systems adapt to new data by updating covariance matrices, which summarize how feature dimensions vary together. This paper introduces ITSPACE, a fast method for moving a low-rank covariance estimate toward a target covariance under the Bures-Wasserstein distance, the exact optimal-transport distance between Gaussian distributions. The method keeps every intermediate covariance valid and low-rank, and it has a closed-form update with theoretical descent guarantees. Experiments show that ITSPACE reaches accurate covariance alignment much faster than several alternatives and can improve downstream adaptation under controlled covariance shifts.