Paper: Principled Architecture-aware Scaling of Hyperparameters (ICLR-2024)

Reviewer comment (Experiments):
The rankings of architectures may be influenced by various factors during training, including random seeds. It raises the question of how the authors determined that their revised learning rates and initializations would have the most significant impact on their observation of updated rankings.

Author response:
Thank you for the question! We agree that numerous factors can influence architecture rankings.

To delve deeper, we analyze how the random seed affects these rankings. Please kindly refer to Appendix E.4 in our updated paper.

Specifically, NAS-Bench-201 reports network performance trained with three different random seeds (seeds = [777, 888, 999]). On CIFAR-100, we created a plot similar to the middle column of Figure 4, but it also displays pairwise ranking correlations among those three random seeds. This plot is included in Appendix E.4. In summary, although we observe different ranking correlations between seeds, they are consistently higher (i.e., correlations are more consistent) than those produced by our method.
This confirms that changes in network rankings by our architecture-aware hyperparameters are *meaningful*, which can train networks to better performance.

Reviewer comment (Evaluation):
NAS-Bench typically selects learning rates closer to those utilized by "good architectures" that are in practical use, reflecting an intentional bias towards favoring such architectures over "bad architectures." However, in Section 4.3, there is no discernible evidence that the new learning rate/initialization approach improves upon the performance of the top-performing architectures, while it appears to only boost many of the "bad architectures." Given this, it raises the question of why this new learning rate/initialization strategy is relevant for practical NAS if it cannot enhance the performance of the best-found architectures.

Author response:
Thank you! We would like to point out that in the left column of Figure 4, on both CIFAR-100 and ImageNet16-120, the improvement in the high-accuracy regime is either better than or as good as that in the low-accuracy regime.

Reviewer comment (Theory):
The present paper introduces a novel Directed Acyclic Graph (DAG) tool to derive "architecture-aware" hyperparameters. However, a fundamental question remains: What are the specific failure modes of the μP when it encounters intricate network topologies? Is the limitation attributed to a violation of Desiderata L.1, or are there other contributing factors at play?

Author response:
Thanks for your question! In the second paragraph of Section 2.2, we specifically discussed two core differences between $\mu$P and our method (and they are also reasons for $\mu$P’s pitfall in our scenario): 1) the non-trivial dependence of the learning rate on network depth, and 2) the complicated graph topologies of architectures. We will merge Figure 8 in Appendix E.3 into Figure 4.

Reviewer comment (Writing):
The clarity of the main body could be improved a lot, via providing a concise summary of main derivations, potentially by reducing the repetitive criticism of architecture search.

Author response:
We have summarized our main derivations at the beginning of Section 3. We will further reduce the emphasis on the architecture search part and will update our draft accordingly in our camera ready.

Reviewer comment (Evaluation):
The empirical results are also not entirely convincing. Figure 2 exhibits a weak correlation with questionable linearity and Figure 3 has a limited range of learning rates. The proposed improvements in Figure 4 seem to be primarily in lower accuracy regimes, raising doubts about absolute improvement on the top-performer architectures (which are of the most interest).

Author response:
We respectfully disagree that a 0.838 correlation in Figure 2 is weak. Figure 3 spans an even wider range of learning rates than Figure 2. In the left column of Figure 4, for both CIFAR-100 and ImageNet16-120, the improvement in the high-accuracy regime is as good as, or better than, that in the low-accuracy regime.

Reviewer comment (Novelty):
The discussion of prior art is limited, particularly regarding weight initialization and learning rates, and it's unclear how the proposed method advances over existing insights in the literature.

Author response:
In the second paragraph of Section 2.2, we provide a detailed discussion of closely related works.

Our work advances over existing insights: Previous studies have observed and developed heuristics about the relationship between learning rates and network structures. However, our core advantage over prior work lies in precisely characterizing the non-trivial dependence of learning rates on network depth, DAG topologies, and convolutional kernel size. This goes beyond previous heuristics and has practical implications.

While our underlying principle may seem to echo patterns from earlier studies, such as using smaller learning rates for larger models, we provide the exact scale of this dependence (-3/2 for depth, and the summation of depths across the network’s graph structure). As confirmed by reviewers nJvL and X2JZ, our work offers a deeper understanding of how to determine initial learning rates for specific architectures and establishes a link between learning rate scaling and layer depth for the first time.

Reviewer comment (Theory):
It is stated "the final layer weights are initialized with variance $1/n^2$ instead of $1/n$" on page 5, and I know this is from $\mu$P but it is never stated in this paper why.

Author response:
We adhere to the variance of $1/n^2$ for the output layer, primarily to align with the desiderata outlined in Section J.2.1 of the referenced work ($\mu$P). This approach is chosen to preserve the O(1) size of pre-activations. It also accommodates the observation that "the input layer is updated much more slowly than the output layer" as mentioned in their Section 5.

Reviewer comment (Theory):
Why is the MSE loss introduced in equation 6? The experiments are mostly classification results using a cross-entropy loss.

Author response:
Thanks for these suggestions!
* We derive our principle using the MSE loss, mainly for its clarity in analyzing the changes of the model’s output $z^{(L+1)}$ (Eq. 21).
* $L_p$ is defined in Sec. 3.3 (after Eq. 7) and is also depicted in Fig. 1.

Reviewer comment (Writing):
Equation 8 contains $L_p$ but $L_p$ is only defined in the Appendix, it should be defined near the equation.

Author response:
Thanks for these suggestions!
* We derive our principle using the MSE loss, mainly for its clarity in analyzing the changes of the model’s output $z^{(L+1)}$ (Eq. 21).
* $L_p$ is defined in Sec. 3.3 (after Eq. 7) and is also depicted in Fig. 1.

Reviewer comment (Experiments):
For each network architecture sampled in the NAS-Bench experiments, the network learning rate is first tuned at small scale. Was the learning rate also tuned on the networks being trained without this hyperparameter transfer method? This may already be in the experiments section or the Appendices and I may have missed it.

Author response:
The learning rate used in the baseline (x-axis in Fig. 4 left column) was not tuned. It was originally pre-determined in NAS-Bench-201.

Reviewer comment (Theory):
Why are the maximal LRs found by experiment in Section 4.1 and 4.2 not the same as those computed by the theory? Is there any way they could be brought closer together?

Author response:
Maximal learning rates (LRs) found through experiments are not exactly the same as those predicted by our theory, mainly due to the more complicated training dynamics that are not fully captured by our analysis.
Bring them closer together requires detailed analysis of the gradient descent beyond the first step.
Despite this, our work demonstrates that analyzing only the first gradient descent step can still yield highly accurate learning rate predictions, with correlations exceeding 0.8.

Reviewer comment (Experiments):
How was the experiment in Sections 4.1 and 4.2 performed? I guess you find the optimal learning rate at a small scale on a set of networks, then scale the network up using the method described in the paper to scale the initalization and learning rates, then perform a grid search to find the optimal learning rates to compare against?

Author response:
Yes, our experiment pipeline is exactly the same as those three steps we stated at the beginning of Section 4.