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

• URL: http://arxiv.org/abs/2506.15025v1
• Date: Tue, 17 Jun 2025 23:57:30 GMT
• Title: Optimal Embedding Learning Rate in LLMs: The Effect of Vocabulary Size
• Authors: Soufiane Hayou, Liyuan Liu,
• Abstract summary: We provide a theoretical analysis of the effect of vocabulary size on training dynamics.<n>We show that as vocabulary size increases, the training dynamics emphinterpolate between the $mu$P regime and another regime.<n>Our analysis reveals that in the LV regime, the optimal embedding LR to hidden LR ratio should roughly scale as $Theta(sqrtwidth)$.
• Score: 12.916861128475272
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: 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, $\mu 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 $\mu$P showed impressive results in practice, recent empirical studies have reported conflicting observations when applied to LLMs. One limitation of the theory behind $\mu$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 \emph{interpolate between the $\mu$P regime and another regime that we call Large Vocab (LV) Regime}, where optimal scaling rules are different from those predicted by $\mu$P. Our analysis reveals that in the LV regime, the optimal embedding LR to hidden LR ratio should roughly scale as $\Theta(\sqrt{width})$, surprisingly close to the empirical findings previously reported in the literature, and different from the $\Theta(width)$ ratio predicted by $\mu$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.

Related papers

• Parallel Scaling Law for Language Models [45.799251718923614]
  We introduce the third and more inference-efficient scaling paradigm: increasing the model's parallel computation during both training and inference time.<n>We theoretically propose a new scaling law and validate it through large-scale pre-training, which shows that a model with $P$ parallel streams is similar to scaling the parameters by $O(log P)$ while showing superior inference efficiency.
  arXiv  Detail & Related papers  (2025-05-15T16:24:45Z)
• Predictable Scale: Part I -- Optimal Hyperparameter Scaling Law in Large Language Model Pretraining [58.26575378840226]
  We show that optimal learning rate follows a power-law relationship with both model parameters and data sizes, while optimal batch size scales primarily with data sizes.<n>This work is the first work that unifies different model shapes and structures, such as Mixture-of-Experts models and dense transformers.
  arXiv  Detail & Related papers  (2025-03-06T18:58:29Z)
• Improving LLM General Preference Alignment via Optimistic Online Mirror Descent [57.622821649679786]
  Reinforcement learning from human feedback (RLHF) has demonstrated remarkable effectiveness in aligning large language models (LLMs) with human preferences.<n>In this paper, we drop the Bradley-Terry (BT) model assumption and study LLM alignment under general preferences, formulated as a two-player game.<n>We show that our approach achieves an $O(T-1)$ bound on the duality gap, improving upon the previous $O(T-1/2)$ result.
  arXiv  Detail & Related papers  (2025-02-24T05:24:52Z)
• LESA: Learnable LLM Layer Scaling-Up [57.0510934286449]
  Training Large Language Models (LLMs) from scratch requires immense computational resources, making it prohibitively expensive.<n>Model scaling-up offers a promising solution by leveraging the parameters of smaller models to create larger ones.<n>We propose textbfLESA, a novel learnable method for depth scaling-up.
  arXiv  Detail & Related papers  (2025-02-19T14:58:48Z)
• RoSTE: An Efficient Quantization-Aware Supervised Fine-Tuning Approach for Large Language Models [53.571195477043496]
  We propose an algorithm named Rotated Straight-Through-Estimator (RoSTE)<n>RoSTE combines quantization-aware supervised fine-tuning (QA-SFT) with an adaptive rotation strategy to reduce activation outliers.<n>Our findings reveal that the prediction error is directly proportional to the quantization error of the converged weights, which can be effectively managed through an optimized rotation configuration.
  arXiv  Detail & Related papers  (2025-02-13T06:44:33Z)
• Scaling Laws with Vocabulary: Larger Models Deserve Larger Vocabularies [46.440917272424315]
  Research on scaling large language models (LLMs) has primarily focused on model parameters and training data size, overlooking the role of vocabulary size. We propose three complementary approaches for predicting the compute-optimal vocabulary size. Adopting our predicted optimal vocabulary size consistently improves downstream performance over commonly used vocabulary sizes.
  arXiv  Detail & Related papers  (2024-07-18T15:58:54Z)
• Fine-Tuning Language Models with Just Forward Passes [92.04219196752007]
  Fine-tuning language models (LMs) has yielded success on diverse downstream tasks, but as LMs grow in size, backpropagation requires a large amount of memory. We propose a memory-efficient zerothorder (MeZO) to operate in-place, thereby fine-tuning LMs with the same memory footprint as inference.
  arXiv  Detail & Related papers  (2023-05-27T02:28:10Z)
• nanoLM: an Affordable LLM Pre-training Benchmark via Accurate Loss Prediction across Scales [65.01417261415833]
  We present an approach to predict the pre-training loss based on our observations that Maximal Update Parametrization (muP) enables accurate fitting of scaling laws. With around 14% of the one-time pre-training cost, we can accurately forecast the loss for models up to 52B. Our goal with nanoLM is to empower researchers with limited resources to reach meaningful conclusions on large models.
  arXiv  Detail & Related papers  (2023-04-14T00:45:01Z)
• Training with Multi-Layer Embeddings for Model Reduction [0.9046327456472286]
  We introduce a multi-layer embedding training architecture that trains embeddings via a sequence of linear layers. We show that it allows reducing d by 4-8X, with a corresponding improvement in memory footprint, at given model accuracy.
  arXiv  Detail & Related papers  (2020-06-10T02:47:40Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.