Research map / Optimization / Optimizers
On a Neural Implementation of Brenier's Polar Factorization
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes a neural implementation of Brenier's polar factorization using ICNNs with low-rank quadratic layers, an amortized measure-preserving map, and bridge-matching-based inverse map, applied to topography gradients, MNIST loss landscapes, and LMC sampling.
Problem
Brenier's polar factorization decomposes any vector field into a gradient of a convex potential composed with a measure-preserving map, but no practical high-dimensional implementation existed. The paper targets learning this factorization from samples and using its inverse to study non-convex lands
Approach
Parameterize the convex potential as an ICNN with added positive-definite (diagonal plus low-rank) quadratic layers, estimate the measure-preserving map either implicitly via the convex conjugate or with an auxiliary network, and learn a stochastic inverse map via augmented bridge matching to sample pre-images; combine this with Langevin dynamics to sample non-convex landscapes.
Key results
- On the ICNN benchmark with a Gaussian target in d=32, the proposed architecture reaches 0.031 L2 unexplained variance versus 0.38 for FQuad and 0.72 for Linear.
- For the Chamonix elevation gradient, the Sinkhorn divergence between the learned pushforward and the target (0.36) is of the same order as between two target batches (0.31).
- On the MNIST classifier loss gradient, the cosine similarity between the target gradient and the gradient at the inverse-map samples averages 0.81 ± 0.20.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use low-rank + diagonal quadratic ICNN layers instead of Linear ICNN layers | ICNN OT map estimator · Gaussian OT, d=32 and d=128 | L2 unexplained variance percentage: 0.72 \color{gray}\pm 0.05 -> 0.031 \color{gray}\pm 0.003; 1.26 \color{gray}\pm 0.08 -> 0.082 \color{gray}\pm 0.002 | Helped |
| Use low-rank + diagonal quadratic ICNN layers instead of FQuad layers | ICNN OT map estimator · Gaussian OT, d=32 and d=128 | L2 unexplained variance percentage: 0.38 \color{gray}\pm 0.03 -> 0.031 \color{gray}\pm 0.003; 0.97 \color{gray}\pm 0.06 -> 0.082 \color{gray}\pm 0.002 | Helped |
| Use low-rank + diagonal quadratic ICNN layers instead of Linear ICNN layers | ICNN OT map estimator · Gaussian mixture OT, d=32 and d=128 | Sinkhorn divergence: 548 \color{gray}\pm 75 -> \mathbf{234 \color{gray}\pm 42}; 246 \times 1e^1 \color{gray}\pm 42\times 1e^1 -> \mathbf{104 | Helped |
| Use low-rank + diagonal quadratic ICNN layers instead of FQuad layers | ICNN OT map estimator · Gaussian mixture OT, d=32 and d=128 | Sinkhorn divergence: 452 \color{gray}\pm 77 -> \mathbf{234 \color{gray}\pm 42}; 176 \times 1e^1 \color{gray}\pm 14\times 1e^1 -> \mathbf{104 | Helped |
| Use NPF to estimate the gradient transport | NPF · Chamonix elevation gradients | Sinkhorn divergence, predicted vs target gradient distribution: 0.36 vs 0.31 | No clear effect |
| Use NPF's learned measure-preserving map Mξ | NPF · Chamonix elevation gradients | Sinkhorn divergence, mapped source vs source: 0.0022 vs 0.0015 | No clear effect |
| Use NPF reconstruction ∇uθ ∘ Mξ | NPF · Chamonix elevation gradients | Reconstruction L2 distance: 0.96 | Helped |
| Use the learned inverse map Iψ | NPF · Chamonix elevation gradients | Sinkhorn divergence, inverse-map samples vs approximate pre-image: 0.046 vs 0.039 | No clear effect |
| Use NPF to estimate the gradient transport | NPF · MNIST classifier loss gradients | Sinkhorn divergence, predicted vs target gradient distribution: 97.4 \color{gray}\pm 6.5 vs 93.6 \color{gray}\pm 6.2 | No clear effect |
| Use the learned inverse map Iψ | NPF · MNIST classifier loss gradients | Sinkhorn divergence, inverse-map samples vs approximate pre-image: 107.3 \color{gray}\pm 0.7 vs 105.7 \color{gray}\pm 0.5 | No clear effect |
| Use the learned inverse map Iψ to reconstruct classifier weights | NPF · MNIST classifier loss gradients | Cosine similarity: 0.81 \color{gray}\pm 0.20 | Helped |
| Train NPF using gradient descent trajectories rather than uniform samples over the whole space | MNIST classifier critical-point sampling · MNIST classifier | Classifier accuracy of generated critical points | Helped |
Surprising findings
- Smooth activations (ELU) were needed in the potential parameterization because ReLU caused convergence problems for conjugate solvers in high dimension.
- Samples generated by the inverse map from zero-gradient targets have low classifier accuracy (around 10%, random-weight level), so uniform sampling mostly reveals bad critical points.
Limitations
- Learned polar factorization is less accurate in parts of the MNIST weight space, as the loss landscape is chaotic.
- Sampling weights are not reweighted by their probability, so the inverse generative model does not produce exact samples from the target density.
- Computing the convex conjugate requires an optimization routine at every evaluation of M.
Experimental setup
Datasets: MNIST Elevation data via the elevation Python package (Chamonix, London, Cyprus) Synthetic Gaussian and Gaussian-mixture distributions
Benchmarks: MNIST Chamonix elevation gradients London elevation gradients Cyprus elevation gradients
Metrics: Sinkhorn divergence L2 unexplained variance percentage L2 distance between G and grad u o M cosine similarity classifier accuracy
Baselines: Linear ICNN (Amos et al.) FQuad ICNN
Builds on: Brenier polar factorization theorem input convex neural networks augmented bridge matching Makkuva et al. neural OT solver
Closest papers by meaning
- Low-Rank Optimal Transport through Factor Relaxation with Latent Coupling (NeurIPS 2024 · similarity 0.84)
- Provable Acceleration of Nesterov's Accelerated Gradient for Asymmetric Matrix Factorization and Linear Neural Networks (NeurIPS 2024 · similarity 0.83)
- Minimum Norm Interpolation Meets The Local Theory of Banach Spaces (ICML 2024 · similarity 0.83)
- PoLAR: Polar-Decomposed Low-Rank Adapter Representation (NeurIPS 2025 · similarity 0.83)
- How to Make the Gradients Small Privately: Improved Rates for Differentially Private Non-Convex Optimization (ICML 2024 · similarity 0.83)
- Sparsistency for inverse optimal transport (ICLR 2024 · similarity 0.83)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.