Output Embedding Centering for Stable LLM Pretraining

Abstract: Pretraining of LLMs 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.

• The paper shows that OEC, via μ-centering and μ-loss, significantly stabilizes LLM pretraining by mitigating logit divergence.
• It presents a mathematical framework linking anisotropic output embeddings to divergence, revealing limitations of traditional z-loss regularization.
• Experimental results confirm that OEC reduces learning rate sensitivity, enabling efficient training even at larger learning rates.

Analysis of "Output Embedding Centering for Stable LLM Pretraining"

Introduction

The paper "Output Embedding Centering for Stable LLM Pretraining" investigates the prevalent issue of training instability in LLMs during pretraining, specifically focusing on output logit divergence. This instability, often encountered with large learning rates, results in inefficient computation and suboptimal models. The authors argue that the conventional z-loss method only addresses the symptoms and propose Output Embedding Centering (OEC) as a method that tackles the root cause of the problem, enhancing training stability across various learning rates.

Theoretical Foundations

The paper starts by thoroughly exploring the nature of instability in LLMs, especially the divergence of output logits attributed to the structure of the output embeddings. Anisotropic embeddings, where embeddings do not spread uniformly in the hidden space, are identified as a key cause of instability. The authors provide a mathematical framework showing how the non-centered nature of these embeddings contributes to unbounded growth in logits.

z-loss Limitations

Traditional z-loss is scrutinized for its limitations in addressing logit divergence. While it stabilizes training to some extent by adding a regularization term that penalizes the extremity in logits, it doesn't prevent singular negative logit divergence.

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).

New Mitigation Strategies: OEC

$\mu$-centering and $\mu$-loss

The authors introduce two strategies under the umbrella of OEC: $\mu$-centering and $\mu$-loss. $\mu$-centering is a deterministic operation that reorients the output embeddings to maintain a centered position around zero. In contrast, $\mu$-loss is a regularization method that applies a penalty based on the deviation of the embeddings from the center. The theoretical analysis demonstrates that these methods successfully suppress both the positive and negative divergence of individual logits.

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$.

Experimental Evaluation

The authors conduct experiments using modern Transformer architectures across different model sizes and learning rates to validate the efficacy of OEC. The results highlight:

• Training Stability: Both $\mu$-centering and $\mu$-loss show marked improvements in training stability, allowing convergence at larger learning rates where z-loss and baseline models fail.
• Learning Rate Sensitivity: The novel OEC techniques reduce sensitivity to learning rates, outperforming traditional methods by permitting stable training over a broader range of hyperparameters.
• Computation Efficiency: The computational overhead introduced by OEC is minimal, demonstrating practical viability.

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.

Hyperparameter Sensitivity

Further analysis reveals that $\mu$-loss requires less meticulous tuning compared to z-loss, offering a robust and user-friendly option for practitioners. This reduced sensitivity to regularization hyperparameters suggests a more adaptable integration into diverse training pipelines.

Conclusion

The study establishes a robust theoretical and empirical foundation for using OEC to mitigate training instabilities in LLMs. By targeting the core issue of embedding anisotropy, OEC significantly enhances stability and efficiency, paving the way for optimized LLM deployment in practical applications. Future work could focus on scaling the approach to even larger model architectures and exploring its effects on different types of neural network instabilities.

Overall, the paper provides a compelling argument for reevaluating current stabilization strategies in LLMs, potentially influencing future developments in AI system design.