# A lift for input-convex neural net training

arXiv: https://arxiv.org/abs/2605.24274 · HTML: https://mlsanity.com/paper/2605.24274/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 33 of its 50 references that are in the map.
- First posted on arXiv 22 May 2026.

## Builds on (most cited first, 30 of 33)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [A Style-Based Generator Architecture for Generative Adversarial Networks](https://mlsanity.com/paper/1812.04948/) (2018)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Visualizing the Loss Landscape of Neural Nets](https://mlsanity.com/paper/1712.09913/) (2017)
- [HyperNetworks](https://mlsanity.com/paper/1609.09106/) (2016)
- [Masked Autoregressive Flow for Density Estimation](https://mlsanity.com/paper/1705.07057/) (2017)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Optimal transport mapping via input convex neural networks](https://mlsanity.com/paper/1908.10962/) (2019)
- [Adding Gradient Noise Improves Learning for Very Deep Networks](https://mlsanity.com/paper/1511.06807/) (2015)
- [Stochastic modified equations and adaptive stochastic gradient algorithms](https://mlsanity.com/paper/1511.06251/) (2015)
- [Adversarial Regularizers in Inverse Problems](https://mlsanity.com/paper/1805.11572/) (2018)
- [Do Neural Optimal Transport Solvers Work? A Continuous Wasserstein-2 Benchmark](https://mlsanity.com/paper/2106.01954/) (2021)
- [Convex Potential Flows: Universal Probability Distributions with Optimal Transport and Convex Optimization](https://mlsanity.com/paper/2012.05942/) (2020)
- [Supervised Training of Conditional Monge Maps](https://mlsanity.com/paper/2206.14262/) (2022)
- [On the representation and learning of monotone triangular transport maps](https://mlsanity.com/paper/2009.10303/) (2020)
- [Asymptotic Optimality in Stochastic Optimization](https://mlsanity.com/paper/1612.05612/) (2016)
- [Conditional score-based diffusion models for Bayesian inference in infinite dimensions](https://mlsanity.com/paper/2305.19147/) (2023)
- [Parameterized Machine Learning for High-Energy Physics](https://mlsanity.com/paper/1601.07913/) (2016)
- [HINT: Hierarchical Invertible Neural Transport for Density Estimation and Bayesian Inference](https://mlsanity.com/paper/1905.10687/) (2019)
- [On amortizing convex conjugates for optimal transport](https://mlsanity.com/paper/2210.12153/) (2022)
- [Incremental Learning in Diagonal Linear Networks](https://mlsanity.com/paper/2208.14673/) (2022)
- [Efficient Neural Network Approaches for Conditional Optimal Transport with Applications in Bayesian Inference](https://mlsanity.com/paper/2310.16975/) (2023)
- [Multivariate Quantile Function Forecaster](https://mlsanity.com/paper/2202.11316/) (2022)
- [Principled Weight Initialisation for Input-Convex Neural Networks](https://mlsanity.com/paper/2312.12474/) (2023)
- [An Optimal Transport Formulation of Bayes' Law for Nonlinear Filtering Algorithms](https://mlsanity.com/paper/2203.11869/) (2022)
- [Invertible Monotone Operators for Normalizing Flows](https://mlsanity.com/paper/2210.08176/) (2022)
- [Deep unfolding as iterative regularization for imaging inverse problems](https://mlsanity.com/paper/2211.13452/) (2022)
- [Improving Fairness and Mitigating MADness in Generative Models](https://mlsanity.com/paper/2405.13977/) (2024)
- [A primal-dual algorithm for image reconstruction with input-convex neural network regularizers](https://mlsanity.com/paper/2410.12441/) (2024)
- [Sample and Map from a Single Convex Potential: Generation using Conjugate Moment Measures](https://mlsanity.com/paper/2503.10576/) (2025)

## Abstract (arXiv)

Input-convex neural nets parametrize the convex potentials of density models and transport maps, and their convexity requires the inter-layer weights to be non-negative. Projected gradient descent enforces this by projecting after each step, and due to mini-batch noise the boundary is re-crossed indefinitely, which leads to an active set the projection never identifies. The differentiable alternative, direct softplus, optimizes a free latent weight through a softplus positivity map whose derivative attenuates the gradient exponentially where the weight is negative---the shoulder---so a coordinate that reaches it stays for an exponentially long time. To keep this unconstrained parametrization without its slow escape, we propose the lift, which replaces the free latent weight by a learnable slack plus an unconstrained network---the body---that takes a permutation-invariant summary of the training batch as input. The latent weight thus varies with the batch before the positivity map, and couples to the gradient formed on it. We show that this coupling enters the variance of the update to the latent weight at first order in the fluctuation, and that the slack, the batch dependence and the shared batch are each needed for it to act. Where the coupling aligns positively with the loss curvature, that variance is larger under the lift than under direct softplus, and a coordinate leaves the shoulder sooner. We compare the lift with the two existing methods on several applications. Where a constrained weight of direct softplus reaches the shoulder and does not leave, the lift fits the target more closely and reaches the same reconstruction about three times sooner. Where almost none reaches it, the methods agree.