# Understanding the training of infinitely deep and wide ResNets with Conditional Optimal Transport

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 36 of its 69 references that are in the map.
- First posted on arXiv 19 Mar 2024.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Identity Mappings in Deep Residual Networks](https://mlsanity.com/paper/1603.05027/) (2016)
- [Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning](https://mlsanity.com/paper/1602.07261/) (2016)
- [A Convergence Theory for Deep Learning via Over-Parameterization](https://mlsanity.com/paper/1811.03962/) (2018)
- [Wide Neural Networks of Any Depth Evolve as Linear Models Under Gradient Descent](https://mlsanity.com/paper/1902.06720/) (2019)
- [Gradient Descent Finds Global Minima of Deep Neural Networks](https://mlsanity.com/paper/1811.03804/) (2018)
- [A Mean Field View of the Landscape of Two-Layers Neural Networks](https://mlsanity.com/paper/1804.06561/) (2018)
- [On the Global Convergence of Gradient Descent for Over-parameterized Models using Optimal Transport](https://mlsanity.com/paper/1805.09545/) (2018)
- [Fixup Initialization: Residual Learning Without Normalization](https://mlsanity.com/paper/1901.09321/) (2019)
- [On the Equivalence between Kernel Quadrature Rules and Random Feature Expansions](https://mlsanity.com/paper/1502.06800/) (2015)
- [Convergence Analysis of Two-layer Neural Networks with ReLU Activation](https://mlsanity.com/paper/1705.09886/) (2017)
- [On the linearity of large non-linear models: when and why the tangent kernel is constant](https://mlsanity.com/paper/2010.01092/) (2020)
- [How Much Over-parameterization Is Sufficient to Learn Deep ReLU Networks?](https://mlsanity.com/paper/1911.12360/) (2019)
- [A Rigorous Framework for the Mean Field Limit of Multilayer Neural Networks](https://mlsanity.com/paper/2001.11443/) (2020)
- [A Mean-Field Optimal Control Formulation of Deep Learning](https://mlsanity.com/paper/1807.01083/) (2018)
- [Representation formulas and pointwise properties for Barron functions](https://mlsanity.com/paper/2006.05982/) (2020)
- [Learning deep linear neural networks: Riemannian gradient flows and convergence to global minimizers](https://mlsanity.com/paper/1910.05505/) (2019)
- [A Mean-field Analysis of Deep ResNet and Beyond: Towards Provable Optimization Via Overparameterization From Depth](https://mlsanity.com/paper/2003.05508/) (2020)
- [Gradient descent with identity initialization efficiently learns positive definite linear transformations by deep residual networks](https://mlsanity.com/paper/1802.06093/) (2018)
- [On the Proof of Global Convergence of Gradient Descent for Deep ReLU Networks with Linear Widths](https://mlsanity.com/paper/2101.09612/) (2021)
- [Nonsmooth Implicit Differentiation for Machine Learning and Optimization](https://mlsanity.com/paper/2106.04350/) (2021)
- [On the Convergence of Gradient Descent Training for Two-layer ReLU-networks in the Mean Field Regime](https://mlsanity.com/paper/2005.13530/) (2020)
- [Momentum Residual Neural Networks](https://mlsanity.com/paper/2102.07870/) (2021)
- [On the Global Convergence of Training Deep Linear ResNets](https://mlsanity.com/paper/2003.01094/) (2020)
- [Analysis of a Two-Layer Neural Network via Displacement Convexity](https://mlsanity.com/paper/1901.01375/) (2019)
- [Conditional Optimal Transport on Function Spaces](https://mlsanity.com/paper/2311.05672/) (2023)
- [Overparameterization of deep ResNet: zero loss and mean-field analysis](https://mlsanity.com/paper/2105.14417/) (2021)
- [Convergence of gradient descent for deep neural networks](https://mlsanity.com/paper/2203.16462/) (2022)

## Built on by (most cited first, 8 of 8)

- [Conditional Wasserstein Distances with Applications in Bayesian OT Flow Matching](https://mlsanity.com/paper/2403.18705/) (2024)
- [Dynamic Conditional Optimal Transport through Simulation-Free Flows](https://mlsanity.com/paper/2404.04240/) (2024)
- [Flow Matching: Markov Kernels, Stochastic Processes and Transport Plans](https://mlsanity.com/paper/2501.16839/) (2025)
- [The Hidden Width of Deep ResNets: Tight Error Bounds and Phase Diagram](https://mlsanity.com/paper/2509.10167/) (2025)
- [Optimal and Diffusion Transports in Machine Learning](https://mlsanity.com/paper/2512.06797/) (2025)
- [Training Infinitely Deep and Wide Transformers](https://mlsanity.com/paper/2605.17660/) (2026)
- [Dropout and Random Gradient Masking Are Asymptotically Equivalent in Large ResNets](https://mlsanity.com/paper/2607.16761/) (2026)
- [Global Exponential Convergence of Two-Layer Linear Network Training](https://mlsanity.com/paper/2610.09356/) (2026)

## Abstract (arXiv)

We study the convergence of gradient flow for the training of deep neural networks. If Residual Neural Networks are a popular example of very deep architectures, their training constitutes a challenging optimization problem due notably to the non-convexity and the non-coercivity of the objective. Yet, in applications, those tasks are successfully solved by simple optimization algorithms such as gradient descent. To better understand this phenomenon, we focus here on a ``mean-field'' model of infinitely deep and arbitrarily wide ResNet, parameterized by probability measures over the product set of layers and parameters and with constant marginal on the set of layers. Indeed, in the case of shallow neural networks, mean field models have proven to benefit from simplified loss-landscapes and good theoretical guarantees when trained with gradient flow for the Wasserstein metric on the set of probability measures. Motivated by this approach, we propose to train our model with gradient flow w.r.t. the conditional Optimal Transport distance: a restriction of the classical Wasserstein distance which enforces our marginal condition. Relying on the theory of gradient flows in metric spaces we first show the well-posedness of the gradient flow equation and its consistency with the training of ResNets at finite width. Performing a local Polyak-Lojasiewicz analysis, we then show convergence of the gradient flow for well-chosen initializations: if the number of features is finite but sufficiently large and the risk is sufficiently small at initialization, the gradient flow converges towards a global minimizer. This is the first result of this type for infinitely deep and arbitrarily wide ResNets.