ドロップアウトをドロップするな：効率的なLLM学習と推論のためのレイヤースパーシティの最適化

要旨

レイヤードロップアウト（stochastic depthとも呼ばれる）は、言語および視覚Transformerの両方において、より高速な学習、より高い精度、およびゼロショット層刈り込みに対する頑健性を実現することが示されてきた。しかし、モデルとデータセットが大規模化するにつれて、ドロップアウト（特にレイヤードロップアウト）は大規模言語モデル（LLM）の事前学習レシピからほぼ姿を消している。一部の先行研究はドロップアウトが精度を低下させ得ると報告しているが、この影響を定量化した包括的な研究は存在せず、ましてや緩和した研究もない。本研究では、最先端のLLM学習においてレイヤードロップアウトを使用すべきであることを示し、学習時および学習後の利点に関するベストプラクティスとスケーリング分析を確立する。具体的には、最適な層分布、時間的スケジューリング、およびオプティマイザのハイパーパラメータを用いると、同じ学習FLOPsにおいてレイヤードロップアウトはより低い損失をもたらす。所定の学習ステップ数に対して、LLMは学習FLOPsを最大25%削減しつつ、より低いまたは同等の検証損失を達成できる。さらに、レイヤードロップアウトは、早期終了（early exit）、中間層スキップ、自己投機的復号（self-speculative decoding）など、学習後の重要な最適化を可能にし、無視できる程度の精度低下で最大1.5倍の推論高速化を実現する。271Mから8.2Bパラメータのモデル、最大160Bトークンのデータセットにわたる2400以上の学習実験を通じて、これらの知見が大規模学習体制にも確実に拡張されることを実証する。すべての事前学習実験はCerebras CS-3システム上で実施した。

English

Layer dropout (a.k.a. stochastic depth) has been shown to enable faster training, higher accuracy, and robustness to zero-shot layer pruning in both language and vision transformers. However, as models and datasets have scaled, dropout - particularly layer dropout - has largely disappeared from large language models (LLMs) pre-training recipes. While some prior work has reported that dropout can degrade accuracy, no comprehensive study has quantified, let alone mitigated, this effect. In this study, we show that layer dropout should be used in state-of-the-art LLM training, establishing best practices and scaling analysis for both training and post-training benefits. Concretely, with optimal layer distribution, time schedule, and optimizer hyperparameters, we observe that at the same training FLOPs layer dropout leads to lower loss. For a given number of training steps, LLMs can achieve lower or similar validation loss while saving upto 25% of training FLOPs. Moreover, layer dropout enables significant post-training optimizations, such as early exit, intermediate-layer skipping, and self-speculative decoding, yielding up to 1.5x inference speedup with negligible accuracy loss. Across more than 2400 training experiments, spanning models from 271M to 8.2B parameters and datasets up to 160B tokens, we demonstrate that these findings extend reliably to large-scale training regimes. All pre-training experiments were run on Cerebras CS-3 systems.