# A Riemannian Mean Field Formulation for Two-layer Neural Networks with Batch Normalization

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 20 references that are in the map.
- First posted on arXiv 17 Oct 2021.

## Builds on (most cited first, 16 of 16)

- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [Layer Normalization](https://mlsanity.com/paper/1607.06450/) (2016)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [Self-Normalizing Neural Networks](https://mlsanity.com/paper/1706.02515/) (2017)
- [Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks](https://mlsanity.com/paper/1602.07868/) (2016)
- [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)
- [How Does Batch Normalization Help Optimization?](https://mlsanity.com/paper/1805.11604/) (2018)
- [Understanding Batch Normalization](https://mlsanity.com/paper/1806.02375/) (2018)
- [Theoretical Analysis of Auto Rate-Tuning by Batch Normalization](https://mlsanity.com/paper/1812.03981/) (2018)
- [Towards Understanding Regularization in Batch Normalization](https://mlsanity.com/paper/1809.00846/) (2018)
- [Machine Learning from a Continuous Viewpoint](https://mlsanity.com/paper/1912.12777/) (2019)
- [Modeling from Features: a Mean-field Framework for Over-parameterized Deep Neural Networks](https://mlsanity.com/paper/2007.01452/) (2020)
- [Riemannian approach to batch normalization](https://mlsanity.com/paper/1709.09603/) (2017)
- [Mean Field Limit of the Learning Dynamics of Multilayer Neural Networks](https://mlsanity.com/paper/1902.02880/) (2019)
- [Mean-field Analysis of Batch Normalization](https://mlsanity.com/paper/1903.02606/) (2019)

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

- [Understanding the Generalization Benefit of Normalization Layers: Sharpness Reduction](https://mlsanity.com/paper/2206.07085/) (2022)

## Abstract (arXiv)

The training dynamics of two-layer neural networks with batch normalization (BN) is studied. It is written as the training dynamics of a neural network without BN on a Riemannian manifold. Therefore, we identify BN's effect of changing the metric in the parameter space. Later, the infinite-width limit of the two-layer neural networks with BN is considered, and a mean-field formulation is derived for the training dynamics. The training dynamics of the mean-field formulation is shown to be the Wasserstein gradient flow on the manifold. Theoretical analysis are provided on the well-posedness and convergence of the Wasserstein gradient flow.