シンプルでスケーラブルな大規模言語モデルの継続的事前学習戦略

Simple and Scalable Strategies to Continually Pre-train Large Language Models

March 13, 2024

著者: Adam Ibrahim, Benjamin Thérien, Kshitij Gupta, Mats L. Richter, Quentin Anthony, Timothée Lesort, Eugene Belilovsky, Irina Rish

cs.AI

要旨

大規模言語モデル（LLM）は通常、数十億のトークンで事前学習されますが、新しいデータが利用可能になると、そのプロセスを最初からやり直すことが一般的です。より効率的な解決策は、これらのモデルを継続的に事前学習することであり、再学習と比較して大幅な計算リソースを節約できます。しかし、新しいデータによって引き起こされる分布シフトは、通常、以前のデータでの性能低下や新しいデータへの適応不良を招きます。本研究では、学習率（LR）の再ウォーミング、LRの再減衰、および以前のデータのリプレイを組み合わせたシンプルでスケーラブルな方法が、すべての利用可能なデータでゼロから再学習した場合の性能と同等であることを、最終的な損失と言語モデル（LM）評価ベンチマークを用いて示します。具体的には、405Mパラメータのモデルスケールで、一般的に使用される2つのLLM事前学習データセット間の弱いが現実的な分布シフト（英語→英語）と、より強い分布シフト（英語→ドイツ語）について、大規模なデータセット（数千億トークン）を用いてこれを示します。大規模実験のために弱いが現実的なシフトを選択し、10BパラメータのLLMにおいても、我々の継続学習戦略が再学習ベースラインと同等であることを確認しました。我々の結果は、LLMがシンプルでスケーラブルな継続学習戦略を通じて成功裏に更新できることを示しており、再学習ベースラインと同等の性能を、計算リソースの一部のみで達成できることを実証しています。最後に、先行研究にインスパイアされ、LR再ウォーミングによって引き起こされる忘却を回避し、固定されたトークンバジェットに縛られない、コサイン学習率スケジュールの代替案を提案します。

English

Large language models (LLMs) are routinely pre-trained on billions of tokens, only to start the process over again once new data becomes available. A much more efficient solution is to continually pre-train these models, saving significant compute compared to re-training. However, the distribution shift induced by new data typically results in degraded performance on previous data or poor adaptation to the new data. In this work, we show that a simple and scalable combination of learning rate (LR) re-warming, LR re-decaying, and replay of previous data is sufficient to match the performance of fully re-training from scratch on all available data, as measured by final loss and language model (LM) evaluation benchmarks. Specifically, we show this for a weak but realistic distribution shift between two commonly used LLM pre-training datasets (EnglishrightarrowEnglish) and a stronger distribution shift (EnglishrightarrowGerman) at the 405M parameter model scale with large dataset sizes (hundreds of billions of tokens). Selecting the weak but realistic shift for larger-scale experiments, we also find that our continual learning strategies match the re-training baseline for a 10B parameter LLM. Our results demonstrate that LLMs can be successfully updated via simple and scalable continual learning strategies, matching the re-training baseline using only a fraction of the compute. Finally, inspired by previous work, we propose alternatives to the cosine learning rate schedule that help circumvent forgetting induced by LR re-warming and that are not bound to a fixed token budget.