GradInit: Learning to Initialize Neural Networks for Stable and Efficient Training

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.


Reviews and author responses:

 I thank the authors for the response that has addressed my concerns. I would like to keep my positive review. 

 Thanks for the reply. I raise my score accordingly.

The paper proposes an initialization that can be applied on multiple architecture (ResNet, VGG, w/ or w/o BN, etc) and optimizers. 
To 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.   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. 

Quality: 
Pros: I think the overall quality is good. 

1. The proposed initialization performs the baselines in multiple tasks and architecture.

2. It can get rid of the warmup trick in PostLN Transformers. 

3. 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. 

4. The extra computation for finding the initialization is reasonable (less than 1%).

Cons:

1. 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.

2. 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?

Clarity: The paper is well-organized.

Significance: 

I am not sure the significance of this work, since I doubt there may be some simple substitutions. 

1. 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.

2. 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].
 
Overall, 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. 

[1] Training BatchNorm and Only BatchNorm: On the Expressive Power of Random Features in CNNs. The authors mention the limitation in the checklist, and maybe they would like to explore in future work.

This work introduces an initialisation strategy called GradInit.
GradInit aims to learn the optimal standard deviation for initial weights
to improve learning speed and stability when training the network to convergence.
Experiments show that GradInit compares favorably to other initialisation strategies.
Moreover, emprical evidence is presented that suggests that GradInit
is able to replace learning rate warmup when training transformers
and enables to build deeper ResNets than with other methods.  Comments
--------

### Originality

This work presents an original contribution, 
but could be a bit more precise when citing related work:

 - Initialisation methods do not (claim to) resolve vanishing gradients (line 20).
   After all, Glorot et al. [1] only study the variance of the gradients,
   rather than the magnitude of the gradients (as is normally meant with vanishing gradients).
   Also note that He et al. [2] do not(!) take into account the backward propagation.
 - The base for most static initialisation methods would be (LeCun et al., 1998).
   This work proposes to rescale the variance using the number of inputs to a layer.
   Although this method clearly builds upon this idea (e.g. line 33), 
   this work is nowhere cited in the manuscript.
 - The citation on line 58 should be (LeCun et al., 1998).
   Glorot et al. [1] also use the fan-out of the layer 
   to account for variance propagation in the backward pass.
 - The most closely related work, MetaInit [16], is properly cited
   and the main differences are clearly explained.

### Quality

There seem to be no general problems with the theoretical results.
All experiments were run multiple times to account for randomness.
Also, most experiments include error bars to give insight in the variability of the results.
A few issues are listed below:

 1. The method described by algorithm 1 does not solve the constraint problem from eq. (1).
    It should be explicitly mentioned that this problem is solved heuristically,
    or alternatively, the algorithm could be reformulated to guarantee that the constraint is satisfied.
 2. From figure 1, it appears that the weights for BN layers were also initialised.
    However, there is no method to initialise BN layers with Kaiming initialisation.
    The most reasonable way to initialise BN layers would be
    to initialise the shift parameter with 0 and the scale parameter with 1, 
    as proposed by Ioffe et al. [9].
    Since I do not seem to find this information in the description,
    I am assuming this is the used technique.
    Since GradInit clearly does initialise the BN layer parameters differently,
    this raises the question whether this could explain the bigger improvements for BN networks.
    Therefore, I think it would be useful to include an ablation study
    where BN layers are ignored by the GradInit procedure and use the default initialisation.
    This should make clear whether GradInit provides the stability
    or the fact that BN layer parameters are being initialised.
    

### Clarity

Overall, this paper is very well written.
However, a few things are not completely clear or confusing:

 1. I do not understand why the SGD update is projected to a ball with radius $\gamma$.
    If we ignore the constraint from eq. (1), 
    I would expect the first update with SGD to be 
    $\mathcal{A}[\boldsymbol{g}] = \boldsymbol{g}$.
    If the constraint should be incorporated in the update,
    the projection should only occur if $\|\boldsymbol{g}\|_2 > \gamma$.
    Also, it does not seem possible to incorporate the constraint in the Adam update.
    Since the constraint seems to be checked in the final algorithm (lines 142-145),
    I don't see why the constraint should affect the updates.
    Moreover, the SGD update seems to be in contradiction with eq. (2).
 2. Let $\hat{y} = \alpha_{M-1} \boldsymbol{W} f(x; \boldsymbol{\theta}\_{\boldsymbol{m}}) + \alpha_M \boldsymbol{b}$
    be the output logits of the network.
    After updating the parameters with a constant term, $\eta \mathcal{A}$ (cf. line 139),
    the logits become
    $\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}$.
    For any commonly used loss function, we have
    $\frac{\partial \ell(x; \boldsymbol{\theta})}{\partial \boldsymbol{m}} = (\hat{y} - y) \frac{\partial \hat{y}}{\partial \boldsymbol{m}}$,
    where $y$ are the labels corresponding to the inputs $x$.
    It turns out that $\frac{\partial \hat{y}}{\partial \boldsymbol{m}} = \frac{\partial \hat{y}'}{\partial \boldsymbol{m}}$,
    indicating that the scale factors of the final layer are optimised directly,
    i.e. $L(\tilde{S}; \boldsymbol{\theta}_\boldsymbol{m})$ is minimised instead of the expression in eq. (1).
    I can't quite figure out whether this would benefit or harm the method.
    Therefore, I was wondering whether the authors have any insights to share in this regard.
 3. It is not clear to me why the case for $S = \tilde{S}$ is not included in table 1.
    I am also not entirely sure if I understand
    how increasing the gradient norm could be an exploit for the loss function in this case (line 164-165).
    After all, this should already be countered by the constraint, shouldn't it?
 4. The exact initialisation procedure for BN layers seems to be missing in the manuscript (see Q2).

### Signficance

Improving initialisation can save training time and other resources.
Although the initialisation is more involved than static approaches,
it does provide computational benefits over methods like MetaInit.
Experimental results also indicate that performance can be 
slightly, but significantly improved by using GradInit instead of other methods.
If the issues/unclarities listed above can be tackled,
I think this could be a useful contribution to the community.

References
----------

 - LeCun, Y., Bottou, L., Orr, G. B., & Müller, K.-R. (1998). Efficient BackProp. 
   In G. B. Orr & K.-R. Müller (Eds.), Neural Networks: Tricks of the Trade (1st ed., pp. 9–50). Springer.

Update after rebuttal
---
The 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.
 I do not think this work could have a direct societal impact.

 Thank you for acknowledging the contributions of our work and your detailed discussions about the theories. Below we focus on addressing your concerns.

1. Further discussions about related works
    - Thanks for ​​reminding us of the pioneering work of LeCun et al. We will clearly recognize that the rescaling idea traces back to this work in our next version, and change our citations in Line 58 to mention that LeCun et al. were earlier to use fan-out to control variance propagation. We will be more rigorous when discussing other related works in our next version.
    - We agree that the solution of Glorot et al. does not explicitly solve vanishing gradients, but we think they were attempting to solve this problem and the good empirical results indicate the vanishing gradient problem is avoided in practice. In Section 4.2.1 of [1], they started by introducing the problem that just after initialization, back-propagated gradients become smaller as one moves from the output layer towards the input layer. Gradient vanishing is an extreme case of this phenomenon. Then, they derive the variance of the zero-mean Gaussian initializations so that the variance of the gradient does not vanish for feedforward networks. When the variance vanishes, the gradient also vanishes, since this means $[n\text{Var}[W]]^d$ is small, and the (zero-mean Gaussian) weight norm is also small with high probability. Therefore, justified by the good empirical results, avoiding vanishing gradient variance somehow avoids vanishing gradients in most cases.
    - He et al. [2] did consider backpropagation, in their “Backward Propagation Case” subsection of Section 2.2, though they also claimed that it does not make significant difference whether to use the fan-in or fan-out for initialization.

2. How does Algorithm 1 solve the constrained problem?
    - We will mention this explicitly in our next version. It is expensive to enforce the constraint strictly for every minibatch, due to the need of computing the second-order gradients. Therefore, we use Algorithm 1 to solve the problem heuristically. Together with other designs such as the batch mixing strategy (see Table 6 in the appendix for ablation study), in practice, we find Algorithm 1 to enforce the constraint quite reliably. The constraint usually becomes satisfied within 10% to 20% of  the total GradInit iterations and seldom violated in later steps. 

3. The importance of improving the initialization of BN layers.
    - For all the networks with BN, we initialize their scale parameters to 1 and shift parameters to 0. GradInit also starts from this.
    - Note we also have results for networks without BNs, and GradInit still improves their results (e.g., ResNet-110 w/o BN in Table 3).
    - Your intuition is right, it is important for GradInit to optimize the initial scale parameters of BN. For networks with the Conv->BN architectures, the gradient of the activations are invariant to the norm of the weights (see the third equation in Section 3.3 of the original BN paper, $\frac{\partial \text{BN}((aW)u)}{\partial u}=\frac{\partial \text{BN}(Wu)}{\partial u}$). Therefore, although the weight norms do affect the gradient of the weights, rescaling the weights alone will not affect the backpropagation of the gradient of these networks, and its effect on initial conditioning is limited. 
    - Following your suggestion, we conducted experiments with VGG-19 (w/ BN) and ResNet-110 (w/ BN). The results are very close to Kaiming initialization.
|  Model            |             Kaiming             |             GradInit            |        GradInit (Fix BN)        |
|--------------------|:-------------------------------:|:-------------------------------:|:-------------------------------:|
|                        |         Acc_0 / Acc_best        |         Acc_0 / Acc_best        |         Acc_0 / Acc_best        |
| VGG-19 (w/ BN)     | 12.6 $\pm$ 0.6 / 94.4 $\pm$ 0.1 | 47.8 $\pm$ 1.8 / 95.1 $\pm$ 0.1 | 13.1 $\pm$ 0.9 / 94.6 $\pm$ 0.1 |
| ResNet-110 (w/ BN) | 23.2 $\pm$ 0.9 / 95.0 $\pm$ 0.2 | 38.2 $\pm$ 0.9 / 95.4 $\pm$ 0.1 | 24.7 $\pm$ 3.1 / 94.7 $\pm$ 0.3 |


4. Discussions about the constraint (Clarity (1))
    - You are right: if the target optimizer is SGD, then $\mathcal{A}[\mathbf{g}]=\mathbf{g}$, and the constraint in this case is using $\ell_2$ norm.
    - If the target optimizer is Adam, then $\mathcal{A}[\mathbf{g}]=\text{sign}(\mathbf{g})$, and the constraint will use $\ell_1$ norm, according to the analysis in Sec. 3.3.
    - In Algorithm 1, after computing the gradients, we first check if the gradient norm constraint is violated (line 6). In the case of SGD, if $||\\mathbf{g}||\_2 > \\gamma$, then we compute the gradient of $||\\mathbf{g}||\_2$ w.r.t. the scale factors $\\mathbf{m}$, i.e., $\\nabla\_{\\mathbf{m}} ||\\mathbf{g}||_2$, and take a decent step on $\\mathbf{m}$ (line 7). In practice, this descent step is taken with an Adam optimizer which records the exponential moving averages of the (squared) gradients of $\\mathbf{m}$. The constraint is enforced so that the initial gradient step updates the loss within a reasonable range (Sec. 3.3). 
    - The SGD update does not contradict with Eq. 2. In this case, $L(S;\theta - \eta \mathcal{A}[\mathbf{g}]) \approx L(S;\theta) + \nabla_{\theta} L(S;\theta) [(\theta- \eta\mathbf{g}) - \theta] \approx L(S;\theta)  - \eta ||\mathbf{g}||_2^2$, from which we can see Eq. 2 holds.

5. In GradInit, is the gradient of the last layer’s scaling factor equal to the gradient of $L(\tilde{S};\theta_{\mathbf{m}})$?
    - They are not equal in most cases. From the equations, it seems that you are assuming the loss $\ell(x;
\theta)=\frac{1}{2}(\hat{y} - y)^2$. We look at the gradient of $\alpha_{M-1}$. We will use $[\cdot]_i$ to denote the $i$-th entry of a matrix or vector. The gradients of the two losses w.r.t. the scale factor on the weights of the last layer are

$$ \frac{\partial \ell(x;\theta)}{\partial \alpha_{M-1}} = (\hat{y}-y)\frac{\partial \hat{y}}{\partial \alpha_{M-1}} =  (\hat{y}-y)  \sum_{i}[Wf(x;\theta_{\mathbf{m}})]_i,$$ 

$$ \frac{\partial \ell(x;\theta-\eta \mathcal{A})}{\partial \alpha_{M-1}} =  (\hat{y}'-y)\frac{\partial \hat{y}'}{\partial \alpha_{M-1}}  = (\hat{y}’-y) \sum_{i}  [Wf(x;\theta_{\mathbf{m}} - \eta\mathcal{A})]_i,$$ 

from which we can see these two gradients are not equal in most cases, since (usually) $f(x;\theta_{\mathbf{m}} - \eta\mathcal{A})\neq f(x;\theta_{\mathbf{m}})$.

6. Result for $S=\\tilde{S}$ in Table 1? Results for $S=\tilde{S}$. Why is increasing the gradient norm undesirable?
    - We give the full comparison in the table below. The model being tested is VGG-19 (w/o BN), same as in Table 1.
    - First, if $S=\tilde{S}$, then we are optimizing a different objective, which is to minimize the loss on the same minibatch after one gradient descent step. A descent is guaranteed for this case as long as the learning rate is not too large. Empirically, we find the network tends to exploit this by increasing its gradients so that the descent step is large and the loss decreases more. 
    - Second, with Algorithm 1, if the optimization of GradInit often increases the gradient norm, then a larger portion of its descending steps on $\mathbf{m}$ is to minimize the gradient norm and GradInit becomes more similar to just controlling the gradient norm. Both factors reduce the efficacy of this alternative.

| Overlapping Ratio            | First-epoch Acc. |  Best Acc.         |
| ------------                          |     --------------       |       ------------     |
|     0                                   |  21.9$\pm$ 4.4 |   94.5 $\pm$ 0.1 | 
|     0.5                                |  29.3$\pm$ 0.6 |   94.7 $\pm$ 0.1 | 
|     1                                   |  28.7$\pm$ 0.9 |   94.5 $\pm$ 0.1 | 


 Thank you for acknowledging the practical importance of our method. We focus on addressing your concerns below.

1. Gradient variance and weight norms of ResNet-50 on ImageNet
    - Please see the plot at the anonymous link: https://ibb.co/PWf3yMZ. We plotted 48 convolutional layers (excluding the first convolution outside the residual blocks for clarity) and 1 linear layer in the plots of “Linear Weights”, as well as all 49 BN layers (first one following the first convolution layer outside the residual blocks, next 48 inside the residual blocks). Each residual block has 3 convolution and 3 BN layers, unlike the ResNet-110 on CIFAR10 where each residual block has 2 convolutional layers and 2 BN layers. For each residual block, the output dimension is higher than the intermediate dimensions, so the weight norms of the third BN/Conv layer is higher.
  - The trends in this plot is consistent with Figure 1 on CIFAR10. All the linear weights except for the final FC layer are increased. Inside each residual block in the deeper layers, the scaling factor of the first BN is preserved or slightly increased while the scales of next two BNs are reduced. 

2. The process of tuning $\tau$
    - We do a grid search on $\tau$ based on the test accuracy after training for one epoch. 

3. What about adding a scalar before each weight and training?
    - It might be a bit surprising, but we found this method to destabilize training. In the following table, we denote this method as WS (Weight Scales). We multiply each weight matrix with a learnable scaler for VGG-19 (w/ BN) and ResNet-110 (w/ BN) models on CIFAR10. We use the same hyperparameters as other experiments (a cosine learning rate schedule with initial learning rate 0.1 by default). We found the VGG-19 model not able to converge unless we reduce the initial learning rate to 0.01, which obtained worse final accuracy. The ResNet-110 model’s $Acc_0$ was 10% for 2 of the 4 runs. 
    - We hypothesize that the bad results are due to a mismatch between the scales/norms of the gradients of those scalars and the weights. To make it work, we may need to set different learning rates for the scalars and the weights, which adds to the difficulty of hyperparameter tuning.

| Model               |   VGG-19 (w/ BN, lr=0.1)   | VGG-19 (w/ BN, lr=0.01) |  ResNet-110 (w/ BN) |
| -------------------  |       -------------------              |    -------------------              |      -------------------      |
| $Acc_0$ (WS)         |       10 $\pm$ 0           |     50.6  $\pm $ 0.8          |      21.5 $\pm$ 6.9     |
| $Acc_{best}$ (WS)  |       10 $\pm$ 0           |     93.4  $\pm $ 0.1          |     94.7 $\pm$ 0.1     |
| $Acc_0$ (GradInit)  |       47.8 $\pm$ 1.8     |     -                                  |      38.2 $\pm$ 0.9     |
| $Acc_{best}$ (GradInit)  |  95.1 $\pm$ 0.1   |     -                                  |      95.4 $\pm$ 0.1     |

4. What about just tuning the BN weights during the initialization?
    - Surprisingly, this also destabilizes training. Following your suggestion, we use SGD with a constant learning rate 0.1 without weight decay and train the BN weights to minimize the cross entropy loss for classification for one epoch (not using GradInit’s objective). After that, we train all the parameters together with the same hyperparameters as other experiments. 
    - We found that when we begin to train all the parameters together, the training accuracy first drops and then increases in the first epoch, so the resulting $Acc_0$ is not significantly higher than Kaiming initialization. A plausible explanation is that in the early stage, some random weights get activated as the joint training proceeds, which negatively impacts the representation. 
- In the following table, we denote this alternative as BNW (BN Weights).
| Model               |   VGG-19 (w/ BN, lr=0.1)   |  ResNet-110 (w/ BN) |
| -------------------  |       -------------------              |      -------------------      |
| $Acc_0$ (Kaiming)         |   12.6 $\pm$ 0.6   |   23.2 $\pm$ 0.9        |
| $Acc_{best}$ (Kaiming)  |   94.4 $\pm$ 0.1   |   95.0 $\pm$ 0.2        |
| $Acc_0$ (BNW)         |        14.4 $\pm$ 2.1   |   25.4 $\pm$ 3.1        |
| $Acc_{best}$ (BNW)  |      94.4 $\pm$ 0.1    |    94.6 $\pm$ 0.3    |
| $Acc_0$ (GradInit)     |      47.8 $\pm$ 1.8      |   38.2 $\pm$ 0.9      |
| $Acc_{best}$ (GradInit)  |   95.1 $\pm$ 0.1     |   95.4 $\pm$ 0.1     |


 Thank you for acknowledging the significance of our results! We will correct the typos and include more discussions about the working mechanism of Gradinit in our future version.

1. Why does GradInit lead to certain initialization strategies? How does it reduce gradient variance?
    - Intuitively, GradInit reduces the variance by the design of its objective. When the gradient variance is too high, the gradient on minibatch $S$ is less informative of the gradient direction on a different batch $\\tilde{S}$. In this case, the gradient descent step along $\\mathbf{g}_{S}$ is less effective at reducing the loss on $\\tilde{S}$. Oh the other hand, when the variance is small, $\\mathbf{g}\_{S}$ and $\\mathbf{g}\_{\\tilde{S}}$ are more correlated, closer to the expected gradient on the dataset, and the initial gradient descent step is more effective at reducing the loss and $L(\\tilde{S},\\theta - \\eta\mathcal{A}[\mathbf{g}\_{S}])$ tend to be small. Therefore, after minimizing $L(\\tilde{S},\\theta - \\eta\mathcal{A}[\mathbf{g}\_{S}])$, GradInit finds an initialization with reasonably small variance.
    - We give post-hoc analysis of why certain initialization strategies found by GradInit can reduce the variance (e.g., Line 274 to line 288 and Appendix C). The reason why GradInit led to a certain initialization strategy may be an open problem difficult to answer. It is possible that GradInit greedly finds a strategy that reduces the gradient variance while it is optimizing the loss, but there could exist other strategies that can reduce the gradient variance as well.


 Thank you for acknowledging our results and the interesting discussions. We will edit the captions accordingly. We focus on addressing your concerns below.
1. Table 4 has missing entry for Kaiming w/o BN?
    * Without BN, using Kaiming’s initialization, ResNet-50 does not converge on ImageNet.

2. Show that GradInit accelerates training
    - We have provided $Acc_0$ as an indicator for the speed of convergence.
    - We also provide the training curves of VGG-19, ResNet-110 and ResNet-1202 on CIFAR10 through the anonymous link: https://ibb.co/HNLKT3R. We plot the the maximum accuracy achieved at each epoch ($Acc^*(n)=\\max\_{i\\le n} Acc(i)$), which reflects the training time when early stopping is used. Each curve indicate the min, max and mean of $Acc^*(n)$ from 4 experiments with different random seeds. It shows GradInit significantly accelerates the convergence.

3. Why minimizing the loss after the initial step? Why do you expect the initialization to land in a region from where you can get small loss?
    -  The intuition comes from convex optimization. For convex problems, well-conditioned loss functions (e.g., problems with very round contours) have the property that its gradient direction points towards global minima, and it can be efficiently optimized by gradient descent. However, if the loss function is poorly conditioned (with very elongated contours), the gradient directions do not point towards global minima in most cases, and the optimization is slow. Therefore, we aim at improving the conditioning of the network at initialization so that it can be optimized efficiently.
    - For the non-convex losses of neural networks we study, it is likely that weight scaling impacts the conditioning of the loss function.  A poorly scaled initialization results in a poorly conditioned optimization problem; one for which the gradient direction is not very informative and optimization is either slow or ineffective.  By carefully scaling the parameters so that the gradient direction is informative and effective at reducing the loss, we believe that we are creating a well-conditioned initialization from which gradient descent is highly efficient.
    - Our primary goal in this work was to make stochastic optimizers stable and efficient for arbitrary architectures.  As a side effect it does seem that sometimes a good initialization improves generalization at the end of training, especially for deeper networks.  We think that understanding this later effect may require tackling very deep open problems at the center of neural network theory, but we certainly think the phenomena is interesting.

4. What changes would be required to the heuristics and GradInit optimization process when the optimizer is not Adam or SGD?
    - The only change is to derive the expression for the first gradient step of the target optimizer $-\eta \mathcal{A}(\mathbf{g})$. Afterwards, we can plug it into Eq. 1 to solve the same optimization problem, or analyzing the norm constraints according to first-order approximations.

5. 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?
    - Our initial motivation was, as the reviewer says, “to put the initial model in a decent place” so that the objective is locally well-behaved at initialization and the optimization is stable. Surprisingly, we also observe that a side effect of this choice at initialization is usually (although not always) slightly better generalization at the end of the training path.  We do think there is evidence for a generalization boost after training (especially for super deep models like ResNet-1202), however  (as the reviewer points out) the smallness of our gains for some problems makes this conclusion a bit more uncertain.  We would point out, however, that many of the problems we considered are standard benchmark problems for which existing models and hyper-parameters have been fine-tuned over years of work to achieve good results, and for this reason we think competing with, and sometimes beating, existing SOTA implementations does carry some weight.


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. 

The 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.  
Strengths:
---------------
- The formulation to find initial scale parameters to achieve a good initialization.

- 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.).

Weakness:
---------------

- 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.

- Scheme require change as per the choice of optimizer and it would have been better to have an optimizer agnostic scheme.

- 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.

- Table~4, has missing entry for Kaiming w/o BN. Is there any reason why this entry was omitted?

- The paper claims that GradInit accelerates training, yet I see no plots for convergence or training time comparison? 



Questions for Authors:
---------------
- 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. 

- What changes would be required to the heuristics and GradInit optimization process when the optimizer is not Adam or SGD?


- 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.

Writing Clarity:
---------------
Paper is clearly written and easy to understand.

In 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. N/A

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.

Empirically, 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.  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. 

After 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?)

Some minor comments:
1. Line 251: there is no Figure 3 in the main text.
2. Line 277: where is Figure 5?
3. Line 297: should be 'It is' The authors have adequately addressed the limitations and potential negative societal impact of their work

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