Equi-normalization of Neural Networks

Modern neural networks are over-parametrized . In particular , each rectified linear hidden unit can be modified by a multiplicative factor by adjusting input and output weights , without changing the rest of the network . Inspired by the SinkhornKnopp algorithm , we introduce a fast iterative method for minimizing the ` 2 norm of the weights , equivalently the weight decay regularizer . It provably converges to a unique solution . Interleaving our algorithm with SGD during training improves the test accuracy . For small batches , our approach offers an alternative to batchand group- normalization on CIFAR-10 and ImageNet with a ResNet-18 . 1 INTRODUCTION . Deep Neural Networks ( DNNs ) have achieved outstanding performance across a wide range of empirical tasks such as image classification ( Krizhevsky et al. , 2012 ) , image segmentation ( He et al. , 2017 ) , speech recognition ( Hinton et al. , 2012a ) , natural language processing ( Collobert et al. , 2011 ) or playing the game of Go ( Silver et al. , 2017 ) . These successes have been driven by the availability of large labeled datasets such as ImageNet ( Russakovsky et al. , 2015 ) , increasing computational power and the use of deeper models ( He et al. , 2015b ) . Although the expressivity of the function computed by a neural network grows exponentially with depth ( Pascanu et al. , 2013 ; Raghu et al. , 2017 ; Telgarsky , 2016 ) , in practice deep networks are vulnerable to both over- and underfitting ( Glorot & Bengio , 2010 ; Krizhevsky et al. , 2012 ; He et al. , 2015b ) . Widely used techniques to prevent DNNs from overfitting include regularization methods such as weight decay ( Krogh & Hertz , 1992 ) , Dropout ( Hinton et al. , 2012b ) and various data augmentation schemes ( Krizhevsky et al. , 2012 ; Simonyan & Zisserman , 2014 ; Szegedy et al. , 2014 ; He et al. , 2015b ) . Underfitting can occur if the network gets stuck in a local minima , which can be avoided by using stochastic gradient descent algorithms ( Bottou , 2010 ; Duchi et al. , 2011 ; Sutskever et al. , 2013 ; Kingma & Ba , 2014 ) , sometimes along with carefully tuned learning rate schedules ( He et al. , 2015b ; Goyal et al. , 2017 ) . Training deep networks is particularly challenging due to the vanishing/exploding gradient problem . It has been studied for Recurrent Neural networks ( RNNs ) ( Hochreiter et al. , 2001 ) as well as standard feedforward networks ( He et al. , 2015a ; Mishkin & Matas , 2015 ) . After a few iterations , the gradients computed during backpropagation become either too small or too large , preventing the optimization scheme from converging . This is alleviated by using non-saturating activation functions such as rectified linear units ( ReLUs ) ( Krizhevsky et al. , 2012 ) or better initialization schemes preserving the variance of the input across layers ( Glorot & Bengio , 2010 ; Mishkin & Matas , 2015 ; He et al. , 2015a ) . Failure modes that prevent the training from starting have been theoretically studied by Hanin & Rolnick ( 2018 ) . Two techniques in particular have allowed vision models to achieve “ super-human ” accuracy . Batch Normalization ( BN ) was developed to train Inception networks ( Ioffe & Szegedy , 2015 ) . It introduces intermediate layers that normalize the features by the mean and variance computed within the current batch . BN is effective in reducing training time , provides better generalization capabilities after training and diminishes the need for a careful initialization . Network architectures such as ResNet ( He et al. , 2015b ) and DenseNet ( Huang et al. , 2016 ) use skip connections along with BN to improve the information flow during both the forward and backward passes. . . . `` # ... `` # '' # `` # Wk ⇥ 1 ↵nk ⇥↵nk ⇥ 1 ↵1 ⇥↵1 ⇥ 1 ↵2 ⇥↵2 ⇥1/↵2 ⇥1/↵1 ⇥1/↵nk⇥1/ nk ⇥1/ 1 ⇥1/ 2 ⇥ 2⇥ 1 ⇥ 1 nk ⇥ nk ⇥ 1 1 ⇥ 1 2 Wk+1 Figure 1 : Matrices Wk and Wk+1 are updated by multiplying the columns of the first matrix with rescaling coefficients . The rows of the second matrix are inversely rescaled to ensure that the product of the two matrices is unchanged . The rescaling coefficients are strictly positive to ensure functional equivalence when the matrices are interleaved with ReLUs . This rescaling is applied iteratively to each pair of adjacent matrices . In this paper , we address the more complex cases of biases , convolutions , max-pooling or skip-connections to be able to balance modern CNN architectures . However , BN has some limitations . In particular , BN only works well with sufficiently large batch sizes ( Ioffe & Szegedy , 2015 ; Wu & He , 2018 ) . For sizes below 16 or 32 , the batch statistics have a high variance and the test error increases significantly . This prevents the investigation of highercapacity models because large , memory-consuming batches are needed in order for BN to work in its optimal range . In many use cases , including video recognition ( Carreira & Zisserman , 2017 ) and image segmentation ( He et al. , 2017 ) , the batch size restriction is even more challenging because the size of the models allows for only a few samples per batch . Another restriction of BN is that it is computationally intensive , typically consuming 20 % to 30 % of the training time . Variants such as Group Normalization ( GN ) ( Wu & He , 2018 ) cover some of the failure modes of BN . In this paper , we introduce a novel algorithm to improve both the training speed and generalization accuracy of networks by using their over-parameterization to regularize them . In particular , we focus on neural networks that are positive-rescaling equivalent ( Neyshabur et al. , 2015 ) , i.e . whose weights are identical up to positive scalings and matching inverse scalings . The main principle of our method , referred to as Equi-normalization ( ENorm ) , is illustrated in Figure 1 for the fullyconnected case . We scale two consecutive matrices with rescaling coefficients that minimize the joint ` p norm of those two matrices . This amounts to re-parameterizing the network under the constraint of implementing the same function . We conjecture that this particular choice of rescaling coefficients ensures a smooth propagation of the gradients during training . A limitation is that our current proposal , in its current form , can only handle learned skipconnections like those proposed in type-C ResNet . For this reason , we focus on architectures , in particular ResNet18 , for which the learning converges with learned skip-connection , as opposed to architectures like ResNet-50 for which identity skip-connections are required for convergence . In summary , • We introduce an iterative , batch-independent algorithm that re-parametrizes the network within the space of rescaling equivalent networks , thus preserving the function implemented by the network ; • We prove that the proposed Equi-normalization algorithm converges to a unique canonical parameterization of the network that minimizes the global ` p norm of the weights , or equivalently , when p = 2 , the weight decay regularizer ; • We extend ENorm to modern convolutional architectures , including the widely used ResNets , and show that the theoretical computational overhead is lower compared to BN ( ×50 ) and even compared to GN ( ×3 ) ; • We show that applying one ENorm step after each SGD step outperforms both BN and GN on the CIFAR-10 ( fully connected ) and ImageNet ( ResNet-18 ) datasets . • Our code is available at https : //github.com/facebookresearch/enorm . The paper is organized as follows . Section 2 reviews related work . Section 3 defines our Equinormalization algorithm for fully-connected networks and proves the convergence . Section 4 shows how to adapt ENorm to convolutional neural networks ( CNNs ) . Section 5 details how to employ ENorm for training neural networks and Section 6 presents our experimental results . 2 RELATED WORK . This section reviews methods improving neural network training and compares them with ENorm . Since there is a large body of literature in this research area , we focus on the works closest to the proposed approach . From early works , researchers have noticed the importance of normalizing the input of a learning system , and by extension the input of any layer in a DNN ( LeCun et al. , 1998 ) . Such normalization is applied either to the weights or to the activations . On the other hand , several strategies aim at better controlling the geometry of the weight space with respect to the loss function . Note that these research directions are not orthogonal . For example , explicitly normalizing the activations using BN has smoothing effects on the optimization landscape ( Santurkar et al. , 2018 ) . Normalizing activations . Batch Normalization ( Ioffe & Szegedy , 2015 ) normalizes the activations by using statistics computed along the batch dimension . As stated in the introduction , the dependency on the batch size leads BN to underperform when small batches are used . Batch Renormalization ( BR ) ( Ioffe , 2017 ) is a follow-up that reduces the sensitivity to the batch size , yet does not completely alleviate the negative effect of small batches . Several batch-independent methods operate on other dimensions , such as Layer Normalization ( channel dimension ) ( Ba et al. , 2016 ) and Instance-Normalization ( sample dimension ) ( Ulyanov et al. , 2016 ) . Parametric data-independent estimation of the mean and variance in every layer is investigated by Arpit et al . ( 2016 ) . However , these methods are inferior to BN in standard classification tasks . More recently , Group Normalization ( GN ) ( Wu & He , 2018 ) , which divides the channels into groups and normalizes independently each group , was shown to effectively replace BN for small batch sizes on computer vision tasks . Normalizing weights . Early weight normalization techniques only served to initialize the weights before training ( Glorot & Bengio , 2010 ; He et al. , 2015a ) . These methods aim at keeping the variance of the output activations close to one along the whole network , but the assumptions made to derive these initialization schemes may not hold as training evolves . More recently , Salimans & Kingma ( 2016 ) propose a polar-like re-parametrization of the weights to disentangle the direction from the norm of the weight vectors . Note that Weight Norm ( WN ) does require mean-only BN to get competitive results , as well as a greedy layer-wise initialization as mentioned in the paper . Optimization landscape . Generally , in the parameter space , the loss function moves quickly along some directions and slowly along others . To account for this anisotropic relation between the parameters of the model and the loss function , natural gradient methods have been introduced ( Amari , 1998 ) . They require storing and inverting the N × N curvature matrix , where N is the number of network parameters . Several works approximate the inverse of the curvature matrix to circumvent this problem ( Pascanu & Bengio , 2013 ; Marceau-Caron & Ollivier , 2016 ; Martens & Grosse , 2015 ) . Another method called Diagonal Rescaling ( Lafond et al. , 2017 ) proposes to tune a particular reparametrization of the weights with a block-diagonal approximation of the inverse curvature matrix . Finally , Neyshabur et al . ( 2015 ) propose a rescaling invariant path-wise regularizer and use it to derive Path-SGD , an approximate steepest descent with respect to the path-wise regularizer . Positioning . Unlike BN , Equi-normalization focuses on the weights and is independent of the concept of batch . Like Path-SGD , our goal is to obtain a balanced network ensuring a good backpropagation of the gradients , but our method explicitly re-balances the network using an iterative algorithm instead of using an implicit regularizer . Moreover , ENorm can be readily adapted to the convolutional case whereas Neyshabur et al . ( 2015 ) restrict themselves to the fully-connected case . In addition , the theoretical computational complexity of our method is much lower than the overhead introduced by BN or GN ( see Section 5 ) . Besides , WN parametrizes the weights in a polar-like manner , w = g × v/|v| , where g is a scalar and v are the weights , thus it does not balance the network but individually scales each layer . Finally , Sinkhorn ’ s algorithm aims at making a single matrix doubly stochastic , while we balance a product of matrices to minimize their global norm .

Summary: The authors propose a new weight re-parameterization technique called Equi-normalization (ENorm) inspired by the Sinkhorn-Knopp algorithm. The authors show that the proposed method preserve functionally equivalent property in respect of the output of the functions (Linear, Conv, and Max-Pool) and show also that ENorm converges to the global optimum through the optimization. The experimental results show that ENorm performs better than baseline methods on CIFAR-10 and ImageNet datasets.