Exponential convergence rates for Batch Normalization: The power of length-direction decoupling in non-convex optimization

By Jonas Köhler, Hadi Daneshmand, Aurélien Lucchi, Thomas Hofmann, M. Zhou, Klaus Neymeyr. 2018

98 citations recorded by the source corpus.

Normalization techniques such as Batch Normalization have been applied successfully for training deep neural networks. Yet, despite its apparent empirical benefits, the reasons behind the success of Batch Normalization are mostly hypothetical. We here aim to provide a more thorough theoretical understanding from a classical optimization perspective. Our main contribution towards this goal is the identification of various problem instances in the realm of machine learning where % -- under certain assumptions-- Batch Normalization can provably accelerate optimization. We argue that this acceleration is due to the fact that Batch Normalization splits the optimization task into optimizing length and direction of the parameters separately. This allows gradient-based methods to leverage a favourable global structure in the loss landscape that we prove to exist in Learning Halfspace problems and neural network training with Gaussian inputs. We thereby turn Batch Normalization from an effective practical heuristic into a provably converging algorithm for these settings. Furthermore, we substantiate our analysis with empirical evidence that suggests the validity of our theoretical results in a broader context.

Subfields: Artificial Intelligence, Computational Mechanics

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computational Mechanics

Citations in the corpus

An Exponential Learning Rate Schedule for Deep Learning (2019, background): “(Kohler et al., 2018) essentially shows linear model with BN could achieve exponential convergence rate assuming gaussian inputs, but their analysis is for a variant of GD with an inner optimization loop rather than GD itself.”
FedBN: Federated Learning on Non-IID Features via Local Batch Normalization (2021, methodology): “Recent machine learning theory (1)We omit centering neurons as well as learning the shift parameter of BN for the neural network analysis because of the assumption that x is zero-mean and the two layer network setting (Kohler et al., 2019; Salimans & Kingma, 2016).”
DO-Conv: Depthwise Over-Parameterized Convolutional Layer (2020, background): “Normalization layers have been widely used for improving the training of CNNs, while the reason for their success is still being actively studied [31, 4, 21, 20].”
Information Geometry of Orthogonal Initializations and Training (2018, result): “Furthermore, we observe that the rate at which Stiefel and Oblique networks decrease training loss is not dependent on their gradient smoothness, a result which is consistent with the recent analysis of [17].”
Batch normalization provably avoids ranks collapse for randomly initialised deep networks (2020, background): “…“internal covariate shift” hypothesis [17] have appeared in the literature: decoupling optimization of direction and length of the parameters [20], auto-tuning of the learning rate for stochastic gradient descent [3], widening the learning rate range [7], alleviating sharpness of the Fisher…”
Normalization Techniques in Training DNNs: Methodology, Analysis and Application (2020, background): “Besides, the work of [203] demonstrated that BN obtains an accelerated convergence on the (possibly) nonconvex problem of learning half-spaces with Gaussian inputs, from a length-direction decoupling perspective.”