Fixup Initialization: Residual Learning Without Normalization

Normalization layers are a staple in state-of-the-art deep neural network architectures . They are widely believed to stabilize training , enable higher learning rate , accelerate convergence and improve generalization , though the reason for their effectiveness is still an active research topic . In this work , we challenge the commonly-held beliefs by showing that none of the perceived benefits is unique to normalization . Specifically , we propose fixed-update initialization ( Fixup ) , an initialization motivated by solving the exploding and vanishing gradient problem at the beginning of training via properly rescaling a standard initialization . We find training residual networks with Fixup to be as stable as training with normalization — even for networks with 10,000 layers . Furthermore , with proper regularization , Fixup enables residual networks without normalization to achieve state-of-the-art performance in image classification and machine translation . 1 INTRODUCTION . Artificial intelligence applications have witnessed major advances in recent years . At the core of this revolution is the development of novel neural network models and their training techniques . For example , since the landmark work of He et al . ( 2016 ) , most of the state-of-the-art image recognition systems are built upon a deep stack of network blocks consisting of convolutional layers and additive skip connections , with some normalization mechanism ( e.g. , batch normalization ( Ioffe & Szegedy , 2015 ) ) to facilitate training and generalization . Besides image classification , various normalization techniques ( Ulyanov et al. , 2016 ; Ba et al. , 2016 ; Salimans & Kingma , 2016 ; Wu & He , 2018 ) have been found essential to achieving good performance on other tasks , such as machine translation ( Vaswani et al. , 2017 ) and generative modeling ( Zhu et al. , 2017 ) . They are widely believed to have multiple benefits for training very deep neural networks , including stabilizing learning , enabling higher learning rate , accelerating convergence , and improving generalization . Despite the enormous empirical success of training deep networks with normalization , and recent progress on understanding the working of batch normalization ( Santurkar et al. , 2018 ) , there is currently no general consensus on why these normalization techniques help training residual neural networks . Intrigued by this topic , in this work we study ( i ) without normalization , can a deep residual network be trained reliably ? ( And if so , ) ( ii ) without normalization , can a deep residual network be trained with the same learning rate , converge at the same speed , and generalize equally well ( or even better ) ? Perhaps surprisingly , we find the answers to both questions are Yes . In particular , we show : • Why normalization helps training . We derive a lower bound for the gradient norm of a residual network at initialization , which explains why with standard initializations , normalization techniques are essential for training deep residual networks at maximal learning rate . ( Section 2 ) ∗Work done at Facebook . Equal contribution . †Work done at Facebook . Equal contribution . ‡Work done at Facebook . • Training without normalization . We propose Fixup , a method that rescales the standard initialization of residual branches by adjusting for the network architecture . Fixup enables training very deep residual networks stably at maximal learning rate without normalization . ( Section 3 ) • Image classification . We apply Fixup to replace batch normalization on image classification benchmarks CIFAR-10 ( with Wide-ResNet ) and ImageNet ( with ResNet ) , and find Fixup with proper regularization matches the well-tuned baseline trained with normalization . ( Section 4.2 ) • Machine translation . We apply Fixup to replace layer normalization on machine translation benchmarks IWSLT and WMT using the Transformer model , and find it outperforms the baseline and achieves new state-of-the-art results on the same architecture . ( Section 4.3 ) In the remaining of this paper , we first analyze the exploding gradient problem of residual networks at initialization in Section 2 . To solve this problem , we develop Fixup in Section 3 . In Section 4 we quantify the properties of Fixup and compare it against state-of-the-art normalization methods on real world benchmarks . A comparison with related work is presented in Section 5 . 2 PROBLEM : RESNET WITH STANDARD INITIALIZATIONS LEAD TO EXPLODING GRADIENTS . Standard initialization methods ( Glorot & Bengio , 2010 ; He et al. , 2015 ; Xiao et al. , 2018 ) attempt to set the initial parameters of the network such that the activations neither vanish nor explode . Unfortunately , it has been observed that without normalization techniques such as BatchNorm they do not account properly for the effect of residual connections and this causes exploding gradients . Balduzzi et al . ( 2017 ) characterizes this problem for ReLU networks , and we will generalize this to residual networks with positively homogenous activation functions . A plain ( i.e . without normalization layers ) ResNet with residual blocks { F1 , . . . , FL } and input x0 computes the activations as xl = x0 + l−1∑ i=0 Fi ( xi ) . ( 1 ) ResNet output variance grows exponentially with depth . Here we only consider the initialization , view the input x0 as fixed , and consider the randomness of the weight initialization . We analyze the variance of each layer xl , denoted by Var [ xl ] ( which is technically defined as the sum of the variance of all the coordinates of xl . ) For simplicity we assume the blocks are initialized to be zero mean , i.e. , E [ Fl ( xl ) | xl ] = 0 . By xl+1 = xl + Fl ( xl ) , and the law of total variance , we have Var [ xl+1 ] = E [ Var [ F ( xl ) |xl ] ] + Var ( xl ) . Resnet structure prevents xl from vanishing by forcing the variance to grow with depth , i.e . Var [ xl ] < Var [ xl+1 ] if E [ Var [ F ( xl ) |xl ] ] > 0 . Yet , combined with initialization methods such as He et al . ( 2015 ) , the output variance of each residual branch Var [ Fl ( xl ) |xl ] will be about the same as its input variance Var [ xl ] , and thus Var [ xl+1 ] ≈ 2Var [ xl ] . This causes the output variance to explode exponentially with depth without normalization ( Hanin & Rolnick , 2018 ) for positively homogeneous blocks ( see Definition 1 ) . This is detrimental to learning because it can in turn cause gradient explosion . As we will show , at initialization , the gradient norm of certain activations and weight tensors is lower bounded by the cross-entropy loss up to some constant . Intuitively , this implies that blowup in the logits will cause gradient explosion . Our result applies to convolutional and linear weights in a neural network with ReLU nonlinearity ( e.g. , feed-forward network , CNN ) , possibly with skip connections ( e.g. , ResNet , DenseNet ) , but without any normalization . Our analysis utilizes properties of positively homogeneous functions , which we now introduce . Definition 1 ( positively homogeneous function of first degree ) . A function f : Rm → Rn is called positively homogeneous ( of first degree ) ( p.h . ) if for any input x ∈ Rm and α > 0 , f ( αx ) = αf ( x ) . Definition 2 ( positively homogeneous set of first degree ) . Let θ = { θi } i∈S be the set of parameters of f ( x ) and θph = { θi } i∈Sph⊂S . We call θph a positively homogeneous set ( of first degree ) ( p.h . set ) if for any α > 0 , f ( x ; θ \ θph , αθph ) = αf ( x ; θ \ θph , θph ) , where αθph denotes { αθi } i∈Sph . Intuitively , a p.h . set is a set of parameters θph in function f such that for any fixed input x and fixed parameters θ \ θph , f̄ ( θph ) , f ( x ; θ \ θph , θph ) is a p.h . function . Examples of p.h . functions are ubiquitous in neural networks , including various kinds of linear operations without bias ( fully-connected ( FC ) and convolution layers , pooling , addition , concatenation and dropout etc . ) as well as ReLU nonlinearity . Moreover , we have the following claim : Proposition 1 . A function that is the composition of p.h . functions is itself p.h . We study classification problems with c classes and the cross-entropy loss . We use f to denote a neural network function except for the softmax layer . Cross-entropy loss is defined as ` ( z , y ) , −yT ( z− logsumexp ( z ) ) where y is the one-hot label vector , z , f ( x ) ∈ Rc is the logits where zi denotes its i-th element , and logsumexp ( z ) , log ( ∑ i∈ [ c ] exp ( zi ) ) . Consider a minibatch of training examples DM = { ( x ( m ) , y ( m ) ) } Mm=1 and the average cross-entropy loss ` avg ( DM ) , 1 M ∑M m=1 ` ( f ( x ( m ) ) , y ( m ) ) , where we use ( m ) to index quantities referring to the m-th example . ‖ · ‖ denotes any valid norm . We only make the following assumptions about the network f : 1. f is a sequential composition of network blocks { fi } Li=1 , i.e . f ( x0 ) = fL ( fL−1 ( . . . f1 ( x0 ) ) ) , each of which is composed of p.h . functions . 2 . Weight elements in the FC layer are i.i.d . sampled from a zero-mean symmetric distribution . These assumptions hold at initialization if we remove all the normalization layers in a residual network with ReLU nonlinearity , assuming all the biases are initialized at 0 . Our results are summarized in the following two theorems , whose proofs are listed in the appendix : Theorem 1 . Denote the input to the i-th block by xi−1 . With Assumption 1 , we have∥∥∥∥ ∂ ` ∂xi−1 ∥∥∥∥ ≥ ` ( z , y ) −H ( p ) ‖xi−1‖ , ( 2 ) where p is the softmax probabilities and H denotes the Shannon entropy . Since H ( p ) is upper bounded by log ( c ) and ‖xi−1‖ is small in the lower blocks , blowup in the loss will cause large gradient norm with respect to the lower block input . Our second theorem proves a lower bound on the gradient norm of a p.h . set in a network . Theorem 2 . With Assumption 1 , we have∥∥∥∥∂ ` avg∂θph ∥∥∥∥ ≥ 1M‖θph‖ M∑ m=1 ` ( z ( m ) , y ( m ) ) −H ( p ( m ) ) , G ( θph ) . ( 3 ) Furthermore , with Assumptions 1 and 2 , we have EG ( θph ) ≥ E [ maxi∈ [ c ] zi ] − log ( c ) ‖θph‖ . ( 4 ) It remains to identify such p.h . sets in a neural network . In Figure 2 we provide three examples of p.h . sets in a ResNet without normalization . Theorem 2 suggests that these layers would suffer from the exploding gradient problem , if the logits z blow up at initialization , which unfortunately would occur in a ResNet without normalization if initialized in a traditional way . This motivates us to introduce a new initialization in the next section . 3 FIXUP : UPDATE A RESIDUAL NETWORK Θ ( η ) PER SGD STEP Our analysis in the previous section points out the failure mode of standard initializations for training deep residual network : the gradient norm of certain layers is in expectation lower bounded by a quantity that increases indefinitely with the network depth . However , escaping this failure mode does not necessarily lead us to successful training — after all , it is the whole network as a function that we care about , rather than a layer or a network block . In this section , we propose a top-down design of a new initialization that ensures proper update scale to the network function , by simply rescaling a standard initialization . To start , we denote the learning rate by η and set our goal : f ( x ; θ ) is updated by Θ ( η ) per SGD step after initialization as η → 0 . That is , ‖∆f ( x ) ‖ = Θ ( η ) where ∆f ( x ) , f ( x ; θ − η ∂∂θ ` ( f ( x ) , y ) ) − f ( x ; θ ) . Put another way , our goal is to design an initialization such that SGD updates to the network function are in the right scale and independent of the depth . We define the Shortcut as the shortest path from input to output in a residual network . The Shortcut is typically a shallow network with a few trainable layers.1 We assume the Shortcut is initialized using a standard method , and focus on the initialization of the residual branches . Residual branches update the network in sync . To start , we first make an important observation that the SGD update to each residual branch changes the network output in highly correlated directions . This implies that if a residual network has L residual branches , then an SGD step to each residual branch should change the network output by Θ ( η/L ) on average to achieve an overall Θ ( η ) update . We defer the formal statement and its proof until Appendix B.1 . Study of a scalar branch . Next we study how to initialize a residual branch with m layers so that its SGD update changes the network output by Θ ( η/L ) . We assume m is a small positive integer ( e.g. , 2 or 3 ) . As we are only concerned about the scale of the update , it is sufficiently instructive to study the scalar case , i.e. , F ( x ) = ( ∏m i=1 ai ) x where a1 , . . . , am , x ∈ R+ . For example , the standard initialization methods typically initialize each layer so that the output ( after nonlinear activation ) preserves the input variance , which can be modeled as setting ∀i ∈ [ m ] , ai = 1 . In turn , setting ai to a positive number other than 1 corresponds to rescaling the i-th layer by ai . Through deriving the constraints for F ( x ) to make Θ ( η/L ) updates , we will also discover how to rescale the weight layers of a standard initialization as desired . In particular , we show the SGD 1For example , in the ResNet architecture ( e.g. , ResNet-50 , ResNet-101 or ResNet-152 ) for ImageNet classification , the Shortcut is always a 6-layer network with five convolution layers and one fully-connected layer , irrespective of the total depth of the whole network . update to F ( x ) is Θ ( η/L ) if and only if the initialization satisfies the following constraint : ∏ i∈ [ m ] \ { j } ai x = Θ ( 1√ L ) , where j ∈ arg min k ak ( 5 ) We defer the derivation until Appendix B.2 . Equation ( 5 ) suggests new methods to initialize a residual branch through rescaling the standard initialization of i-th layer in a residual branch by its corresponding scalar ai . For example , we could set ∀i ∈ [ m ] , ai = L− 1 2m−2 . Alternatively , we could start the residual branch as a zero function by setting am = 0 and ∀i ∈ [ m − 1 ] , ai = L− 1 2m−2 . In the second option , the residual branch does not need to “ unlearn ” its potentially bad random initial state , which can be beneficial for learning . Therefore , we use the latter option in our experiments , unless otherwise specified . The effects of biases and multipliers . With proper rescaling of the weights in all the residual branches , a residual network is supposed to be updated by Θ ( η ) per SGD step — our goal is achieved . However , in order to match the training performance of a corresponding network with normalization , there are two more things to consider : biases and multipliers . Using biases in the linear and convolution layers is a common practice . In normalization methods , bias and scale parameters are typically used to restore the representation power after normalization.2 Intuitively , because the preferred input/output mean of a weight layer may be different from the preferred output/input mean of an activation layer , it also helps to insert bias terms in a residual network without normalization . Empirically , we find that inserting just one scalar bias before each weight layer and nonlinear activation layer significantly improves the training performance . Multipliers scale the output of a residual branch , similar to the scale parameters in batch normalization . They have an interesting effect on the learning dynamics of weight layers in the same branch . Specifically , as the stochastic gradient of a layer is typically almost orthogonal to its weight , learning rate decay tends to cause the weight norm equilibrium to shrink when combined with L2 weight decay ( van Laarhoven , 2017 ) . In a branch with multipliers , this in turn causes the growth of the multipliers , increasing the effective learning rate of other layers . In particular , we observe that inserting just one scalar multiplier per residual branch mimics the weight norm dynamics of a network with normalization , and spares us the search of a new learning rate schedule . Put together , we propose the following method to train residual networks without normalization : Fixup initialization ( or : How to train a deep residual network without normalization ) 1 . Initialize the classification layer and the last layer of each residual branch to 0 . 2 . Initialize every other layer using a standard method ( e.g. , He et al . ( 2015 ) ) , and scale only the weight layers inside residual branches by L− 1 2m−2 . 3 . Add a scalar multiplier ( initialized at 1 ) in every branch and a scalar bias ( initialized at 0 ) before each convolution , linear , and element-wise activation layer . It is important to note that Rule 2 of Fixup is the essential part as predicted by Equation ( 5 ) . Indeed , we observe that using Rule 2 alone is sufficient and necessary for training extremely deep residual networks . On the other hand , Rule 1 and Rule 3 make further improvements for training so as to match the performance of a residual network with normalization layers , as we explain in the above text.3 We find ablation experiments confirm our claims ( see Appendix C.1 ) . 2For example , in batch normalization gamma and beta parameters are used to affine-transform the normalized activations per each channel . 3It is worth noting that the design of Fixup is a simplification of the common practice , in that we only introduce O ( K ) parameters beyond convolution and linear weights ( since we remove bias terms from convolution and linear layers ) , whereas the common practice includes O ( KC ) ( Ioffe & Szegedy , 2015 ; Salimans & Kingma , 2016 ) or O ( KCWH ) ( Ba et al. , 2016 ) additional parameters , where K is the number of layers , C is the max number of channels per layer and W , H are the spatial dimension of the largest feature maps . Our initialization and network design is consistent with recent theoretical work Hardt & Ma ( 2016 ) ; Li et al . ( 2018 ) , which , in much more simplified settings such as linearized residual nets and quadratic neural nets , propose that small initialization tend to stabilize optimization and help generalizaiton . However , our approach suggests that more delicate control of the scale of the initialization is beneficial.4

Summary: A method is presented for initialization and normalization of deep residual networks. The method is based on interesting observations regarding forward and backward explosion in such networks with the standard Xavier or (He, 2015) initializations. Experiments with the new method show that it is able to learn with very deep networks, and that its performance is on a par with the best results obtained by other networks with more explicit normalization.