🤖 AI Summary
This work addresses the instability commonly observed in large language model pretraining under high learning rates, which often manifests as divergent output logits. The study uncovers the geometric origin of this phenomenon: the output embedding matrix drifting away from the origin. To resolve this, the authors propose Output Embedding Centering (OEC), a novel strategy that enforces the output embeddings to remain centered at the origin either through μ-centering—a deterministic recentering operation—or μ-loss, a regularization term incorporated into the training objective. By directly mitigating the root cause of logit divergence, OEC substantially enhances training stability and convergence. Empirical results demonstrate that OEC outperforms the existing z-loss under high learning rates, with μ-loss exhibiting greater robustness to hyperparameter choices.
📝 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 {\mu}-centering, or a regularization method called {\mu}-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 {\mu}-loss is significantly less sensitive to regularization hyperparameter tuning than z-loss.