W-LALM: A Linearized Augmented Lagrangian Method for Constrained Optimization in Wasserstein Space

Many problems in inference, sampling, and scientific calibration require learning a probability law subject to linear expectation constraints. The natural extension of Wasserstein proximal gradient (WPG) to this problem enforces the constraints inside its Jordan–Kinderlehrer–Otto (JKO) subproblem, which can turn an otherwise tractable backward step into a constrained transport problem. Primal–dual sampling methods avoid this difficulty by moving feasibility into a multiplier update, but their guarantees cover averaged or distributional quantities rather than the full primal–dual sequence, and they do not exploit composite objectives whose nonsmooth part admits a tractable JKO resolvent. We propose the Wasserstein linearized augmented Lagrangian method (W-LALM), which keeps the original unconstrained JKO step and handles feasibility through an explicit step on the smooth part of the augmented Lagrangian together with a multiplier update. Our analysis is built on transport plans rather than transport maps, since optimal Monge maps need not exist for the discrete measures used in particle implementations. This representation yields a plan-based Lyapunov descent estimate, convergence of the full primal–dual sequence to a saddle point under generalized-geodesic convexity, and an $O(1/\epsilon )$ iteration complexity for the squared residual of a plan-valued Wasserstein Karush–Kuhn–Tucker (KKT) condition that we introduce. W-LALM matches or improves upon the terminal objective of the strongest feasible baseline, while being 8–38$\times$ faster than the constrained-JKO variant of WPG and up to two orders of magnitude faster than the Euclidean augmented-Lagrangian baseline.