Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: Authors study the problem of finding (learning) the best neural network initialization, and propose GradCosine - a measure of fitness for neural network initialization, based the similarity between individual sample gradients. This measure can then be optimized using and iterative procedure denoted as NIO (neural initialization optimization). Authors explain how GradCosine relates to model training and show a the relation between this quantity and the density of sample-wise local optima ([52] in the paper).
From a theoretical point of view, the paper demonstrates that minimizing GradCosine results in favourable bounds on the training loss. For practical applications, authors offer a sub-batch version of GradCosine that can be computed more efficiently.
strengths_and_weaknesses: To the best of my understanding, the theoretical analysis appears sound, based on reasonable assumptions. [That said, I am no expert in the area of learned initialization.] The main experiments look reasonable, though I would recommend making some extra comparisons.

1. in Table 5, you evaluate how NIO trains SWIN transformer without warmup. At least one baseline (GradInit[53]) also claims to work in that setting. Perhaps it would be best to compare NIO against that.

2. Table 4 does not report standard deviations, while all prior experiments do. It might be useful to include them - or explain why they are missing.


On an unrelated note, I must applaud authors for supplying a Dockerfile in their supplementary code. Publishing the corresponding docker container will make it easier for future researchers to reproduce this work and build on it, even if the required libraries break compatibilities.


### Typos / nitpicking

> L215  making the sub-batches overlapped with each other **stables** the optimization

perhaps a typo? "stables" -> "stabilizes"


> L250 Each model is trained **for** four times with different seeds. 

consider removing "for"


> NIO is able to produce a better initialization that benefits model performance agnostic of architecture and dataset.

[nit] this is concluded at an early section, where the only evidence is training cnn-only models on CIFAR-10/100. As a weak suggestion, I would recommend making this conclusion later, once you demonstrate NIO performance for transformers and ImageNet.


questions: Note: the question / suggestion below bears little significance. If authors are constrained by response time, please ignore it.

To the best of my understanding, even batched GradCosine currently requires more GPU memory than the objectives of MetaInit or GradInit (per model parameter). This could be a limitation when applying NIO for large models, e.g. transformer language models.
I wonder if it is possible to reduce the memory usage algorithmically.

Suppose that most of that memory is contributed by gradients w.r.t. weight matrices in linear layers, whether in convolutions or transformer projections.

For linear layers, one can observe that __one-sample gradients w.r.t. weight matrices are low-rank__.
In the simplest case, an MLP linear layer with batch size 1 will always receive rank-1 gradient - due to the fact that the gradient w.r.t. weight matrix is a product of two matrices of shape (1 x in_features) and (out_features x 1), where both "1"s correspond to batch size.
For convolution and attention layers, this will similarly result in low-rank gradients.

This trick is described in more detail in [1], though i believe that it was invented prior to that work.

Put simply, if your gradient w.r.t. weight matrix was computed with batch size 1, you can compute the pairwise products for GradCosine without computing the gradients w.r.t. weight, using only the gradients w.r.t. activations. Furthermore, you could avoid storing activations for all layers by re-computing them on the fly[2].


[1] https://arxiv.org/abs/2110.11309

[2] https://arxiv.org/abs/1604.06174v2



limitations: To the best of my knowledge, authors have sufficiently addressed the limitations of their work.
As for the societal impact, this specific paper contains fundamental research in deep learning, thus it is hard to foresee its societal impact.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: The paper first introduces a quantity, the cosine similarity of sample-wise local optima to evaluate the model performance at the initialization. They theoretically proved that their proposed quantity is the upper bound of both the training and generalization error under certain assumptions. Based on this theoretical finding, they approximate the sample-wise optimum with the first-order approximation to make the quantity differentiable and tractable. As a result, they simplify the upper bound quantity and achieve the initialization by maximizing the quantity with the gradient-based method. Their empirical results show that they can achieve better performance on various datasets and network structures compared to other initialization methods.
strengths_and_weaknesses: Strengths:
The paper is well organized and easy to follow. The theoretical analysis seems correct and motivated. Their empirical results are quite good, especially on CIFAR datasets. 

Weaknesses:
It will be better if we can have some experiments to show that the proposed cosine similarity of sample-wise local optima is useful or better than the previous methods. Since the initialization method is motivated by this quantity, it’s better to make the usefulness of this quantity clear. 

The citation in Appendix A for the proof of Lemma 1 seems wrong. 

questions: 1. Can we have some experiments to show that the proposed cosine similarity of sample-wise local optima is more suitable for evaluating the initialization quality?

2. Does the initialization framework work for multiple optimization methods? For example, like sharpness-aware-minimization?

3. It seems the method is dataset-dependent, will we have the same benefits if we initialize the parameter of the networks with a different dataset? For example, training on cifar-10 but initialized with cifar100 or ImageNet?

4. Does this only work for initialization? Can we repeatedly use it during the training and will that be better?

limitations: The authors fairly addressed the limitations and potential negative societal impact of their work.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: This work proposes a new initialization method for neural network. The authors first use Fig. 1 to illustrate drawbacks of the sample-wise local optima density $\Phi_{S,l}$ that adopts Manhattan distance between the pair-wise local optima, and lead to a Cosine similarity of sample-wise local optima $\Theta_{S,l}$. Then the authors introduce to approximate local optimum by one step gradient descent and approximate $\Theta_{S,l}$ by Eq. 4 (GradCosine). Finally, the initialized weights $\theta_0$ are obtained by maximizing GradCosine and GradNorm. Extensive experiments are conducted to verify the efficacy of this approach, showing that GradCosine surpasses MetaInit and GradInit on multiple datasets in both accuracy and speed.
strengths_and_weaknesses: Strengths:
1) The writing is logical and fluent.
2) The experimental results are good.
3) The motivation is interesting. Nevertheless, I have one question about it, please see the first point in Question.

Weakness (Question):
1) In Line 28, the authors claim that “these methods merely use the first order training dynamic as the main optimization objective”. However, MetaInit[8] and GradInit[53] both adopts gradient descent method.
2) It seems that the theoretical analysis in Sec. 3.2 cannot show that $\Theta_{S,l}$ has a tighter bound than $\Phi_{S,l}$. Then I think Sec. 3.2 is a little redundant, which can be put into the supplementary.
3) Typo: ‘GradCoisne’ in Line 6.

questions: 1. The authors claim that (Line 114-115): The optimization path from the initialization to the local optima of the two samples in Fig. 1(d) are more consistent compared to that in Fig. 1(c). However, I notice that $\theta_0$ in Fig. 1(c) is closer to $\theta_1^*, \theta_2^*$ than $\theta_0$ in Fig.1 (d), resulting in a larger angle $\angle(\theta_1^*-\theta_0, \theta_2^*-\theta_0)$. My point is that the angle cannot reflect the quality of optimization path since an initial point that is closer to the optimal point might as a larger angle.
2. I have one concern on the approximation in Eq. 3. For initialized weights $\theta_0$, such an approximation seems not reasonable. I hope the authors could discuss the rationality of it.

limitations: The limitations are addressed.
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: This paper works on the initialization of neural network parameter based on theoretically inspired optimization algorithm. It is to substitute the previous network initialization algorithm, such as Kaiming's method, [52], etc. The proposed approach is closely related to [52] by introducing the cosine similarity of sample-wise gradient for optimizing the initial network parameter. The proposed approach is applied to Resnet, DenseNet, WideResNet, and transformer. The results show that the proposed initialization method achieved better final network training results than the Kaiming's method, and the other compared learning-based methods.  
strengths_and_weaknesses: Strength:

1. The paper theoretically dig into the network initialization based on the difference of network parameter initialization and sample-wise local minima, and derive a quantity to approximate the upper-bound of training and test error bounds. The paper further approximate and minimize  the bound with constraint on gradient magnitude.  The proposed method is reasonble and inspired by theoretical analysis.

2. The proposed method is evaluated on network training and show improved results with the proposed network parameter initialization method.

Weakness:

My major questions and concerns are on the novelty over [52] and insufficient evaluation details. 

1. The  theoretical analysis of this work is similar to the analysis of [52], though the conclusion of this paper is different measuring by the cosine similarity of sample-wise gradients. 

2. The proposed approach is subject to an optimization problem of (6) adaptive to mini-batch based implementation. For feasibility in implementation, it makes some assumption and introduced the hyper-parameters, e.g., \lambda. How is the performence affected by the hyper-parameters of the optimization algorithm, including, e.g., \lambda, and number of iterations, etc ?

3.   Due to the randomness, e.g., the mini-batch, the performance should be reported by running a same task in several times, and the mean the variance of the performance should be more insightful to compare the different methods. 

4. Because of the introduced additional computational cost in optimization, the approach is reported by running 100 iterations. How about the comparisons (speed and performance) if all the compared methods are aligned in number of running iterations for initialization?
questions: Please see my questions on the novelty, experimental justifications, etc., in the weakness. 
limitations: The paper clearly state the limitations of this approach.
review_title: null

Meta-review: metareview: The paper introduces a new procedure to initialize the optimisation in training process of DNN models, including the recent ViT architecture. All the reviewers recommend acceptance and appreciate the promising empirical results backed by the strong theoretical foundations. AC recommends acceptance as well.
confidence: Less certain
Decision: Accept