ZerO Initialization: Initializing Residual Networks with only Zeros and Ones

1 INTRODUCTION . An important question in training deep neural networks is how to initialize the weights . When Rumelhart et al . first introduced backward propagation , they initialized the weights randomly to break symmetries among different parameters and activations . Currently , random weight initialization is the de-facto practice across all architectures and tasks . However , choosing the variance of the initial distribution or equivalently the scale of the initial weights is a delicate balance when training deep neural networks . Too large a scale can lead to excessive amplification of the activations propagating through the network , resulting in exploding gradients during training . On the other hand , if the weights are initialized too small the activations may not propagate at all , resulting in vanishing gradients , which is especially an issue as the depth of the network grows . He et al . ( a ) ; Glorot and Bengio study the propagation of variance in the forward and backward passes at initialization , using it as a proxy for the magnitudes of the gradients . By choosing scale parameters that depend on the width of the neural networks , to upper bound the variance , they design initialization schemes that avoid exploding gradients . However , this still does not fully address the issue of vanishing gradients because it is not easy to estimate the lower bounds on the variance . In order to design networks robust to vanishing gradients at initialization , He et al . ( b ) introduced skip connections that take the activations of a layer , and add them to the next layer . Given enough skip connections , the activations can , in principle , always be propagated through the entire network . Thus , skip connections allow us to mitigate the vanishing gradient problem and possibly avoid it completely . Zhang et al . ( a ) ; Yang and Schoenholz observe that initializing the residual branches ( i.e. , the branch that is not the skip connection ) at ( or close to ) zero benefits signal propagation and optimization . To achieve this , Zhang et al . ( a ) initialize the last layer of each residual branch to zero , which is 1We share our code via an anonymous link here . also shown to stabilize training without batch normalization . However , previous studies rely on random weights to propagate signals to layers without skip connections , which still requires the consideration of the weight variance at initialization . Our approach : In this work , we show that all randomness in the weight initialization can be removed , resulting in a fully deterministic initialization scheme ZerO , which initializes residual networks with only zeros and ones . As illustrated in Figure 1 , ZerO achieves it by adding extra skip connections to enable non-zero signal propagation throughout the network . Additionally , it applies the Hadamard transform to avoid the training degeneracies in the latent space when there is a width expansion across layers . Moreover , our approach does not add any new learnable parameters to the standard architectures and does not involve any significant computational overhead . We find that ZerO achieves state-of-the-art performance over various image classification tasks . Our contributions are summarized as follows : 1 . We carefully identify the training difficulties that occur when initializing residual networks to zero , including 1 ) the dead neuron problem—addressed by employing skip connections at each layer , and 2 ) the degeneracies due to width expansion—addressed by applying the Hadamard transform . 2 . We propose ZerO initialization for ResNet as an example to show how skip connections enable a deterministic initialization . Our experiments show that ResNet with ZerO achieves state-of-theart results over various image classification tasks 3 . We find ZerO works well without batch normalization during training . With proper regularization , ZerO beats random initialization methods such as Kaiming initialization while almost matching the batch normalization baseline . 4 . We observe that compared to random initialization methods , the deterministic ZerO initialization achieves 20 % -40 % lower standard deviation over repeated experiments , and thus ZerO enjoys the advantage of high reproducibility . 2 WHEN DOES A ZERO-INITIALIZED NEURAL NETWORK WORK ? . In this section , we discuss when a zero-initialized neural network works by identifying the training difficulties that occur when initializing the weights to zero . We also propose adequate methodologies addressing these problems in order to achieve zero initialization without damaging the performance . We use the fully connected neural network as a proxy to illustrate the difficulties and our idea . Consider a problem of learning a function ŷ = F ( x ) with parameters W by minimizing the squared error L = 12‖y − ŷ‖22 . ( x , y ) ∈ RNx×Ny represent the sample data , and the learning is accomplished via gradient descent . F ( x ) is a fully connected neural network with L layers : F ( x ) = ϕ ( WL ) ◦ ... ◦ϕ ( W 1 ) ◦x , where ◦ is composition operator , W l ∈ RP l×Ql for l ∈ 1 , ... , L , and ϕ is an nonlinearity . We ignore bias terms , and denote xl = ϕ ( W l xl−1 ) as the activations of the lth layer . We also define W li , : and W l : ,j to be ith row and jth column of the matrix W l . 2.1 SKIP CONNECTIONS SOLVE DEAD NEURON PROBLEM . For simplicity , we first assume that F ( x ) is a linear network with a fixed dimension , such that Nx = Ny = P l = Ql for any l , and ϕ is an identity mapping . In this case , if we initialize the weights to zero , every layer will output zero constantly no matter what the inputs are . On the other hand , every W l receives zero derivatives in the backward pass , and thus there is no weight that can move out of the zero during the entire training . This extreme case of the vanishing gradient problem , well-known under the name dead neuron problem , is caused by zero-initialized F ( x ) only generating zero signals at initialization . Skip connections help propagate signals through the network and thus mitigate the vanishing gradient problem . A natural approach towards practical zero initialization is therefore to equip every layers of the network with skip connections . This ensures that the signal is propagated even with zero weights , with each layer applying a simple identity function to the activations . For example , equipping each layer of F ( x ) with a skip connection , we obtain F ( x ) = ( WL + I ) ... ( W 1 + I ) x . The derivatives with respect to weight matrices are then given as ∂L ∂W 1 = ( ( WL + I ) ... ( W 2 + I ) ) > ( y − ŷ ) x > , ∂L ∂W l = ( ( WL + I ) ... ( W l+1 + I ) ) > ( y − ŷ ) ( ( W l−1 + I ) ... ( W 1 + I ) x ) > for l ∈ 2 , ... , L− 1 , ∂L ∂WL = ( y − ŷ ) ( ( WL−1 + I ) ... ( W 1 + I ) x ) > . As zero-initialized weights initialize F ( x ) as an identity mapping , the loss L at initialization is the error when directly using the inputs as the predicted outputs ŷ = x . Therefore , the derivatives at initialization uniformly become ∂L ∂W l = ( y − x ) x > for every W l. This is because the skip connections are applied to every layer , and thus the error derivatives flow back to every matrix equivalently . With the help of the skip connections , each weight matrix W l receives non-zero derivatives , which break the dead neuron problem for the entire network . This also works for residual networks with nonlinearity , as long as the derivative of the nonlinearity exists at zero . Although there exists nonlinearity that is non-differentiable at zero such as Relu ( Xu et al . ) , this can be solved easily by utilizing its subderivatives at zero . For example , we define the derivative of Relu to be its subderivative one at zero , which allows the signals to be identically propagated through the Relu . We will apply this special modification to Relu in our later experiments . We note that in our example above , at the first iteration , all weight matrices receive the same derivatives . However , as training progresses , the gradient that a weight matrix receives depends on its position in the network , and thus the full capacity of the network is utilized . 2.2 WIDTH EXPANSION LEADS TO TRAINING DEGENERACY . For the deep residual networks discussed above , the dimensions of all hidden layers are equal to the input and output dimensions . However , for modern residual networks , the hidden dimensions are usually varying and larger than the input dimension , such as the large channel dimension in intermediate layers of ResNet . We find that for these networks , directly initializing weights to zero leads to another kind of dead neuron problem , which can not be avoided even with skip connections . To distinguish it from the previous dead neuron problem , we will refer to this problem as training degeneracy . To simplify the presentation , we explain this phenomenon in the setting where the nonlinearities are applied to the outputs of skip connections from previous layers . However , an equivalent phenomenon is present in the case where skip connections skip the nonlinearity as well , which is the variant we are using in Section 3 . To understand the cause of this degeneracy , we consider a simple 2-layer residual network : F ( x ) = Relu ( W 2 + I2 ) ◦ Relu ( W 1 + I1 ) x where W 1 , I1 ∈ R2×1 and W 2 , I2 ∈ R2×2 . Nx and Ny are changed to 1 and 2 accordingly , and we use Relu as the activation . Because the dimension is increased in the network , we apply zero padding to the skip connection I1 to match the dimension , where I10,0 = 1 and I 1 1,0 = 0 . This duplicates the previous activations while padding additional dimensions using zero , which is a standard parameter-free method applied in ResNet ( He et al. , c ) . For such a network , there exists a training degeneracy when directly initializing all the weights to zero . As illustrated in Figure 2 , at initialization , the second dimension of the first layer activations x11 receives a zero signal padded by I 1 , and it propagates the zero signal to ŷ as well . During backpropagation , ∂L ∂W 21,1 and ∂L ∂W 20,1 are zero because x11 = 0 . Also , ∂L ∂W 11,0 is zero as it receives the zero derivative from ∂L ∂x11 . Because x11 and ∂L ∂x11 keep receiving the zero signal , W 11,0 , W 2 0,1 and W 2 1,1 are bound to be zero for the entirety of training , restricting the expressivity of the network . 2.3 HADAMARD TRANSFORM AVOIDS TRAINING DEGENERACY . The reason for the degeneracy is that the skip connection I1 with zero padding is only able to map the input x into a low-dimensional subspace of x1 . As illustrated in Figure 2 , when initializing all the weights to zero , I1 maps the single dimension x to a horizontal line ( dashed line ) in the twodimensional space . Because this subspace is aligned with the standard basis , its linear dimension is invariant under application of the Relu , or any component-wise nonlinearity . As the activations x1 stay in a 1-dimensional subspace span , its associated derivatives stay in such a subspace as well . This eventually causes the zero-initialized weights W 1 to be updated into a low-dimensional space , as they solely depend on the low-dimensional derivatives . To ensure that the first hidden space receives high-dimensional activations x1 , we propose to apply a Hadamard transform H1 after I1 to spread the information into a new space . The Hadamard transform is an example of the generalized family of Fourier transforms that performs an orthogonal linear operation on 2m real numbers ( Pratt et al. ) . It consists of a Hadamard matrix and a normalization factor . A Hadamard matrix is defined as follows : Definition 1 ( Hadamard matrix ) . For any Hadamard matrix Hm ∈ R2 m×2m where m is a positive integer , we define H0 = 1 by the identity , and the matrix with large m is defined recursively : Hm = ( Hm−1 Hm−1 Hm−1 −Hm−1 ) = 1 1 1 1 . . . 1 −1 1 −1 . . . 1 1 −1 −1 . . . 1 −1 −1 1 . . . . . . . . . . . . . . . . . . ∈ R2m×2m . We rescale the Hadamard matrix by 2− ( m−1 ) /2 , resulting in the orthonormal Hadamard transform . The Hadamard transform rotates the standard basis into a new basis , such that each element of the new basis is equally weakly aligned with every element of the standard basis . For example , in a two-dimensional space , the Hadamard transform rotates the standard basis by an angle of 45 degree , as illustrated in Figure 2 2 . 2We use a column-order-reversed Hadamard matrix in the example to simplify the presentation . This is equivalent to using a standard Hadamard matrix with order-reversed I1 , without affecting our conclusions . When applying component-wise nonlinearities to the spaces spanned by the elements of this basis , their linear dimension increases . In our example , after applying the Hadamard matrix H1 to I1x , it forms a new space that is not aligned with the standard basis . Under the application of the Relu , the new space ( i.e. , H1I1x ) is spread into a 2-dimensional span in x1 , which ensures that neither x10 nor x11 stay at zero forever . Therefore , with the Hadamard transform , the activation x 1 1 can propagate the non-zero signals to the next layer . This ensures the derivatives of W 21,1 and W 2 0,1 to be non-zero , and thus breaks the training degeneracy .

Summary: In this paper, the authors propose some modifications to the classic ResNet architecture and initialization scheme, so that the network can be trained with fully deterministic initialization, even without any normalization layers. The authors motivate their design choices by identifying and analyzing the problems with the forward and backward dynamics of some naive ideas. Then the authors present their approach and explain why each choice they make can solve the specific problems they have presented. Then they evaluate their proposed method along with some baselines on two popular image classification datasets, showing that their deterministic initialization method is on par with the ones with random weight initialization. Finally, they present related work and conclusion.