Research map / Learning theory and algorithms / Optimal transport
Learning Elastic Costs to Shape Monge Displacements
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 proximal-gradient method to compute provably optimal Monge maps for elastic (translation-invariant) costs, and a bilevel loss to learn a subspace-elastic cost parameter from samples.
Problem
Monge optimal transport with structured costs: generating ground-truth OT maps for elastic costs and learning the regularizer parameter (e.g., a low-dimensional displacement subspace) from source and target samples.
Approach
For any elastic cost h = l2^2/2 + gamma*tau, compute h-transforms via proximal gradient iterations using only prox_tau, yielding provably optimal ground-truth maps; introduce subspace costs tau(z)=||A^perp z||^2 with Stiefel-constrained A, learned by a bilevel loss <P*(theta), R(theta)> optimized via Riemannian gradient descent.
Key results
- Subspace recovery error often below one percent for ground-truth A* after 1000 Riemannian gradient steps on synthetic tasks.
- On synthetic data, the MBO estimator with positive gamma significantly outperforms the l2^2 cost (gamma=0) for any range of gamma.
- On single-cell data (d=256), the learned elastic subspace cost improves predictive performance over the l2^2 baseline, measured even in squared-Euclidean distance.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use a positive-γ elastic cost instead of the unregularized ℓ2² cost | MBO estimator · Synthetic ground-truth displacement tasks (ℓ1 and subspace costs, d=5 and d=10) | MSE | Helped |
| Use an elastic subspace cost instead of the ℓ2² cost | MBO estimator · Single-cell RNA-sequencing perturbation tasks | Sinkhorn divergence | Helped |
| Increase the share of displacement inertia captured by the p* subspace (ablation) | Subspace recovery pipeline · Synthetic ground-truth subspace recovery | Normalized squared-residual error | Helped |
| Set estimated subspace dimension p̂ to 1.25 p* rather than p* (ablation) | Subspace recovery pipeline · Synthetic ground-truth subspace recovery | Normalized squared-residual error | Helped |
| Vary γ for the ℓ1 regularizer (ablation) | MBO estimator · Synthetic ground-truth displacement task with ℓ1 cost | MSE | Helped |
| Vary γ for the subspace regularizer (ablation) | MBO estimator · Synthetic ground-truth displacement task with subspace cost | MSE | Helped |
Surprising findings
- Promoting displacements in a much lower-dimensional subspace improves predictions even when evaluated with the squared-Euclidean distance.
Limitations
- Sample-complexity result for subspace costs is stated only for fixed gamma and known A; large-gamma rate is left as future work
- Parameterization restricted to Stiefel-manifold A; alternative parameterizations left for future work
- Most underperforming recovery experiments could be improved with early stopping
Experimental setup
Datasets: synthetic Gaussian point clouds single-cell RNA sequencing drug perturbation data (188 drugs, 5 selected)
Benchmarks: synthetic ground-truth Monge maps single-cell drug perturbation transport
Metrics: MSE ratio vs gamma normalized subspace residual error Sinkhorn divergence
Baselines: squared-Euclidean cost MBO estimator (gamma=0) l2^2 Sinkhorn baseline
Builds on: Sinkhorn algorithm MBO estimator Brenier theorem OTT-JAX
Closest papers by meaning
- Monge, Bregman and Occam: Interpretable Optimal Transport in High-Dimensions with Feature-Sparse Maps (ICML 2023 · similarity 0.88)
- The Monge Gap: A Regularizer to Learn All Transport Maps (ICML 2023 · similarity 0.87)
- Optimization Can Learn Johnson Lindenstrauss Embeddings (NeurIPS 2024 · similarity 0.84)
- Sparsistency for inverse optimal transport (ICLR 2024 · similarity 0.83)
- High-Dimensional Calibration from Swap Regret (NeurIPS 2025 · similarity 0.82)
- Improving Neural Optimal Transport via Displacement Interpolation (ICLR 2025 · similarity 0.82)
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.