Stable Language Model Pre-training by Reducing Embedding Variability (2024)

Abstract: Stable pre-training is essential for achieving better-performing language models. However, tracking pre-training stability is impractical due to high computational costs. We study Token Embedding Variability as a simple proxy to estimate pre-training stability. We theoretically and empirically demonstrate that Multi-head Low-Rank Attention acts as a fundamental approach to reducing instability. This is supported by empirical findings on variants on GPT-2, demonstrating improved stability and lower perplexities, even at deeper layer counts.

## introduction improving large language models (llms) typically involves increasing model size, especially through greater depth @xcite @xcite . however, this approach often causes instability during pre-training, indicated by sudden spikes in loss @xcite , while stable pre-training typically leads to stronger performance under controlled training configurations @xcite . such instability can lead to catastrophic divergence or degradation, underscoring the importance of assessing pre-training stability @xcite @xcite . the conventional methods for monitoring the pre-training stability are computationally expensive @xcite , such as observing the gradient variance which needs additional o(nd) for gradient matrix g @xmath0: tev distribution for opt, pythia, llama-2, and gpt-2 reveals that as model size grows, both µ tev and σ tev decrease. this trend correlates with better model performance, as reduced noisy gradients lead to higher pre-training stability and improved performance. for a fair comparison, pythia 6.9b and 12b were excluded due to their different vocabulary sizes. the singular values of the second-order derivative of the loss with respect to model parameters @xcite @xcite . further details will be addressed in appendix a. we address both issues by dissecting the token embedding layer: we theoretically and empirically substantiate that the standard deviation of token embedding in the embedding layer, denoted token embedding variability (tev), can be a simple and efficient proxy for estimating pre-training stability in models with pre-layer normalization @xcite @xcite as it best reflects the level of gradient noise (i.e., gradient variance). we demonstrate a correlation between tev and language model performance by evaluating opt @xcite , pythia @xcite , llama-2 @xcite , and gpt-2 @xcite (figure 1 ). furthermore, we introduce factorized multi-head attention projection matrices (i.e.,, multi-head low-rank attention; mlra) as a fundamental method to mitigate pre-training instability. we empirically show that pre-training gpt-2 @xcite with mlra effectively lowers tev and achieves higher downstream performance with better pre-training stability, aligning with the theoretical analysis of tev. ## mitigating tev with factorization we propose low-rank factorized attention projection matrices (i.e., multi-head low-rank attention; mlra) as a simple way of lowering tev mean and variances, improving pre-training stability and performance. ## experiments we demonstrate the significance of tev in section 4.1 and the effectiveness of mlra on pre-training stability and performance in section 4.2. ## conclusion this paper shows that token embedding variability (tev) can be used as a simple and efficient proxy for pre-training stability, avoiding the high cost of monitoring gradient variance. theoretical analysis reveals that factorized multi-head attention projection matrices (i.e., mlra) reduce gradient explosion. empirically, mlra lowers tev mean and variance, improves stability, and outperforms gpt-2 and σreparam in reducing zero-shot perplexity, particularly in deeper models.