Accelerating Training of Batch Normalization: A Manifold Perspective
This work provides a theoretical and practical improvement for optimizing deep neural networks using Batch Normalization, potentially benefiting researchers and practitioners in deep learning by accelerating training and improving generalization.
The paper addresses the issue of multiple equivalent networks due to Batch Normalization's invariance to positive linear re-scaling, which causes first-order methods to converge to different local optima. The authors propose a quotient manifold, the PSI manifold, where equivalent weights are treated as the same element, and demonstrate that their manifold-based optimization algorithms accelerate training compared to Euclidean space methods.
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 \emph{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.