Output Embedding Centering for Stable LLM Pretraining

TL;DR

This work identifies output logit divergence during large-scale pretraining as arising from anisotropic, shifted output embeddings. It introduces output embedding centering (OEC) with two variants, $\mu$-centering (deterministic, hyperparameter-free) and $\mu$-loss (regularization), to center the mean output embedding around zero and bound logits. The methods consistently outperform the standard z-loss by improving training stability and reducing learning-rate sensitivity, with $\mu$-loss offering robust hyperparameter behavior. The findings offer a principled approach to stabilizing LLM pretraining by addressing the embedding geometry rather than merely treating symptoms, with the code released for reproducibility.

Abstract

Pretraining of large language models is not only expensive but also prone to certain training instabilities. A specific instability that often occurs for large learning rates at the end of training is output logit divergence. The most widely used mitigation strategy, z-loss, merely addresses the symptoms rather than the underlying cause of the problem. In this paper, we analyze the instability from the perspective of the output embeddings' geometry and identify its cause. Based on this, we propose output embedding centering (OEC) as a new mitigation strategy, and prove that it suppresses output logit divergence. OEC can be implemented in two different ways, as a deterministic operation called μ-centering, or a regularization method called μ-loss. Our experiments show that both variants outperform z-loss in terms of training stability and learning rate sensitivity. In particular, they ensure that training converges even for large learning rates when z-loss fails. Furthermore, we find that μ-loss is significantly less sensitive to regularization hyperparameter tuning than z-loss.

Figures (5)

• Figure 1: Left: z-loss from Eq. (\ref{['eq:zloss']}) without the factor $10^{-4}$. The vertical dashed line corresponds to $Z=1$, at which the z-loss reaches $0$ (indicated by the horizontal dashed line). Center: Illustration of Anisotropic Embeddings and the effect of $\mu$-centering. The purple arrow represents the mean embedding $\mu$. Right: Histogram of dot products $e_i \mathbin{\vcenter{\hbox{$\m@th\bullet$}}} \mu$ for a trained model with a standard language modeling head. The dotted, black line represents 0, while the purple and green dashed lines indicate $\| \mu \|^2 = 4.9$ and the extrema of the dot product, respectively. In the example, we have $B_- = 7.8$ and $B_+ = 4.7$, which means that the condition for reduced output logit bounds, Eq. (\ref{['eq:theorem_oec_condition']}), is fulfilled: $B_{\rm ratio} = 0.82 \leq 1$.
• Figure 2: Main results. Top: Dependency of the loss $\mathcal{L}$ on the learning rate $\eta$. Bottom: Dependency of the learning rate sensitivity LRS on the model size $N$.
• Figure 3: Additional results. The plots show the dependency of the logits mean (top left), logits standard deviation (top right), mean embedding norm (bottom left) and maximum absolute logit (bottom right) on the learning rate.
• Figure 4: Hyperparameter dependency of $\mu$-loss (left) and z-loss (right). The top plots show loss $\mathcal{L}$ vs. learning rate $\eta$, while the bottom plots show learning rate sensitivity vs. model size $N$. The results correspond to (i) and (ii) in Tab. \ref{['tab:overview_comparison']}, respectively.
• Figure 5: Histogram of $B_{\rm ratio}$ for all baseline models with a standard language modeling head.

Theorems & Definitions (14)

• Lemma 1
• proof
• Proposition 2
• proof
• Lemma 3
• proof
• Lemma 4
• proof
• Proposition 5
• proof