Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

Summary (Overview)

• Develops a quantitative framework for evaluating hyperparameter transfer quality using three complementary metrics: loss predictability error ( $\mathcal{E}$ ), transfer robustness exponent ( $\kappa$ ), and asymptotic loss degradation ( $\mathcal{R}(\infty)$ )
• Isolates the embedding layer learning rate as the primary driver of $\mu$P's advantage over SP when training GPT-style Transformers with AdamW, contradicting the notion that the full set of $\mu$P conditions is necessary
• Demonstrates that SP with a corrected embedding learning rate (SP+Embd) matches $\mu$P across all three transfer metrics, providing practitioners a simple fix for reliable transfer
• Reveals that insufficient embedding layer learning rate causes training instabilities, not just slower convergence—a counterintuitive finding since one would expect later layers to compensate
• Shows weight decay improves loss predictability but degrades asymptotic performance in fixed-step settings and hurts transfer robustness in compute-optimal (fixed token-per-parameter) settings

Introduction and Theoretical Foundation

Background and Motivation

Training large neural networks requires careful hyperparameter tuning, which becomes prohibitively expensive as models scale to trillions of parameters. Hyperparameter transfer addresses this by finding optimal hyperparameters at small scales and extrapolating them to larger scales via scaling laws or structural parameterization choices.

Two main approaches exist:

1. Fitted scaling laws: Fit functional forms to predict how optimal learning rate $\eta^*$ scales with model/data size
2. Structural parameterization: Design the model so that training dynamics remain invariant across scales (e.g., Maximal Update Parameterization, $\mu$P)

Theoretical Foundation: Neural Network Parameterizations

Following Yang and Hu [48], networks are parameterized using four scalar exponents $\{a_l, b_l, c_l, d_l\}$ per layer $l$:

• $a_l$ controls forward pass scaling: $\mathbf{h}^{(l+1)} = n^{-a_l} W^{(l+1)} \phi(\mathbf{h}^{(l)})$
• $b_l$ scales initialization variance: $W^{(l)} \sim \mathcal{N}(0, n^{-2b_l})$
• $c_l$ scales layer-wise learning rate: $\eta^{(l)} = \eta \cdot n^{-c_l}$
• $d_l$ scales weight decay: $\lambda^{(l)} = \lambda \cdot n^{-d_l}$

Standard Parameterization (SP) requires activations at initialization to not blow up or vanish, imposing one constraint per layer. Maximal Update Parameterization ($\mu$P) imposes a stronger condition: both activations and their updates must be width-independent.

Key Differences Between SP and $\mu$P

The two parameterizations differ in four key ways (Table 1):

Theoretical Gaps

The theoretical derivation of $\mu$P makes assumptions that don't hold in practice:

1. Finite training steps in the infinite width limit (practice is the opposite)
2. Full alignment between weight updates and activations (rarely achieved)
3. Fixed dataset size (practice scales data with parameters)

Despite these violations, $\mu$P still exhibits high-quality transfer—motivating this investigation.

Methodology

Experimental Setup

• Architecture: GPT-style decoder-only Transformers
• Dataset: FineWeb-Edu [36]
• Optimizer: AdamW with Warmup-Stable-Decay (WSD) schedule (20% warmup, 60% stable, 20% decay)
• Training: Fixed 10,000 steps, batch size 1024 (1M tokens/step)
• Width scaling: $n \in [128, 2048]$, head dimension fixed at $d = 64$

Three Transfer Metrics

1. Loss Predictability Error ($\mathcal{E}$)

Models the loss landscape as a joint scaling law:

$$L(\nu; n) = L^*(\infty) + An^{-\alpha} + \frac{1}{2}Cn^{\gamma} \cdot \left(\nu - \nu^*(\infty) - Bn^{-\beta}\right)^2$$

where $\nu = \log_2 \eta$. The error is computed as:

$$\mathcal{E} = \frac{1}{N_{\nu} N_n} \sum_{i,j} \left[ L(\nu_i; n_j) - \hat{L}(\nu_i; n_j) \right]^2$$

2. Transfer Robustness Exponent ($\kappa$)

Defined as $\kappa = \alpha - 2\beta + \gamma$. The parameterization exhibits robust transfer if:

$$\kappa = \alpha - 2\beta + \gamma \leq 0$$

Negative $\kappa$ means the loss landscape flattens with width, so prediction errors diminish at scale. Positive $\kappa$ amplifies errors, degrading transfer reliability.

3. Asymptotic Loss Degradation ($\mathcal{R}(\infty)$)

$$\mathcal{R}(\infty) = L^*(\infty) - L^*_{\text{best}}(\infty) \geq 0$$

Measures the performance gap at scale relative to the best possible parameterization.

Ablation Strategy

To isolate which of the four differences between SP and $\mu$P matters most, the authors performed all 16 ablations, systematically modifying one component at a time:

• SP+Embd: SP with $\Theta(1)$ embedding learning rate
• $\mu$P-Embd: $\mu$P with $\Theta(1/n)$ embedding learning rate
• Similar modifications for attention scale, LayerNorm LR, and last-layer initialization

Empirical Validation / Results

SP vs. $\mu$P: Surprising Similarity

Despite SP exhibiting noisier loss curves (training instabilities), both parameterizations show:

• Comparable asymptotic loss gap $\mathcal{R}(\infty)$
• Converging optimal learning rates $\nu^*(n)$ to finite asymptotic values
• Negative robustness exponents (robust transfer)

Key difference: SP's predictability error $\mathcal{E}$ is ~3× larger than $\mu$P's, making SP transfer unreliable in practice despite working in principle.

The Embedding Layer Learning Rate is Critical

Figure 2 demonstrates the key finding:

• SP+Embd ($\Theta(1)$ embedding LR): Eliminates training instabilities, yields smooth $\mu$P-like curves
• $\mu$P-Embd ($\Theta(1/n)$ embedding LR): Reintroduces SP-like instabilities

Figure 3 shows SP+Embd matches $\mu$P across all three metrics, while other modifications (attention scaling, LayerNorm LR, last-layer initialization) contribute little or even hurt transfer.

When Does the Embedding Layer Matter Most?

Switching experiments reveal early training is most critical:

• $\mu$P → slow embedding at step $t_{\text{switch}}$: Causes persistent loss gap (0.1-0.2) that grows for earlier switches
• SP → fast embedding at step $t_{\text{switch}}$: Eliminates instabilities and improves performance

Freezing the embedding at initialization hurts both parameterizations, but $\mu$P much more than SP, showing later layers do not compensate for an untrained embedding.

Effect of Weight Decay

Fixed-step setting (Figure 5a-c):

• $\mathcal{R}(\infty)$ monotonically increases with $\lambda$ (from ~0.01 to ~1 at large $\lambda$)
• Small weight decay improves $\mathcal{E}$ for stable parameterizations
• At large $\lambda$, $\mathcal{E}$ converges to ~0.01 across all parameterizations

Compute-optimal setting (fixed TPP = 20):

• Most parameterizations achieve near-zero $\mathcal{R}(\infty)$
• $\kappa$ degrades sharply with increasing weight decay, converging to $\kappa \approx 0$
• The $\eta \cdot \lambda = \Theta(1)$ convention appears inadequate when training steps scale as $\Theta(n^2)$

Theoretical and Practical Implications

Why Does Embedding LR Matter So Much?

The authors argue the embedding layer sits at the network boundary with no upstream processing to compensate for poor training. The $\Theta(1)$ embedding learning rate is required for activation updates to be $\Theta(1)$—this first-layer update has only a single term contributing, so its scaling must be correct.

By contrast, the last-layer initialization constraint is weaker:

• It contributes to only one of three terms in the function update $\Delta f_1$
• Its importance depends on alignment assumptions that may not hold in practice
• Learning rate warmup compensates for the higher sharpness from larger initialization

Practical Recommendations

1. Practitioners using SP can recover $\mu$P-comparable transfer by simply setting the embedding layer learning rate to $\Theta(1)$ instead of $\Theta(1/n)$
2. Slow embedding training may be an overlooked source of training instabilities in practice
3. Weight decay scaling conventions need reconsideration in compute-optimal regimes

Generalization Beyond Transformers

In Section D, the authors show CNNs trained on CIFAR-100 exhibit the same pattern: SP with $\Theta(1)$ input-layer learning rate matches $\mu$P's transfer quality. This suggests the first/last layer special role is architecture-agnostic, arising from their boundary position in the network.

Conclusion

This work provides both a quantitative framework for evaluating hyperparameter transfer and a surprising mechanistic insight: the primary practical benefit of $\mu$P over SP for AdamW-trained Transformers comes from training the embedding layer fast enough.

Key takeaways:

1. The three metrics ($\mathcal{E}$, $\kappa$, $\mathcal{R}(\infty)$) provide complementary diagnostic lenses for transfer quality
2. The full $\mu$P prescription is excessive—SP+Embd suffices
3. Embedding layer learning rate is critical for both stability and transfer quality

Future directions:

• Extending analysis to other optimizers (SGD, Muon) whose update geometries differ
• Weight-tied embedding cases require careful treatment (naive SP needs $1/n$ output multiplier)
• Finding correct weight decay scaling in compute-optimal regimes where training horizon scales as $\Theta(n^2)$

Limitations: Experiments limited to decoder-only Transformers with fixed depth, scaled to 1B parameters, trained with AdamW on a single dataset (FineWeb-Edu), with single random seed per configuration.