# Accelerating Training of Batch Normalization: A Manifold Perspective

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

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 25 of its 32 references that are in the map.
- First posted on arXiv 8 Jan 2021.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [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)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Group Normalization](https://mlsanity.com/paper/1803.08494/) (2018)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Stochastic First- and Zeroth-order Methods for Nonconvex Stochastic Programming](https://mlsanity.com/paper/1309.5549/) (2013)
- [How Does Batch Normalization Help Optimization?](https://mlsanity.com/paper/1805.11604/) (2018)
- [Cheap Orthogonal Constraints in Neural Networks: A Simple Parametrization of the Orthogonal and Unitary Group](https://mlsanity.com/paper/1901.08428/) (2019)
- [An Exponential Learning Rate Schedule for Deep Learning](https://mlsanity.com/paper/1910.07454/) (2019)
- [Path-SGD: Path-Normalized Optimization in Deep Neural Networks](https://mlsanity.com/paper/1506.02617/) (2015)
- [First-order Methods for Geodesically Convex Optimization](https://mlsanity.com/paper/1602.06053/) (2016)
- [Orthogonal Weight Normalization: Solution to Optimization over Multiple Dependent Stiefel Manifolds in Deep Neural Networks](https://mlsanity.com/paper/1709.06079/) (2017)
- [Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform](https://mlsanity.com/paper/2002.01113/) (2020)
- [Trivializations for Gradient-Based Optimization on Manifolds](https://mlsanity.com/paper/1909.09501/) (2019)
- [Theoretical Analysis of Auto Rate-Tuning by Batch Normalization](https://mlsanity.com/paper/1812.03981/) (2018)
- [Global rates of convergence for nonconvex optimization on manifolds](https://mlsanity.com/paper/1605.08101/) (2016)
- [Unified Convergence Analysis of Stochastic Momentum Methods for Convex and Non-convex Optimization](https://mlsanity.com/paper/1604.03257/) (2016)
- [Riemannian approach to batch normalization](https://mlsanity.com/paper/1709.09603/) (2017)
- [G-SGD: Optimizing ReLU Neural Networks in its Positively Scale-Invariant Space](https://mlsanity.com/paper/1802.03713/) (2018)
- [WNGrad: Learn the Learning Rate in Gradient Descent](https://mlsanity.com/paper/1803.02865/) (2018)
- [Symmetry-invariant optimization in deep networks](https://mlsanity.com/paper/1511.01754/) (2015)
- [Capacity Control of ReLU Neural Networks by Basis-path Norm](https://mlsanity.com/paper/1809.07122/) (2018)
- [Projection Based Weight Normalization for Deep Neural Networks](https://mlsanity.com/paper/1710.02338/) (2017)

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

- [Understanding Gradient Descent on Edge of Stability in Deep Learning](https://mlsanity.com/paper/2205.09745/) (2022)
- [Symmetry in Neural Network Parameter Spaces](https://mlsanity.com/paper/2506.13018/) (2025)
- [Improving Learning to Optimize Using Parameter Symmetries](https://mlsanity.com/paper/2504.15399/) (2025)
- [Continuous-time Riemannian SGD and SVRG Flows on Wasserstein Probabilistic Space](https://mlsanity.com/paper/2401.13530/) (2024)

## Abstract (arXiv)

Batch normalization (BN) has become a critical component across diverse deep neural networks. The network with BN is invariant to positively linear re-scale transformation, which makes there exist infinite functionally equivalent networks with different scales of weights. However, optimizing these equivalent networks with the first-order method such as stochastic gradient descent will obtain a series of iterates converging to different local optima owing to their different gradients across training. To obviate this, we propose a quotient manifold PSI manifold, in which all the equivalent weights of the network with BN are regarded as the same element. Next, we construct gradient descent and stochastic gradient descent on the proposed PSI manifold to train the network with BN. The two algorithms guarantee that every group of equivalent weights (caused by positively re-scaling) converge to the equivalent optima. Besides that, we give convergence rates of the proposed algorithms on the PSI manifold. The results show that our methods accelerate training compared with the algorithms on the Euclidean weight space. Finally, empirical results verify that our algorithms consistently improve the existing methods in both convergence rate and generalization ability under various experimental settings.