🤖 AI Summary
Practical implementation of Brenier’s polar decomposition theorem in machine learning—decomposing an arbitrary vector field into the gradient of a convex potential and a measure-preserving map—is hindered by the latter’s frequent non-injectivity, rendering its inverse unreliable. Method: We propose the first differentiable, learnable neural decomposition framework: (1) parameterize the convex potential via input-convex neural networks (ICNNs), compute its gradient using explicit convex conjugation and conjugate gradient optimization; (2) adopt a dual-path architecture—auxiliary networks estimate the measure-preserving map while stochastic generators mitigate ill-posedness in inversion; (3) enforce measure preservation through neural optimal transport theory. Contribution/Results: Experiments demonstrate significantly improved convergence in non-convex optimization and high-fidelity sampling from non-log-concave densities, establishing, for the first time, the practical feasibility of Brenier decomposition in high-dimensional, non-convex settings.
📝 Abstract
In 1991, Brenier proved a theorem that generalizes the polar decomposition for square matrices -- factored as PSD $ imes$ unitary -- to any vector field $F:mathbb{R}^d
ightarrow mathbb{R}^d$. The theorem, known as the polar factorization theorem, states that any field $F$ can be recovered as the composition of the gradient of a convex function $u$ with a measure-preserving map $M$, namely $F=
abla u circ M$. We propose a practical implementation of this far-reaching theoretical result, and explore possible uses within machine learning. The theorem is closely related to optimal transport (OT) theory, and we borrow from recent advances in the field of neural optimal transport to parameterize the potential $u$ as an input convex neural network. The map $M$ can be either evaluated pointwise using $u^*$, the convex conjugate of $u$, through the identity $M=
abla u^* circ F$, or learned as an auxiliary network. Because $M$ is, in general, not injective, we consider the additional task of estimating the ill-posed inverse map that can approximate the pre-image measure $M^{-1}$ using a stochastic generator. We illustrate possible applications of Brenier's polar factorization to non-convex optimization problems, as well as sampling of densities that are not log-concave.