Research map / Architectures and efficiency / RNNs and normalization
Principled Architecture-aware Scaling of Hyperparameters
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Derives architecture-aware initialization (in-degree fan-in scaling) and maximal-update learning rates (scaling with DAG path depths and CNN kernel size) for arbitrary DAG networks, and shows these re-scaled training recipes change NAS-Bench-201 accuracies and rankings.
Problem
Default initializations and learning rates are chosen independently of network architecture (depth, connectivity, kernel size), so architectures in NAS benchmarks may be under-trained and rankings unreliable.
Approach
Use an in-degree-aware initialization that preserves forward pre-activation variance through arbitrary DAG topologies, then apply the mu-P maximal-update principle to derive learning rates that scale as (sum over paths of L_p^3)^(-1/2), with an extra 1/q factor for CNN kernel size q.
Key results
- Estimated maximal learning rates correlate with grid-searched ground truth at r=0.838 for topology-scaled MLPs and r=0.856 for CNNs on CIFAR-10.
- Re-scaled training improves NAS-Bench-201 test accuracy on average by 0.44 (CIFAR-10), 1.24 (CIFAR-100), and 2.09 (ImageNet16-120).
- Network rankings change substantially under re-scaled training, especially for top-performing architectures, with Kendall-Tau correlations lower than those between random seeds.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Rescale each architecture's learning rate and initialization instead of using the fixed benchmark training protocol | NAS-Bench-201 architectures · CIFAR-10 | average test accuracy: +0.44 | Helped |
| Rescale each architecture's learning rate and initialization instead of using the fixed benchmark training protocol | NAS-Bench-201 architectures · CIFAR-100 | average test accuracy: +1.24 | Helped |
| Rescale each architecture's learning rate and initialization instead of using the fixed benchmark training protocol | NAS-Bench-201 architectures · ImageNet16-120 | average test accuracy: +2.09 | Helped |
| Use architecture-aware initialization and learning-rate scaling instead of μP heuristics | NAS-Bench-201 architectures · CIFAR-10, CIFAR-100, and ImageNet-16-120 | accuracy | Helped |
| Scale learning rates according to graph topology | MLPs · CIFAR-10 MLPs with different topologies | correlation with grid-searched maximal learning rates: r=0.838 | Helped |
| Scale learning rates according to graph topology and kernel size | CNNs · CIFAR-10 CNNs with different topologies and kernel sizes | correlation with grid-searched maximal learning rates: r=0.856 | Helped |
| Scale learning rates by depth | Feedforward MLPs · CIFAR-10 MLPs of different depths | correlation with grid-searched maximal learning rates: r=0.962 | Helped |
| Rescale learning rates and initializations for each architecture | NAS-Bench-201 architectures · Architecture rankings | ranking agreement with the default benchmark | Hurt |
| Rescale learning rates and initializations for each architecture | NAS-Bench-201 architectures · Pairwise architecture performance gaps | accuracy separation | Hurt |
| Apply depth-wise learning-rate scaling to GELU MLPs | GELU MLPs · CIFAR-10 MLPs of different depths | correlation with grid-searched maximal learning rates: r=0.920 | Helped |
| Apply topology-wise learning-rate scaling to GELU MLPs | GELU MLPs · CIFAR-10 MLPs with different graph topologies | correlation with grid-searched maximal learning rates: r=0.680 | Helped |
| Apply the learning-rate scaling principle to GELU CNNs | GELU CNNs · CIFAR-10 GELU CNNs | correlation with grid-searched maximal learning rates: r=0.949 | Helped |
Surprising findings
- Architecture rankings in NAS-Bench-201 can be substantially reordered by simply retraining with architecture-aware learning rates and initialization.
- The muP heuristic applied directly to NAS-Bench-201 architectures yields inferior results compared with the proposed scaling.
Limitations
- Width-dependent scaling of the learning rate is not addressed.
- Depth dependence is characterized only for the first step / early training, not later training.
- Depth-dependence with normalization layers is not covered.
Experimental setup
Datasets: CIFAR-10 CIFAR-100 ImageNet16-120 ImageNet
Benchmarks: NAS-Bench-201 CIFAR-10 CIFAR-100 ImageNet16-120 ImageNet
Metrics: test accuracy Pearson correlation between estimated and ground-truth maxim Kendall-Tau ranking correlation pairwise performance gaps
Baselines: muP default NAS-Bench-201 training protocol
Builds on: muP (maximal update parameterization) NAS-Bench-201 He initialization
Closest papers by meaning
- Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit (ICLR 2024 · similarity 0.88)
- Scaling Exponents Across Parameterizations and Optimizers (ICML 2024 · similarity 0.88)
- Magnitude Invariant Parametrizations Improve Hypernetwork Learning (ICLR 2024 · similarity 0.87)
- Tensor Programs VI: Feature Learning in Infinite Depth Neural Networks (ICLR 2024 · similarity 0.87)
- On the Parameterization of Second-Order Optimization Effective towards the Infinite Width (ICLR 2024 · similarity 0.87)
- $\boldsymbol{\mu}\mathbf{P^2}$: Effective Sharpness Aware Minimization Requires Layerwise Perturbation Scaling (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.