Principled Architecture-aware Scaling of Hyperparameters

We truly thank reviewer X2JZ's time and effort in reviewing our paper!

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

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.

Q2: I could not discern 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.

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.

We truly thank reviewer drLU's time and effort in reviewing our paper!

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

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.

Q2: The strategy of finding a base maximal learning rate for one epoch may not be meaningful for practical training cycles, where learning rates follow complex schedules.

In Section 4.3, we strictly followed the cosine learning rate schedule used in NAS-Bench-201 to train different networks. Additionally, we discuss the practical implications of finding the base maximal learning rate for one epoch at the bottom of page 7.

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

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.

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

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.

We truly thank reviewer nJvL's time and effort in reviewing our paper!

Q1: 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?

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.

Q2: Additionally, further clarification is needed regarding the assumptions made and the empirical comparisons drawn, especially across various architectural configurations. For example, it was not clearly mentioned until late, that the authors’ theory cannot apply to normalization layers.

Thank you for pointing that out! In the conclusion section, we have clarified some of our limitations, including the existence of normalization layers. We will include additional discussions in our camera-ready version.

We truly thank reviewer BP74's time and effort in reviewing our paper!

Q1: Update scale during training.

Thanks for the question! We further conduct experiments to verify the O(1) update scale for pre-activations.

Specifically, using our settings in Section 4.1 (MLP with topology scaling), we measure the change in a network’s output $z^{(L+1)}$ (Equation 7) before and after the first gradient descent step (averaged over 10 minibatches, on CIFAR-10), and compare the norm of this change across different MLP architectures. Based on our analysis, the baseline (standard initialization plus a fixed learning rate) cannot preserve this O(1) update scale when scaling up networks and will also give a large variance across different architectures.

Experiments show that $\mathbb{E}\left[\left(\Delta z_{i}^{(L+1)}\right)^2\right]$ for the baseline is 5.59 (std = 4.48), while our topology-aware initialization and learning rate yield a norm of 1.70 (std = 1.01). This result indicates that the success of our hyperparameter scaling is based on the preservation of the O(1) update scale across models.

Q2: Reference implementation ($C$ and $L_p$).

• $C$ is initially set as 2 for ReLU in standard initialization, and is further updated to be architecture-aware (see Eq. 4) in our paper.
• Calculating $L_p$ is straightforward: after representing the architecture as a computational graph, we derive its affinity matrix. Similar to Fig. 1, in this matrix, each element indicates the type of layer ($W$) used to connect two features ($x$ and $z^{(1)}, z^{(2)}, \cdots, z^{(L+1)}$); “0” indicates no connections (layers) between hidden features, “1” indicates a skip connection, and “2” indicates other parameterized layers (with ReLU activation). This affinity matrix reveals: 1) the number of paths (consecutive layers) from the input to the output (i.e., $P$); 2) the number of parameterized layers on each path (i.e., $L_p$). We will ensure comprehensive documentation in the README.md file in our released code.

Q3: The final layer weights are initialized with variance 1/n^2 instead of 1/n.

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.

Q4: Minor presentation issues.

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.

Q5: Was the learning rate also tuned on the networks being trained without this hyperparameter transfer method?

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.

Q6: Why are the maximal LRs found by experiment in Section 4.1 and 4.2 not the same as those computed by the theory?

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.

Q7: How was the experiment in Sections 4.1 and 4.2 performed?

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