Spike No More: Stabilizing the Pre-training of Large Language Models

Sho Takase, Shun Kiyono, Sosuke Kobayashi, Jun Suzuki

Keywords: stable training, llm, pre-training

TL;DR: Theoretical analysis to prevent loss spikes during LLM pre-training

Abstract:

Loss spikes often occur during pre-training of large language models.
The spikes degrade the performance of large language models and sometimes ruin the pre-training.
Since the pre-training needs a vast computational budget, we should avoid such spikes.
Based on the assumption that the loss spike is caused by the sudden growth of the gradient norm, we explore factors to keep the gradient norm small through an analysis of the spectral norms of the Jacobian matrices for the sub-layers.
Our findings suggest that stabilizing the pre-training process requires two conditions: small sub-layers and large shortcut.
We conduct various experiments to empirically verify our theoretical analyses.
Experimental results demonstrate that methods satisfying the conditions effectively prevent loss spikes during pre-training.

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the COLM Code of Ethics on https://colmweb.org/CoE.html

Author Guide: I certify that this submission complies with the submission instructions as described on https://colmweb.org/AuthorGuide.html

Submission Number: 1395