Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

2026-05-20 · Source: Artificial Intelligence · Field: Technology & Digital — Artificial Intelligence & Machine Learning · Depth: Expert, quick

Summary

A new framework quantifies hyperparameter transfer for large language models (LLMs) using three metrics: scaling law fit quality, robustness to extrapolation errors, and asymptotic loss penalty from parameterization choice. This research investigates why Maximal Update ($μ$P) outperforms standard parameterization (SP) in learning rate transfer, particularly with AdamW. The primary finding reveals that $μ$P's advantage is largely due to its maximized learning rate for the embedding layer. In SP, the embedding layer's learning rate often bottlenecks training, causing instabilities. Aligning SP's embedding layer learning rate with $μ$P's by increasing it by a factor of width significantly stabilizes training and enhances hyperparameter transfer. Additionally, weight decay is shown to improve scaling law fits but negatively impacts extrapolation robustness in fixed token-per-parameter settings.

Key takeaway

For Machine Learning Engineers training large language models with AdamW, you should critically evaluate your embedding layer's learning rate. If using standard parameterization, increasing your embedding layer's learning rate by a factor of width, similar to Maximal Update (\$μ\$P), can significantly improve training stability and hyperparameter transfer. This adjustment can mitigate bottlenecks and enhance the scalability of your optimization efforts, making your models more robust to extrapolation.

Key insights

Maximal Update (\$μ\$P) improves LLM hyperparameter transfer by maximizing the embedding layer's learning rate, a bottleneck in standard parameterization.

Principles

• Hyperparameter transfer can be quantified by scaling law fit, extrapolation robustness, and loss penalty.
• Embedding layer learning rate is critical for LLM training stability and transfer.
• Weight decay impacts scaling law fit and extrapolation robustness differently.

Method

A framework quantifies hyperparameter transfer using three metrics: scaling law fit quality, robustness to extrapolation errors, and asymptotic loss penalty due to parameterization.

In practice

• Increase embedding layer learning rate in SP to match \$μ\$P for stability.
• Evaluate hyperparameter transfer using the proposed three metrics.
• Consider weight decay's dual impact on scaling law fits and extrapolation.

Topics

• Hyperparameter Transfer
• Large Language Models
• Maximal Update (\$μ\$P)
• Embedding Layers
• Learning Rate Optimization
• AdamW

Best for: Research Scientist, AI Engineer, NLP Engineer, AI Scientist, Machine Learning Engineer

Related on AIssential

• Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate — Maximizing the embedding layer learning rate is crucial for effective hyperparameter transfer in large language model training.
• A Quantitative Experimental Repeated Measures Study of Training Dynamics in a Small Llama Style Language Model Under a Compute-Aware Token Budget — Compute-constrained LLM training trajectories reveal non-monotonic degradation and instability, challenging endpoint-only evaluation.
• Completed Hyperparameter Transfer across Modules, Width, Depth, Batch and Duration — Complete(d)P enables effective per-module hyperparameter transfer across diverse scaling axes for large models.
• The Impact of Dimensionality on the Stability of Node Embeddings — Node embedding stability varies with dimensionality, but optimal stability does not always align with peak performance.
• Synthetic Benchmarks Overstate Forward-Forward Scaling: Real-Data Limits of Layer-Local Training — Forward-Forward learning's real-data scaling and memory efficiency are significantly overstated by synthetic benchmarks, revealing limitations compared to backpropagation.
• Learning to Recall with Transformers Beyond Orthogonal Embeddings — Transformer storage capacity scales multiplicatively with sample size, embedding dimension, and sequence length under random embeddings.

Open in AIssential →

Editorial summary, takeaway, and curation by AIssential. Original article published by Artificial Intelligence.