GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training
Authors: Chen Zhu
Renkun Ni
Zheng Xu
Kezhi Kong
W Ronny Huang
Tom Goldstein
Venue: NeurIPS 2021 Poster (2021)
Keywords: Initialization
Transformers
Convolutional Networks

Abstract: Innovations in neural architectures have fostered significant breakthroughs in language modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-parameter choices and training instability if the network parameters are not properly initialized. A number of architecture-specific initialization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on a simple heuristic; the norm of each network layer is adjusted so that a single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing a scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many convolutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer architecture for machine translation, enabling training it without learning rate warmup using either Adam or SGD under a wide range of learning rates and momentum coefficients. Code is available at https://github.com/zhuchen03/gradinit.

GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training
Chen Zhu University of Maryland chenzhuêcs. umd. edu
Renkun Ni University of Maryland rn9zm#cs, umd. edu
Zheng Xu Google Research xuzhengdgoogle..co com
Kezhi Kong University of Maryland kongecs umd. edu
Ronny Huang Google Research wrh@google. com
Tom Goldstein University of Maryland tomg umd. edu
Abstract
Innovations in neural architectures have fostered significant breakthroughs in lan- guage modeling and computer vision. Unfortunately, novel architectures often result in challenging hyper-paramettr choices and training instability if the network parameters are not properly initialized. A number of architectur--speiff initial- ization schemes have been proposed, but these schemes are not always portable to new architectures. This paper presents GradInit, an automated and architecture agnostic method for initializing neural networks. GradInit is based on simple heuristic; the norm of each network layer is adjusted so that single step of SGD or Adam with prescribed hyperparameters results in the smallest possible loss value. This adjustment is done by introducing scalar multiplier variable in front of each parameter block, and then optimizing these variables using a simple numerical scheme. GradInit accelerates the convergence and test performance of many volutional architectures, both with or without skip connections, and even without normalization layers. It also improves the stability of the original Transformer archi- tecture for machine translation, enabling training it without learning rate warmup using either Adam SGD under wide range of learning rates and momentum coefficients. Code is available at hitps:://// //gi gigghu...... wtt.coom//ca..e wwwmm/huchen03//// 3a000////radiii..
Introduction
The initialization of network parameters has a strong impact on the training stability and performance of deep neural networks. Initializations that prevent gradient explosion/vanishing in back propagation played key role early successes with feed-forward networks [ Even with cleverly designed initialization rules, complex models with many layers multiple branches instability. For example, the original Transformer model does not converge without learning rate warmup using the default initialization 4 O RoBERTa N and GPT-3
still suffer from
[8] have to tune the 89 parameter of Adam for stability when the batch size is large. Recent innovations have shown that inchitecture-speciff initializations, which carefully derived maintain stability, can promote convergence without needing normalization layers (5 [2 Unfortunately. the reliance on analyti- cally derived initializations makes it difficult realize the benefits of these methods when performing architecture search, training networks with branched or heterogeneous components, or proposing altogether new architectures.
In this work, we propose a simple method for learning the initialization of a network with any architecture. Typically, initialization schemes draw parameters independently from zero-mean distribution, with the variance of each distribution set to pre-determined values depending on the
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

dimensions of the layers 12. Rather than deriving a closed-fore expression for the these distribution parameters, our method re-scales each random weight tensor (e.g. convolution kernels) directly by a learned scalar coefficient. This small set of coefficients is optimized to make the first step of stochastic optimizer (e.g. SGD or Adam) as effective as possible at minimizing the training loss, while preventing the initial gradient norm from exploding. In addition, this process is designed to take into account the direction, step size, and stochasticity of the optimizer. Finally, after the variance has been learned for each parameter tensor, the random network parameters are re-scaled and optimization proceeds as normal. We empirically find that our methods can make the initialization fall into a smooth loss region, reduce the inter-sample gradient variance, and accelerates training.
Our proposed method, GradInit, is architecture agnostic, and works with both Adam and SGD optimizers. In the vision domain, we show it accelerates the convergence and test performance of a variety of deep architectures, from the vanilla feed-forward VGG net to ResNet, with or without Batch Normalization. It is efficient and scalable, finding good initializations using less than 1 1% of the total training time in our experiments, and it improves the initialization of ResNet-50 on ImageNet to obtain better final test accuracy. In the language domain, GradInit enables training the original Transformer model K using either Adam or SGD without learning rate warmup for machine translation, which is commonly acknowledged to be difficult 4 [3 As an extreme example of the capabilities of GradInit, we use it to initialize and train a 1202- layer ResNet that achieves significantly higher test accuracy than ResNet- 1 10, which other initialization methods have failed to achieve.
Finally, by visualizing the initial norms and gradient variances of the weights before and after GradInit is applied, we show that GradInit is a useful tool for identifying potential causes for instability at initialization, such as those imposed by normalization layers, and we summarize interesting scale patterns learned by GradInit that can be helpful for designing better initialization rules.
2
Related Work
Controlling the norms of network parameters at initialization has proven to be an effective approach for speeding up and stabilizing training. Glorot and Bengio (I] studied how the variance of features evolves with depth in feed-forward linear neural networks by assuming both activations and weight tensors are independent and identical random variables. They developed a technique in which the variance of each filter scales with its fan-in (the number of input neurons). This style of analysis was later generalized to the case of ReLU networks 2 These two analyses most effective for networks without skip connections or normalization layers. Based on the orthogonal initialization scheme [i Mishkin and Matas [15] proposed an iterative procedure to rescale the orthogonally initialized weights of each layer in feedforward networks so that the activations of that layer have unit variance. However, this method fails prevent the blowup of activations with depth for ResNets [16]. Recently, Gurbuzbalaban and Hu [77 proposed initialization schemes such that the network can provably preserve any given moment of order E (0, 2] for the output of each layer. The motivation is that the stochastic gradient updates can result in htily-taileddnes in the distribution of the network weights with a potentially infinite variance, but finite s-order moment these initialization schemes can only be applied for feed-forward neural networks.
| Again,
For more complex architectures, normalization layers 119 20 and skip connections training dynamics and improved the tfate-of-the-ar Similarly, learning rate warmup is a common trick for training large Transformers [ These methods make training tractable for some models, but do not eliminate the high initial gradient variance that destabilizes training when the network is deep I or when the normalization layers are not carefully positioned
21 stabilized
Several authors have proposed better initializations for networks with skip connections. This is often achieved by replacing the normalization layers with simpler scaling bias operations, and scaling the weight matrices in each layer so that the variance of activations does not increase with depth [ 12]. Similar analysis has been applied to self attention in Transformers [ Without removing the
normalization layers, it is possbile to stabilize the initial parameter updates by introducing carefully initialized learnable scale factors to the skip connections [6] the residual branches 2 However,
such techniques are often restricted to one specific architecture such as ResNets.
Recently, Dauphin and Schoenholz [16] proposed agnk-ago and automatic initialization method, Metalnit, for any neural network achitecture. Metalnit optimized the norms of weight tensors to minimize the "gradient quotient", which measures the effect of curvature near the initial parameters,

---

on minibatches of random Gaussian samples, However, as training data is usually accessible for most tasks of interest, it is simpler and potentially more efficient to use the training data for initialization. Metalnit also involves the gradient of a Hessian-vector product that requires computing a "gradient of the gradient" multiple times in tandem, which is very computationally intensive. Our proposed method distinguishes itself from Metalnit in the following ways: (i) Our method is more computationally efficient. Metalnit involves computing third-order derivatives, results in long computing times and high memory usage. The memory overhead of Metalnit is more of an issue for networks with normalization layers. For the relatively small-scale CIFAR-I0 problem with batch size 64, Metalnit requires three GPUs (RTX 2080Ti), while the proposed GradInit needs just one. (ii) Our method takes the stochasticity of minibatches into consideration. Metalnit uses the local curvature evaluated on single minibatch, which fails to capture the variance of the loss/gradient between two different stochastic minibatches. (iii) Our method considers the training dynamics of different optimization algorithms including the learning rate and the direction of the gradient step, and effectively handles different optimizers including SGD and Adam.
3 Method
We aim to develop an initialization scheme applicable to arbitrary network architectures. Since previous works [ 2 9 I IO [12 have shown that the initial weight norms effectively control the initial gradient norm on average, our method rescales the randomly initialized weight matrices using learnable scale factors |
Using small number of gradient descent steps on these scale factors, the proposed GradInit method chooses the initialization scalars so that the loss after the first gradient step taken by a stochastic optimizer (SGD or Adam) is as low as possible. The process of learning initialization coefficients accounts for the chosen learning rate, optimizer, and other parameters. To prevent gradient explosion, our method enforces constraint that the gradient norm is no larger than a constant 7.
Note that for scale-invariant weights, c.g., convolution kernels before BN layers, rescaling still changes their learning dynamics by changing their effective learning rate [2]] 24. Empirically,
GradInit goes beyond simply preventing exploding or vanishing gradients; it also reduces the gradient variance, making the initialization fall into smooth loss region with small gradient variance that training is fast, see discussion about Figure and comparisons in Figure 2
3.1 Efficient Learning-based Initialization via Constrained Optimization
We begin by filling all the weight matrices {W..... WM } of the network with values drawn from independent zero-mun mean Gaussian distributions, except for the scales and biases of the normalization layers (if any), which are initialized to I and 0 respectively. During the initialization process, we keep {W1.... Wa } constant, but we multiply each W; with a learnable non- -og-angteeve scale factor as (initialized to 1). After initialization, training without the learnable scale factors just as normal. We use m = {ai..... am) } to denote the set of scale factors, and O. {ar W Wi,..., Oas WM) is the set of rescaled weight matrices.
rescale the weights by the learned scale factors, and start
Let T(S: 0) = [ST Eves E(x; 0) be the average loss of the model parameterized by 0 on a minibatch of samples S, where SS is the number of samples in the minibatch. We use Is.e = LL 0) as shorthand for the gradient of 0. During standard training. this gradient is preproessed/precondtttoo by the optimization algorithm A, and then used to update the network parameters. GradInit solves the following constrained optimization problem:
minimize m subject to |19s.... 1lpA I ?:
L(S: 6m (gs.ee....)
(I)
where S and Š are two different minibatches, 17 is prescribed learning rate for the optimization
algorithm A. P.A is the norm associated with A. and y is the upper bound for the norm. For
the first gradient step, Adam uses Algs, [l.ss sign(gs.e... 25 while SGD uses A[g(S; 0m) =
For convenience, we refer to weight vectons/mattices/eenssss as "weight matrices", which includes the scale vectors of the normalization layers, the bias vectors, the weight matrices of the fully connected layers, and the convolution kernels.
3

---

SS; (m))llg((s 1l. We show how to choose and PA without tuning in Section 3.3 discuss the formulation of this problem and how solve it below.
We
3.2 Solving the Constrained Problem
The problem ① is solved using stochastic gradient descent method in which we sample new mini-batches on each iteration. Since the proposed method uses gradient updates to compute the initialization, we dub GradInit. We propose simple solver to optimize objective I in Algorithm || A key feature of our method is that is makes simple approximation: after 9s is computed on the forward pass of an iteration, we treat A[gs.om. Algs.om. on the backward pass. We make this choice to keep computing costs low, and because it is not possible to back-pru proppapt through the non-differentiable sign function for Adam.
constant and do not back-propagate through
Algorithm GradInit for learning the initialization of neural networks.
1: Input: Target optimization algorithm A and learning rate 1 for model training, initial model parameters O.. learning rate of the GradInit scales m, total iterations T, upper bound of the gradient ", lower bound for the initialization scalars 0.01.
2: mi <-
for L to T
4: Sample St from training set.
5:
Le + [8,] ExLes, f(zx; 0m.). 9: VoL.
1
6:
if I1g:llpa > then
7:
m:+1 m. --Vmmllelle..
8:
else
9:
Sample 5, from training set.
10:
Le+1 E8 (1xx; Pm. n(())))
11:
M++1 + m. -VVm Le=1
12:
Clamp m4++1 using a
To enforce the constraint in ( we test whether the constraint is satisfied after computing gS; m. If not, we take gradient descent step to minimize [g 0m) which involves computing second order derivatives. If the constraint is satisfied, then we instead compute a gradient descent step for the loss. In addition, from landing on small values during minimization and keeps the GradInit optimizer stable. In our experiments, find the only layer that ever hit this lower bound is the final FC layer some networks (see the figures in Section 4. We find this procedure converges reliably within 2000 iterations for ImageNet, and fewer than 400 iterations for CIFAR-10, taking less than 1% of the total training time on both problems. We also find it works well to set the step size to values within the range between 10-3 and 10-1 During initialization, the gradient norm constraint is satisfied for the majority of iterations. The choice of 7, PA will be discussed in Section 3.3
set lower bound 0.01 for all a; We find that this prevents scalars
Stochasticity of mini-batching. The objective in T uses two different mini-batches:; S is used to compute the gra- dient, and Š is used to compute the loss. Ideally, S and Š should be independently sampled from the training set capture the randomness of the stochastic optimizer. How- ever, when the network has large initial gradient variance, the gradients on S and Š usually differ lot, and for S, the gradient update step 0. n1 [gs.o... becomes more similar to adding random perturbations to the parameters. We find our objective less effective at accelerating conver- gence this case, as shown by the first-epoch accuracy (Ace)) in Table T
Table Accuracies on CIFAR- 10 using different overlapping ratios of S and S for GradInit.
Model
Brs
Ace
Accbeat
T9
VGG-19
0
21.9 + 4.4
94,5 # 0.1
w/o BN (20.03 M)
0.5
29,3
+ 0.6
94.7 + 0.02
28.7 + 1.0
94.5 = 0.1
On the other hand, the randomness is not captured if 5, and we find empirically that 6m can exploit the loss by increasing the gradient norm and destabilize training in this case (see Table Without excessive tuning, we find that we get more reliable behavior for different architectures when Š is mixture of 50% samples from S and 50% re-sampled training data, and use this setting by default unless otherwise stated.

---

3.3 Setting and Enforcing the Constraint
The constraint in 4 included to prevent the network from minimizing the loss in trivial way by blowing up the initial gradient. In other words, we want the optimizer to achieve small loss by choosing an effective search direction rather than by taking an extremely large step in sub-optimal direction.
Setting P.A and y through first-order approximation. We show that P.A and can be set easily with rule of thumb and without a parameter search. From the first-order approximation, expect the first gradient step to result in a change in the loss on S as following:
nlgssllsss gs 112. if A is SGD, mll.gs..mll l12, if is Adam.
L(S:0m 11llg........
-nAlgs, I gs.om.
(2)
To effectively bound the approximated change in Eq. 2 we choose lpu to be the lg and & norm for SGD and Adam respectively, so when the constraint is satisifed, the maximum change in the loss, according to our local approximation, is 1722 for SGD and ry? for Adam. We recommend setting such that 77. = 0.1 for SGD and ryy = 0.1 for Adam. According to the linear approximations, this limits the gradient magnitude so that the first step of SGD can decrease the loss by at most 0.1. This simple rule was used across all vision and language experiments.
Why constraint and not penalty? Instead of formulating GradInit as a constrained optimization, one can alternatively formulate it as minimizing the objective with gradient is the penalty strength.
penalty: minimize L(S; On - 11.A [gs.e...) + Allgs.e.. II where
m The penalized objective has two drawbacks com- pared to the constrained one in Eq. First, every gradient descent step the penalized objective involves second-order gradients due to the gra- dient regularization, while the constrained form does not need second-ordes gradients when the constraint satisfied. Second, choose good À that works well for all archi- tectures. By contrast, we set 7 by analyzing the first-order approximation mentioned above, and find the same 7 works well for different archi- tectures. The results supporting these two points are given in Table 2
Table 2: Time cost and accuracy (average runs) for running one epoch of regularization/constanne form of GradInit.
VGG-19
VGG-19 w/ BN
ResNet-110 ResNet-1110
Model
w/o BN
w/o BN w/ BN
difficult to
Time (s)
82 vs. 56 100 vs.
169 vs.
vs. 195
À = 10
32.3, 94.6 10.6, 93.1
33.7, 93.9
32.4, 95.2
À = 10-2
30.4, 94.5 10.4, 93.0
36.7. 94.1
32.6, 95.3
À =
18.2, 74.7
38.5, 95.1
30.7, 94.2
36.5, 95.3
y = 1
29.3, 94.7
47.8, 95.1
36.2, 94.6
38.2, 95.4
Experiments
We evaluate GradInit on benchmark datasets for image classification and machine translation tasks. For image classification, five different architectures arc evaluated for CIFARIO 26). and ResNet-50 is evaluated for ImageNet (27]. For machine translation, we use GradInit to find good initializations for Post-LN Transformer without any change to its original architecture IWSLT-14 14 De-En [2X] We observe that the method can remove the necessity any form of learning rate warmup for both Adam and SGD.
We conduct our experiments in PyTorch. We use the fairseq library for machine translation 29 All the experiments on CIFAR-10 and IWSLT-14 DE-EN can run with one single NVIDIA RTX 2080 Ti GPU with 1 11GB of RAM.
GradInit first initializes the weights using Kaiming initialization 2 for all the Conv and FC layers for image classification. For machine translation, we use the default Xavier initialization M. We optimize the scale factors {a,} with Adam 300 using the default momentum parameters.
4.1 Image Datasets with Various Architectures
The introduction of Batch Normalization (BN) 19 and skip connections makes it relatively easy to train common CNNs for image classification to achieve high accuracy. Despite this, we show that
5

---

when the network is very deep. the network is unstable even when both BN and skip connections are used, and GradInit can significantly improve the stability. The results on CIFAR-10 10 are given in Table and results on ImageNet are given in Table 6
4.1.1 Settings
Architectures. On CIFAR-10, we focus on the feedforward VGG net and the prevalent and powerful ResNet, with and without BN layers. For networks without BN, we use learnable biases in all layers. For ResNet, we additionally evaluate deep 1202-layer version. We give results for other architectures (Wide ResNet, DenseNet) in Appendix E due to space limits. We compare with four different methods/settings:: 1) Kaiming Initialization 22; 2) First train the network for one epoch with constant learning rate equal to the starting learning rate, labelled as "+1 epoch (Const. LR)" in Table 3 3) First train the network for one epoch with a linear warmup learning rate, labbeled "+1 epoch (Warmup)" in Table 5 4) Metalnit [66]
On ImageNet, we use the ResNet-50 model (1I.. We compare with Kaiming Initialization, Fi initialization [ and Metalnit. For the ResNet-50 without BN, we follow the architecture of FixUp for fair comparisons, but we still use the original Kaiming initialization as the starting point of GradInit.
Hyperparameters. We set to SGD and n) 0.1 (the same as the base learning rate) for GradInit in all image classification experiments. On CIFAR-10, -0, we train networks with a batch size of 128. We find Metalnit often takes
times as much memory as GradInit. We run GradInit or Metalnit for one epoch on the data, which takes less than % of the total training time. For GradInit, according to our analysis in Section 3 we fix the gradient norm constraint 7 = in all these experiments. Therefore, as in Metalnit, the only hyperparameter that needs to be tuned is the learning rate of the scale factors. We do a grid search on T in the range [10- 10 100-1, and report the results with the best average final test accuracy on 4 runs. After GradInit initialization, we use a learning rate of 0. and the cosine annealing learning rate schedule without restart 3 to train the model for 200 epochs, where the learning rate decays after each iteration and decays to O in the last iteration. Due to their high initial gradient variance (see Figure 6 we have applied gradient clipping (maximum norm is 1) to all non- non-BN networks so that they converge without GradInit under the same schedule.
On ImageNet, we train the ResNet-50 model for epochs with a total batch size of 256 on 4 GPUs. Due to the difference in the library for training and the number of GPUs used, which affects the BN statistics, our baseline top-1 accuracy of ResNet-50 (w/ BN) on ImageNet 0.79% lower than (32) We use SGD with a starting learning rate of 0.1 and decay the learning rate by 10 after the 30th and 60th epoch. We provide additional details in Appendix
4.1.2 Results and Analysis
Table 3:
First epach (Ace)) ) and best test accuracy over epochs (Accsest) for models on CIFAR-10. We report the mean and standard error of the test accuracies in 4 experiments with different random seeds. Best results in each group are in bold.
Model
VGG-19 w/o BN (20.03M)
VGG-19 w/ BN (20.04M)
ResNet-110 w/o BN (1.72M)
ResNet-110 w/ BN (1.73M)
ResNet-1202 w/ BN (19.42M)
(# Params)
Aceı
29.1
12.6 0.6
16.1 1 2.1
23.2 0.9
12.9 1 2.8
Kaiming
Accoest
94.5 0.1
9..4
94.2 = 0.1
95.0 0.2
94.4 = 0.6
+1 epoch
Acei
1.1
19.6 49.6
21.0
32,5 3.8
12.6 - 288
(Const. LR)
Aoquast
9..4
94.5
93.9 + 0.4
94.7 0.3
94.0 -
+1 epoch
Acci
1.2
19.8 2.99
48.7 1.1
1.3
(Warmup)
Anceest
94.4
94.7 0.1
9411
95.1 0.1
95.4 1 0.2
Acci
3.9
35.1 14.6 0.6
1.5
11.7
Metalnit
Accuest
9.1
94.6 0.1
94.2
94.8 0.1
95.0 5 0.5
Acci
0.6
47.8 1.8
0.8
3.9
29.0 + 1.1
GradIni
Accoest
9.17
95.1
15.1
94.6 0.1
9..4
96.2 = 0.1
GradInit further stabilizes feedforward nets with BN. BN does stabilize VGG-19 and allows training without gradient clipping, but with an average first-epoch test accuracy of only 12.57 and an average final test accuracy lower than the version without BN (see Table B) it does not seem to

---

[001 Rediiett midd
[[K]] Baa maggls
[002]
Top row: results of ResNet- 110 on CIFAR-10. Bottom row: results of ResNet- on ImageNet. Left two columns: compare the relative cross-batch gradient variance on the training set for the BN and Conv/FC layers before and after GradInit. Right two columns: weight norms before and after GradInit. Ratio between points in the same layer reflects the scale factor. Note each of the residual blocks has 2 and Conv and BN layers for the ResNet-1 100 and ResNet-50, respectively. The initial relative gradient variance are reduced for all layers except the final linear layer in both settings. The strategics are similar on two different datasets. Within each residual block, the last BN layer has the smallest scaling factors, and the scales of all Conv layers are surprisingly increased. Best viewed in color.
Figure 1:
VGG-22
Reskett110 (w) BNU
ResNet:202
Restitt-10 (wy BN, Adam!
ACCUTCY Ungrage o/ 1an00
80
3 8
70-
6D-
Kelming Credhet
Kamsng Cradin
Kaining Dredtit
Feors
Fooch
Faom
Feoch
Figure 2: Comparing the convergence of Kaiming Initialization and GradInit on CIFAR-IO, for models trained with SGD (left three) and Adam (right).
climinate the instability of Kaiming initialization. As shown in Figure Fi its initial gradient variance is still relatively high compared with GradInit. BN could magnify the gradient variance when the variance of its input features (in the forward pass) smaller than 1 (see Appendix O GradInit reduces the gradient variance by 4 orders of magnitude compared to Kaiming initialization. resulting in significantly higher test accuracy after the first epoch (47.795 vs. 12.57%), which also has an impact on the final test accuracy 13% vs. 94.41 %).. The reduction in gradient variance is achieved mainly by scaling down the weights of the final FC layer and the last BN layers, so that the variance of the activations is reduced in the forward pass. This learned behavior is consistent with the strategy of FixUp, where the final FC layer is initialized to O. Another source of gradient variance reduction is achieved by increasing the weight norms of the remaining Conv and BN layers, so that the variance of the inputs to the BN layers is increased and the gradient magnifying effect of BN is alleviated in from 204.9 to 164.8 for the Conv layers
the backward pass. This reduced the ratio ((g)) in Figure By contrast, UixUp only reduces the weight norms, which may not always be the best solution for networks with normalization layers.
Deep residual networks still need better initializations. We also gain significant improvements from GradInit for ResNet- 10 and ResNet- 1202. In ResNets, the skip connections cause the variance of activations to accumulate as the ResNet goes deeper, even for the version with BN IT.. This issue is more significant when the ResNet scales to 1202 layers, from which initialization, the first-epoch accuracy of ResNet- 1202 is quite low, and the final test accuracy is even worse than the shallower ResNet- 110, matching the observations of He et al. PI Warmup is even more effective than Metalnit at accelerating the convergence and improving the final test accuracy of ResNet- 1202, but GradInit still outperforms its final test accuracy by 0.8%, and the resulting ResNet- 1202 finally achieved higher accuracy than ResNet-110.
that with Kaiming
The learned layer-wise rescaling patterns of GradInit are even more interesting for ResNets with BN. For ResNets with BN, recall that we have two Cony layers and two BN layers each residual block. As shown in Figure [ GradInit learns to increase the weight norms of all the linear layers except for
7

---

the final FC layer, instead of decreasing as for the case without BN (see Figure 6). A more unique pattern is the collaborative behavior of the BN weights. where the second BN in each residual block is usually scaled down while the first BN is always scaled up. In deeper layers, the joint effect of these two BN weights is to downscale the activations and reduce their variance in the forward pass, with a more significant reducing effect as the layers get deeper. Intuitively, the marginal utility of adding new layer decreases with depth. Therefore, for deeper layers, GradInit learns to further downscale the residual branch, and prevents the variance from increasing too much in the forward pass. Inside each residual block, increasing the scale factors of the first BN helps to reduce the magnification effect of the second BN on the gradient; forcing the input activations to the second convolution to have variance larger than I ensures its variance after the following convolution layer does not go below 1. avoiding the magnification effect that the second BN has on the gradient variance. See Appendix for more discussions about the magnifying effect.
Table Comparing the results GradInit with fixed BN scale parameters (Fix BN) and only rescale the BN parameters (Only BN).
Kaiming
GradInit
GradInit (Fix BN)
GradInit (Only BN) Acco
Model
Acco
AcChest
Aeco
Acchest
Acco
AcCiest
Accbest
VGG-19 (w/ BN)
12.6 1 0.6 0.1 47.8 1.8
95.1
0.1
14.4 = 2.1 94.4 + 0.1 0.1 24.7 = 3.1 94.7 1 0.3 25.4 = 94.6 + 0.3
13.1 0.9 94.6 +
ResNet-110 (w/ BN) 23.2 + 0.9 95.0 1 0.2 38.2
0.9 95.4
Table 5: Comparing the results with multiplying each weight matrix with a learnable scaler (Learning Scalars) on CIFARIO. The VGG-19 model is not able to converge unless we reduce the initial learning rate to 0.01, which obtained worse final accuracy. The ResNet-110 model'' Accs was 10% for 2 of the 4 runs.
Model
Learning Scalars
GradInit
Acco
AcCbest
Acco
Accbest
VGG-19 (w/ BN. 1r10.1)
10.0 + 0.0 10.0 0.0 47.8 + 1.8 95.1 0.1
VGG-19 (w/ BN. 0=001)
50.6 + 0.8 93.4 0.1
ResNet- 10 (w/ BN)
21.5 6.9 94.7 0.1 38.2 0.9 95.4 0.1
Generalizing to Adam. Models in previous experiments are trained with SGD. We also consider the case when A is Adam and use AdamW [3 to train the ResNet-1 10 (w/ BN) model on CIFAR-10. Following (4], we use a cosine annealing learning rate schedule with initial learning rate 3 x 10-3 and weight decay 0.2. For GradInit, we set 7 25. The Aecı and Accaest of Kaiming initialization and GradInit are (36.6 + 4.7, 94.9 + 0.1) and (40.2 + 0.2, 95.3 + 0.1), respectively. We also show the per-epoch test accuracy in Figure
The importance of rescaling BN layers. The scale parameters of BN layers usually controls the variance of activations and gradients in the forward and backward passes, while the linear layers right before the BN layers are ccale-invariant Although changing the magnitudes of the scale-invariant layers affect their learning dynamics 23 24], we find it important for GradInit to rescale both BN and other linear layers, as shown Table
The importance of GradInit's objective. GradInit is designed to rescale the layers to solve the constrained optimization problem in Eq. D Simply letting the model to learn rescale the layers cannot improve the results, and sometimes further causes instability, as shown in Table 5 We hypothesize that the bad results with VGG are due gradients of the scalars and the weights. To make this alternative work, we may need to set different learning rates for the scalars and the weights, which adds to the difficulty of hyperparameter tuning. Note we do not learn the scalars when training networks initialized by GradInit.
mismatch between the scales/norms of the
Table 6: Accy/Accbest of ResNet-50 models on ImageNet. Result of Metalnit comes from Dauphin and Schoenholz [0 and we reimplemented the rest.
Kaiming
FixUp
Metalnit
GradInit
w/ BN w/o BN
14.6/75.9
19.2/76.2
18.0/75.7
-/75.4
19.2/75.8
GradInit also accelerates convergence and improves test accuracy of ResNet-50 on ImageNet, with without BN layers, despite having to
GradInit scales to ImageNet. As shown in Table 6.
8

---

use smaller batch size for GradInit than training due to our GPU memory limit. The acceleration achieved by GradInit is even more significant than FixUp, even on the network with the architecture designed for the initialization.
4.2 Training the Original Transformer Model without Warmup
For a Transformer model to converge, either an explicit or implicit learning rate warmup stage needed, especially for the original Transformer architecture. It is observed that this Post-LN architecture tends to outperform the Pre-LN model K while having higher gradient variance at initialization [. Is it believed that this high variance makes warmup stage inevitable. Previous works that removes the warmup stage often involves architectural changes, e.g., removing Layer Normalizations, since can surprisingly cause instability [ Here, we show that with a proper initialization, we can do away with the warmup stage for the original Post-LN Transformer without any modification to the architecture. Table 7 summarizes the architectural changes and best results of methods for improving the initialization of Post-LN Transformers. We compare the stability of the GradInit and Admin initialization methods without warmup in Figure
Table 7: comparison GradInit with with the results from the papers (top 4 rows), and our reimplementation Admin for training the Post-INN Transformer model on IWSLT-14 De-EN dataset. "Standard" refers to training with standard initialization and warmup.
Method
Remove LN
w skiip
Warmup
Optimizer
BLEU
Standard
RAdam
35.6
FixUp
9
Adam
34.5
T-FixUp S
Adam
35.5
Admin 6
RAdam
35.7
Admin
Adam
36.1
Admin
SGD
33.7
GradInit
Adam
36.0
GradInit
Adam
36.1
GradInit
SGD
35.6
Dataset, Architecture, & Hyperparameters.
IWSLT' 14 DE-EN [28] is German to English translation dataset that has 160k training examples.
Our Transformer model is inherited from (3) which is a Post-LN Transformer placing its Layer Normalization after the summation of the skip connection and the residual branch. It has a 512- dimensional word embedding layer and 1024 dimensions in its hidden FFN layer. We also apply GradInit to the variant from Admin [6]. where a learnable vector Wakip is element-wise multiplied with each dimension of the skip connection, but we initialize it to for GradInit. Please refer (6] for how Admin initializes these weights. Following [6. we use a linearly decaying learning rate schedule that decays from the maximum learning rate 7max to as the model trains for 100K iterations. For training with SGD, GradInit. We do a grid search on mmax for Admin and report its best result in Table 7 For training 103h for the objective of GradInit,
set the prescribed lcarning rate Thmax 0.15, and use 1 0.15, 7 for
with Adam, set 1 5 x 10 4, y discussed Section 3.3 We train the initialized model 7/uax and Bg as listed in Figure 3 We evaluate For
that 7/7 IS O(10-2) as
the BLEU score every epoch, and report the best BLEU scores throughout training for each GradInit, we set the maximum number of iterations to 780. By comparison, the warmup stage usually takes 4000 iterations, and we find that if we use 780 steps for warmup, the model does not converge with 7/max 2 3 x 10-4 For 7/max 35.4, worse than Gradlnit's s 36.0, showing the advantage of GradInit against warmup.
x 1 with 780-step warmup, the BLEU score is
Stability after removing warmup for Adam. In Figure 3 the training process becomes more unstable B2 grows larger. From the analysis of RAdam 355 this is because the variance of the gradient has a stronger impact on the adaptive learning rate when B3 is closer to Therefore, the largest Bg < 1 that maintains the performance of the trained model reflects the stability of the initial- ization. We can see GradInit results in more stable models than Admin in general, though their best performance numbers almost the same. In addition,
find Wakip can help stabilize training in extreme hyper parameter settings, e.g., 7mmax 5 x 1 and Bz = 0.995 in Figure 3 GradInit with

---

Wekip obtains a good average BLEU score of 36.0, while without Wakip only succeeded in obtaining a BLEU score 35 for one out of four experiments, resulting in an average BLEU score of 8.9.
We also find the network is unable to be trained without learning rate warmup if we just fix Wakip to its initial value given by Admin and leave the initialization of other parameters un- changed. Nevertheless, with GradInit, we do not need to modify the architecture of Post-LN Transformer to obtain the same good result as Admin For a closer look at the stabilization mechanism, show the weight norms and gra- dient variance at initialization of the original Post-LNN architecture using GradInit and Xavier initialization in Figure of the Appendix. For Xavier initialization, the gradient variance is rel- atively higher for all encoder layers, so Gra- dlnit downscales the encoder layer weights more in general. For the LN weights, GradInit only downscales the final LN of both the encoder and decoder, which reduces the variance of the encoder and decoder during the forward pass. Another strategy GradInit learns is to downscale the weights of the output projection and the FFN layers, so that the residual branch is relatively down-weighted compared with the skip connection, similar to Admin.
Adaii (w/ swal
Geailait (ey sseel
Gredlaii (w/j (Pwa))
363
10-1
31.4
35.4
19.0
11.3
353
11.3
11.3
33.0
33.4
3:-00-1
$1.b
21.
$5.9
11.7
338
51.7
13.6
319
31.7
91.1
21.0
350
359
319
300
910
56.
014
11.0
3500
35.0
13.6
361
560
0.50
0.50
Figure 3: BLEU scores for the PN Transformer without learning rate warmup using Adam on IWSLT- 14 DE-EN under different learning rates fjmax (y axis) and B8 (ar axis). Each result is averaged over 4 experiments.
Removing warmup without architectural change. Another widely observed phenomenon is that adaptive methods such as Adam seem to be much better than SGD for training Transformer- based language models K Table 7 shows that, with GradInit, we can find good initialization for the Post-LN Transformer on IWSLT-14 DE-EN that trains using SGD without learning rate warmup nor gradient clipping, and achieves performance close Adam trained using the same type of learning rate schedule. By comparison, Admin also makes the Transformer trainable with SGD, but the BLEU score is lower than the one initialized with GradInit. By comparing Figures and 10 in the Appendix. we find GradInit for Adam and SGD adopts different rescaling patterns, with the Adam version depending more on downscaling the residual branches through the FFN and output projection layers than the SGD version, and the SGD version downscaling more in the final FFN block of the decoder. This highlights the importance of considering the optimization algorithm A in GradInit, and also indicates the presence of different ways to reduce the initial gradient variance.
5 Conclusion
In this paper. we propose GradInit, gradient-based initialization scheme for any architecture. GradInit reinitializes a network by learning a scale factor for each randomly initialized parameter block of network, so that the training loss evaluated on different minibatch after one gradient step of a specific stochastic optimizer is minimized. Such design takes the stochasticity, the learning rate, and the direction of the optimizer into account, allowing us find better initializations tailored for the optimizer. The initialization learned by GradInit often decreases the gradient variance for most of the parameter blocks. We show that GradInit accelerates the convergence and improves the test performance of a variety of architectures Post-LN Transformer without any form of learning rate warmup, even for SGD, GradInit can be useful tool in the future discovery of better neural architectures that are otherwise discarded due to poor initializations. By analyzing the learned scaling coefficients and their impact on gradient variance, it can also serve guide to design better initialization schemes for complex architectures to shorten the training schedule and save energy.
image classification. It also enables training the
Acknowledgement
This project was supported by the Office of Naval Research, AFOSR MURI program, the DARPA Young Faculty Award, and the National Science Foundation Division of Mathematical Sciences. Additional support was provided by Capital One Bank and JP Morgan Chase.
10

---

References
[1] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In AISTATS, 2010.
[2] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Delving deep into rectifiers: Surpassing human-level performance on imagenet classification. In CVPR, 2015.
[3] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NeurlPS, pages 1998-6088, 2017.
[4] Ruibin Xiong, Yunchang Yang. Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang. Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In ICML, 2020.
[5] Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In ICML, 2020.
[6] Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. EMNLP, 2020.
[7] Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. arXiv preprint 1907 1990.1169.2...... 2019.
[8] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. NeurIPS, 2020.
[9] Hongyi Zhang, Yann N Dauphin, and Tengyu Ma. Fixup initialization: Residual learning without normalization. In ICLR, 2019.
[10] Soham De and Sam Smith. Batch normalization biases residual blocks towards the identity function in deep networks. NeurlPS, 2020.
[11] Andrew Brock, Soham De, and Samuel I. Smith. Characterizing signal propagation to close the performance gap in unnormalized resnets. ICLR, 2021.
[12] Andrew Brock, Soham De, Samuel Smith, and Karen Simonyan. High-performance large- scale image recognition without normalization. arXiv preprint arXiv:2/02.06171, 20211
[13] Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank J Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? NeurIPS, 2020.
[14] Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions to the nonlinear dynamics of learning deep linear neural networks. ICLR, 2014.
[15] Dmytro Mishkin and Jiri Matas. All you need a good init. ICLR, 2016.
[16] Yann N Dauphin and Samuel Schoenholz. Metainit: Initializing learning by learning to initialize.
In NewrIPS, pages 12645-12657, 2019.
[17] Mert Gurbuzbalaban and Yuanhan Hu. Fractional woomett-preserving initialization schemes for training deep neural networks. In International Conference on Artificial Intelligence and Statistics, pages 2233-2241. PMLR, 2021.
[18] Charles H Martin and Michael W Mahoney. Traditional and heavy-tailed self regularization in neural network models. arXiv preprint 17X:0.1901.08276, 2019.
[19] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In ICML, pages 448-456, 2015.
[20] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1rX1v:1607.06450, 2016.

---

[21] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In CVPR, pages 770-778, 2016.
[22] Thomas Bachlechner, Bodhisattwa Prasad Majumder, Huanru Henry Mao, Garrison W Cottrell, and Julian McAuley. Rezero is all you need: Fast convergence at large depth. arXiv preprint arXtv:2003.04887, 2020.
[23] Sanjeev Arora, Zhiyuan Li, and Kaifeng Lyu. Theoretical analysis of auto rate-tuning by batch normalization. International Conference on Learning Representations, 2019.
[24] Ruosi Wan, Zhanxing Zhu, Xiangyu Zhang, and Jian Sun. Spherical motion dynamics of deep neural networks with batch normalization and weight decay. arXiv preprint arXiv:2006.08419, 2020.
[25] Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In ICML., pages 403, 2018.
[26] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[27] Deng, W. Dong, R. Socher, L..J. Li, K. Li, and L.. Fei-Fei. ImageNet: A Large-Scale Hierarchical Image Database. In CVPR, 2009.
[28] Mauro Cettolo, Jan Niehues, Sebastian Stüker, Luisa Bentivogli, and Marcello Federico. Report on the 1 11th iwslt evaluation campaign, iwslt 2014. In /WSLT, volume 57, 2014.
[29] Myle Ott, Sergey Edunov, Alexei Baevski, Angela Fan, Sam Gross, Nathan Ng, David Grangier, and Michael Auli. fairseq: A fast, extensible toolkit for sequence modeling. In NAACL-HLT (Demonstrationss) 2019.
[30] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. ICLR, 2015.
[31] Ilya Loshchilov and Frank Hutter. Sgdr: Stochastic gradient descent with warm restarts. arXiv preprint 1608.06608 2016.
[32] Priya Goyal, Piotr Dollár, Ross Girshick, Pieter Noordhuis, Lukasz Wesolowski, Aapo Kyrola, Andrew Tulloch, Yangqing Jia, and Kaiming He. Accurate, large minibatch sgd: Training imagenet in 1 hour, arXiv preprint arXiv: 1766.02677 2017.
[33] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Conference Learning Representations, 2018.
[34] Chen Zhu, Yu Cheng, The Gan, Furong Huang, Jingjing Liu, and Tom Goldstein. Maxva: Fast adaptation of step sizes by maximizing observed variance of gradients. In Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pages 628-643. Springer, 2021.
[35] Liyunn Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. On the variance of the adaptive learning rate and beyond. ICLR, 2020.
[36] Karen Simonyan and Andrew Zisserman. Very deep convolutional networks for large-scale image recognition. arXiv preprint 11Xiv: 1556.1556, 2014.
[37] Sergey Zagoruyko and Nikos Komodakis. arXiv: 1605.070554 2016.
Wide residual networks. arXiv preprint
[38] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger. Densely connected convolutional networks. In CVPR, 2017.
[39] Terrance Vries and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint 1708. 04552, 2017.
[40] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. ICLR, 2018
[41] Jeremy Bernstein, Yu-X Yu-Xiang Wang, Kamyar Arizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In ICML, 2018.
12

---

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's
contributions and scope? [Yes]
(b) Did you describe the limitations of your work? [Yes] One limitation of our current
work we have checked whether GradInit can improve the training of models from other domains such as speech.
(c) Did you discuss any potential negative societal impacts of your work? [N/A]
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full of assumptions all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experi- mental results (either the supplemental material or as URL)? [Yes]
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes]
(c) Did you report error bars (e.g., with respect the random seed after running experi- ments multiple times)? [Yes]
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)? [Yes]
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) your work uses existing assets, did you eite the creators? [Yes]
(b) Did you mention the license of the assets? [Yes]
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent obtained from people whose data you're you're using/curating? [Yes]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/A]
13


Reviews:
[{"review_id": "QEOE7E6bxe2", "reviewer": "Reviewer_MSFd", "summary": "This paper proposed a new initialization scheme for any neural network architecture, called GradInit. It re-initializes a network by learning a scalar factor for each randomly initialized parameter block, so that the loss evaluated on a different minibatch after one update is minimized. This automated and architecture agnostic method allows the initialization to be tailored to the specific network architecture and optimization algorithms.\n\nEmpirically, this new initialization scheme often decreases the gradient variance. Furthermore, it's shown that GradInit accelerates the convergence and improves the test performance on computer vision tasks through many network architectures. In language domain, it allowed training of the Post-LN Transformer without learning rate warmup using either Adam or SGD.", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "I think this is a good empirical paper that proposed a useful automated initialization scheme that works for any optimizer and any network architecture. This method is more computationally efficient than the previous method MetaInit and also gained better performance in several tasks. GradInit also learned interesting initialization strategies, which can be useful for guiding the design of other initialization schemes. For example, in ResNets with BN, it learned to increase the weight norms of all the linear layers except for the final fully connected layer. \n\nAfter reading this paper, I think the performance of GradInit is pretty convincing, but I am still not sure why it works. In this paper, it seems the possible explanation is that it can reduce the gradient variance. But it's unclear to me how minimizing the loss after one gradient step can lead to an initialization that has smaller gradient noise. It might be good to include more discussions on the working mechanism of GradInit (why it leads to certain initialization strategies? How does it reduce gradient variance?)\n\nSome minor comments:\n1. Line 251: there is no Figure 3 in the main text.\n2. Line 277: where is Figure 5?\n3. Line 297: should be 'It is'", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations and potential negative societal impact of their work", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "HJv4u7pv09U", "reviewer": "Reviewer_DheR", "summary": "The paper proposes an architecture agnostic initialization method, GradInit, wherein neural network initialization are learnt by introducing a scalar multiplier variable in front of each parameter block. These parameters are learnt to make the first step of the optimizer (such as Adam/SGD) achieve minimum possible training loss while avoiding gradient explosions. Once the initialization has been learnt, the network is trained in normal manner using the optimizer. \n\nThe paper claims that using GradInit, the neural network initial weights fall into a smooth loss region, reduces inter-sample gradient variance and accelerates training. Empirical evaluations demonstrate the acceleration achieved by GradInit as well as the gradient norms pre and post initializations by GradInit.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "\nStrengths:\n---------------\n- The formulation to find initial scale parameters to achieve a good initialization.\n\n- Empirical evaluation shows decent improvements and speed-ups over other initialization schemes (at least in terms of the Acc@1, i.e. accuracy after initialization.).\n\nWeakness:\n---------------\n\n- Given that the initialization optimization Eq.~1 is simple, I would have expected the solution to be somewhat elegant as well, but some of the heuristics used are too optimizer specific.\n\n- Scheme require change as per the choice of optimizer and it would have been better to have an optimizer agnostic scheme.\n\n- Table~3, shows that the first epoch accuracy for the scheme \"+1 epoch(warmup)\" is mostly similar or better than the proposed scheme \"GradInit\". But the latter achieves somewhat better final accuracy. This seems somewhat unintuitive and needs an explanation.\n\n- Table~4, has missing entry for Kaiming w/o BN. Is there any reason why this entry was omitted?\n\n- The paper claims that GradInit accelerates training, yet I see no plots for convergence or training time comparison? \n\n\n\nQuestions for Authors:\n---------------\n- Eq.~1 minimizes the loss of the parameters after one step of the gradient with a constraint on the norm of the gradients. I get that you want avoid the gradient explosion and hence the constraint, but why is that you want to minimize the loss in this initialization step? Why do you expect the initialization to land you in a place where you can get small loss, its possible that this minimization lends in a region from where it would be hard for the standard optimization to recover. \n\n- What changes would be required to the heuristics and GradInit optimization process when the optimizer is not Adam or SGD?\n\n\n- I have a very fundamental question regarding the initialization scheme, i.e., in the tables you have presented, in many places, the final accuracy achieved by various initialization schemes are not that different and are very close by. Does this mean that initialization may not have that much of a significance than just putting the initial model into a decent place? I may be wrong in my assumptions and would be interested in hearing your perspective.\n\nWriting Clarity:\n---------------\nPaper is clearly written and easy to understand.\n\nIn the caption for tables, it would be better if the authors highlight the main takeaways, since its not always immediate from looking at the table, what conclusions the paper is drawing.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "N/A", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "8Gc8jDO4x0u", "reviewer": "Reviewer_vKtB", "summary": "This work introduces an initialisation strategy called GradInit.\nGradInit aims to learn the optimal standard deviation for initial weights\nto improve learning speed and stability when training the network to convergence.\nExperiments show that GradInit compares favorably to other initialisation strategies.\nMoreover, emprical evidence is presented that suggests that GradInit\nis able to replace learning rate warmup when training transformers\nand enables to build deeper ResNets than with other methods.", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Comments\n--------\n\n### Originality\n\nThis work presents an original contribution, \nbut could be a bit more precise when citing related work:\n\n - Initialisation methods do not (claim to) resolve vanishing gradients (line 20).\n   After all, Glorot et al. [1] only study the variance of the gradients,\n   rather than the magnitude of the gradients (as is normally meant with vanishing gradients).\n   Also note that He et al. [2] do not(!) take into account the backward propagation.\n - The base for most static initialisation methods would be (LeCun et al., 1998).\n   This work proposes to rescale the variance using the number of inputs to a layer.\n   Although this method clearly builds upon this idea (e.g. line 33), \n   this work is nowhere cited in the manuscript.\n - The citation on line 58 should be (LeCun et al., 1998).\n   Glorot et al. [1] also use the fan-out of the layer \n   to account for variance propagation in the backward pass.\n - The most closely related work, MetaInit [16], is properly cited\n   and the main differences are clearly explained.\n\n### Quality\n\nThere seem to be no general problems with the theoretical results.\nAll experiments were run multiple times to account for randomness.\nAlso, most experiments include error bars to give insight in the variability of the results.\nA few issues are listed below:\n\n 1. The method described by algorithm 1 does not solve the constraint problem from eq. (1).\n    It should be explicitly mentioned that this problem is solved heuristically,\n    or alternatively, the algorithm could be reformulated to guarantee that the constraint is satisfied.\n 2. From figure 1, it appears that the weights for BN layers were also initialised.\n    However, there is no method to initialise BN layers with Kaiming initialisation.\n    The most reasonable way to initialise BN layers would be\n    to initialise the shift parameter with 0 and the scale parameter with 1, \n    as proposed by Ioffe et al. [9].\n    Since I do not seem to find this information in the description,\n    I am assuming this is the used technique.\n    Since GradInit clearly does initialise the BN layer parameters differently,\n    this raises the question whether this could explain the bigger improvements for BN networks.\n    Therefore, I think it would be useful to include an ablation study\n    where BN layers are ignored by the GradInit procedure and use the default initialisation.\n    This should make clear whether GradInit provides the stability\n    or the fact that BN layer parameters are being initialised.\n    \n\n### Clarity\n\nOverall, this paper is very well written.\nHowever, a few things are not completely clear or confusing:\n\n 1. I do not understand why the SGD update is projected to a ball with radius $\\gamma$.\n    If we ignore the constraint from eq. (1), \n    I would expect the first update with SGD to be \n    $\\mathcal{A}[\\boldsymbol{g}] = \\boldsymbol{g}$.\n    If the constraint should be incorporated in the update,\n    the projection should only occur if $\\|\\boldsymbol{g}\\|_2 > \\gamma$.\n    Also, it does not seem possible to incorporate the constraint in the Adam update.\n    Since the constraint seems to be checked in the final algorithm (lines 142-145),\n    I don't see why the constraint should affect the updates.\n    Moreover, the SGD update seems to be in contradiction with eq. (2).\n 2. Let $\\hat{y} = \\alpha_{M-1} \\boldsymbol{W} f(x; \\boldsymbol{\\theta}\\_{\\boldsymbol{m}}) + \\alpha_M \\boldsymbol{b}$\n    be the output logits of the network.\n    After updating the parameters with a constant term, $\\eta \\mathcal{A}$ (cf. line 139),\n    the logits become\n    $\\hat{y}' = (\\alpha_{M-1} \\boldsymbol{W} - \\eta \\mathcal{A}_W) f(x; \\boldsymbol{\\theta}_\\boldsymbol{m} - \\eta \\mathcal{A}) + \\alpha_M \\boldsymbol{b} - \\eta \\mathcal{A}$.\n    For any commonly used loss function, we have\n    $\\frac{\\partial \\ell(x; \\boldsymbol{\\theta})}{\\partial \\boldsymbol{m}} = (\\hat{y} - y) \\frac{\\partial \\hat{y}}{\\partial \\boldsymbol{m}}$,\n    where $y$ are the labels corresponding to the inputs $x$.\n    It turns out that $\\frac{\\partial \\hat{y}}{\\partial \\boldsymbol{m}} = \\frac{\\partial \\hat{y}'}{\\partial \\boldsymbol{m}}$,\n    indicating that the scale factors of the final layer are optimised directly,\n    i.e. $L(\\tilde{S}; \\boldsymbol{\\theta}_\\boldsymbol{m})$ is minimised instead of the expression in eq. (1).\n    I can't quite figure out whether this would benefit or harm the method.\n    Therefore, I was wondering whether the authors have any insights to share in this regard.\n 3. It is not clear to me why the case for $S = \\tilde{S}$ is not included in table 1.\n    I am also not entirely sure if I understand\n    how increasing the gradient norm could be an exploit for the loss function in this case (line 164-165).\n    After all, this should already be countered by the constraint, shouldn't it?\n 4. The exact initialisation procedure for BN layers seems to be missing in the manuscript (see Q2).\n\n### Signficance\n\nImproving initialisation can save training time and other resources.\nAlthough the initialisation is more involved than static approaches,\nit does provide computational benefits over methods like MetaInit.\nExperimental results also indicate that performance can be \nslightly, but significantly improved by using GradInit instead of other methods.\nIf the issues/unclarities listed above can be tackled,\nI think this could be a useful contribution to the community.\n\nReferences\n----------\n\n - LeCun, Y., Bottou, L., Orr, G. B., & Müller, K.-R. (1998). Efficient BackProp. \n   In G. B. Orr & K.-R. Müller (Eds.), Neural Networks: Tricks of the Trade (1st ed., pp. 9–50). Springer.\n\nUpdate after rebuttal\n---\nThe issues that I raised have been properly addressed and convince me that this work is ready for publication. However, the results of the ablation study on tuning the batch-normalisation parameters does give the impression that this contributes a significant part of the improvement. In my eyes, this reduces the significance of this work, since there might be a more principled way to set these parameters (similar to standard initialisation schemes) that is comparable in performance and much cheaper than the proposed method. Considering the resolved issues and the reduced significance, I am not inclined to change my rating.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "I do not think this work could have a direct societal impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "00q5S1ylmR", "reviewer": "Reviewer_tN3S", "summary": "The paper proposes an initialization that can be applied on multiple architecture (ResNet, VGG, w/ or w/o BN, etc) and optimizers. \nTo get the initialization, they firstly introduce a learnable scalar before each weight and train the scalers only during the initialization process (T_0) with a loss. After the loss gets converged, they multiply the weights with scalars, then removing scalers from the architecture as the final initialization. Then they perform the usual training (T_1).  The loss of T_0 is composed of two parts: one makes the loss after the first step of the  T_1 as low as possible, the other enforces the gradient norm of T_1 to be within a constant constraint. ", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: This paper considers two factors of a good initialization: 1. making the loss of the first step of training as low as possible after the initialization. 2. the gradient norm should be small to guarantee a good landscape. I think the second consideration is well-known and has been used in previous initialization methods (e.g. MetaInit considers second-order info). The first consideration is new to me but not surprising. \n\nQuality: \nPros: I think the overall quality is good. \n\n1. The proposed initialization performs the baselines in multiple tasks and architecture.\n\n2. It can get rid of the warmup trick in PostLN Transformers. \n\n3. The changes of gradient variance and weight norms are different for different architectures. I think the proposed initialization is universal and should be more flexible than human-designed ones. \n\n4. The extra computation for finding the initialization is reasonable (less than 1%).\n\nCons:\n\n1. The plots for gradient variance and weight norms (like Fig. 1) are missing for ImageNet on ResNet-50. Since I doubt the pattern for CIFAR and ImageNet on ResNet are quite different based on my previous observations. It would be great if you can provide those.\n\n2. I would like to clarify the process for tuning \\tau. The selection is based on the performance after the initialization or the performance after the whole training?\n\nClarity: The paper is well-organized.\n\nSignificance: \n\nI am not sure the significance of this work, since I doubt there may be some simple substitutions. \n\n1. What about just adding a scalar before each weight and training. This can provide the network flexibility and won't result in much extra computation given the number of scalars.\n\n2. In Fig.1 (i.e. case with BN), we can observe that the change in BN weights is much more than in Linear weights. I wonder that instead of using scalers, what about just tunning the BN weights during the initialization? Given the analysis that BN weights can control the signal variance, improve the landscape and have more expressive power [1].\n \nOverall, I think this initialization is practical and researchers would like to use it in their work. And I will raise my score if my concerns get addressed. \n\n[1] Training BatchNorm and Only BatchNorm: On the Expressive Power of Random Features in CNNs.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors mention the limitation in the checklist, and maybe they would like to explore in future work.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "6"}, "extra_text": {}}]

Author rebuttal:

Meta-review: The reviewers found the paper propose a new initialization method that work better on a few datasets than prior works (metaInit, fixup, etc.). The AC agrees with these assessments, and encourage the authors to incorporate the comments of the reviewers upon acceptance. 
Decision: Accept (Poster)