# Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size

arXiv: https://arxiv.org/abs/2506.15025 · HTML: https://mlsanity.com/paper/2506.15025/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 25 references that are in the map.
- First posted on arXiv 17 Jun 2025.

## Builds on (most cited first, 18 of 18)

- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone](https://mlsanity.com/paper/2404.14219/) (2024)
- [On Lazy Training in Differentiable Programming](https://mlsanity.com/paper/1812.07956/) (2018)
- [Tensor Programs V: Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer](https://mlsanity.com/paper/2203.03466/) (2022)
- [Exponential expressivity in deep neural networks through transient chaos](https://mlsanity.com/paper/1606.05340/) (2016)
- [Deep Information Propagation](https://mlsanity.com/paper/1611.01232/) (2016)
- [Scaling Limits of Wide Neural Networks with Weight Sharing: Gaussian Process Behavior, Gradient Independence, and Neural Tangent Kernel Derivation](https://mlsanity.com/paper/1902.04760/) (2019)
- [Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies](https://mlsanity.com/paper/2407.13623/) (2024)
- [On the Impact of the Activation Function on Deep Neural Networks Training](https://mlsanity.com/paper/1902.06853/) (2019)
- [Tensor Programs VI: Feature Learning in Infinite-Depth Neural Networks](https://mlsanity.com/paper/2310.02244/) (2023)
- [How Does Critical Batch Size Scale in Pre-training?](https://mlsanity.com/paper/2410.21676/) (2024)
- [Depthwise Hyperparameter Transfer in Residual Networks: Dynamics and Scaling Limit](https://mlsanity.com/paper/2309.16620/) (2023)
- [Stable ResNet](https://mlsanity.com/paper/2010.12859/) (2020)
- [On the infinite-depth limit of finite-width neural networks](https://mlsanity.com/paper/2210.00688/) (2022)
- [Infinite-width limit of deep linear neural networks](https://mlsanity.com/paper/2211.16980/) (2022)
- [An Empirical Study of μP Learning Rate Transfer](https://mlsanity.com/paper/2404.05728/) (2024)

## Built on by (most cited first, 6 of 6)

- [Learning Rate Scaling across LoRA Ranks and Transfer to Full Finetuning](https://mlsanity.com/paper/2602.06204/) (2026)
- [A Proof of Learning Rate Transfer under μP](https://mlsanity.com/paper/2511.01734/) (2025)
- [One LR Doesn't Fit All: Heavy-Tail Guided Layerwise Learning Rates for LLMs](https://mlsanity.com/paper/2605.22297/) (2026)
- [Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate](https://mlsanity.com/paper/2605.21486/) (2026)
- [Fast Learning Rate Transfer in Shallow Linear Networks at Growing Training Horizons](https://mlsanity.com/paper/2609.35029/) (2026)
- [MuonIO: Principled Norm-Aware Descent for Embedding Tables and Language Model Heads](https://mlsanity.com/paper/2610.02705/) (2026)

## Abstract (arXiv)

Pretraining large language models is a costly process. To make this process more efficient, several methods have been proposed to optimize model architecture/parametrization and hardware use. On the parametrization side, μ P (Maximal Update Parametrization) parametrizes model weights and learning rate (LR) in a way that makes hyperparameters (HPs) transferable with width (embedding dimension): HPs can be tuned for a small model and used for larger models without additional tuning. While μP showed impressive results in practice, recent empirical studies have reported conflicting observations when applied to LLMs. One limitation of the theory behind μP is the fact that input dimension (vocabulary size in LLMs) is considered fixed when taking the width to infinity. This is unrealistic since vocabulary size is generally much larger than width in practice. In this work, we provide a theoretical analysis of the effect of vocabulary size on training dynamics, and subsequently show that as vocabulary size increases, the training dynamics interpolate between the μP regime and another regime that we call Large Vocab (LV) Regime, where optimal scaling rules are different from those predicted by μP. Our analysis reveals that in the LV regime, the optimal embedding LR to hidden LR ratio should roughly scale as Theta(√width), surprisingly close to the empirical findings previously reported in the literature, and different from the Theta(width) ratio predicted by μP. We conduct several experiments to validate our theory, and pretrain a 1B model from scratch to show the benefit of our suggested scaling rule for the embedding LR.
