[論文レビュー] Continual Pre-Training of Large Language Models: How to (re)warm your model?

Kshitij Gupta, Benjamin Thérien+6 more|arXiv (Cornell University)|Aug 8, 2023

Topic Modeling被引用数 12

ひとこと要約

本論はLLMの継続的事前学習におけるウォームアップ戦略を検討し、再ウォームアップ学習率が下流適応を改善する一方で上流の忘却に影響を与えること、そして最新のチェックポイントからの再ウォームアップがスクラッチからの再学習よりも優れていることを示している。

ABSTRACT

Large language models (LLMs) are routinely pre-trained on billions of tokens, only to restart the process over again once new data becomes available. A much cheaper and more efficient solution would be to enable the continual pre-training of these models, i.e. updating pre-trained models with new data instead of re-training them from scratch. However, the distribution shift induced by novel data typically results in degraded performance on past data. Taking a step towards efficient continual pre-training, in this work, we examine the effect of different warm-up strategies. Our hypothesis is that the learning rate must be re-increased to improve compute efficiency when training on a new dataset. We study the warmup phase of models pre-trained on the Pile (upstream data, 300B tokens) as we continue to pre-train on SlimPajama (downstream data, 297B tokens), following a linear warmup and cosine decay schedule. We conduct all experiments on the Pythia 410M language model architecture and evaluate performance through validation perplexity. We experiment with different pre-training checkpoints, various maximum learning rates, and various warmup lengths. Our results show that while rewarming models first increases the loss on upstream and downstream data, in the longer run it improves the downstream performance, outperforming models trained from scratch$\unicode{x2013}$even for a large downstream dataset.

研究の動機と目的

• 新しいデータが利用可能になったとき、完全再学習よりも安価な代替として継続的事前学習を動機づける。
• 2つの大規模テキストコーパスに対する継続的事前学習の間に、ウォームアップ戦略が忘却と適応に与える影響を調査する。
• 上流と下流の性能のバランスを取るための、ウォームアップ長さ、最大学習率、チェックポイント選択の実用的なガイドラインを特定する。

提案手法

• Pileで事前学習された410MのPythiaモデルを上流データとして使用し、SlimPajama（≈297Bトークン）で継続的事前学習を下流データとして行う。
• 線形ウォームアップとコサイン減衰学習率スケジュールを組み合わせて実験する。
• ウォームアップ長さ（下流データの0–2%）と最大学習率（1.5e-4、3e-4、6e-4）を変化させる。
• 上流（Pile）と下流（SlimPajama）データセットの検証 perplexity によって性能を評価する。
• 最新、最大損失付近の以前のもの、そしてその中間の1つなど、異なる事前学習チェックポイントをテストする。
• 同じスケジュールで継続的事前学習とスクラッチからの訓練を比較する?

実験結果

リサーチクエスチョン

• RQ1 ウォームアップ段階の長さは、継続的事前学習中の下流または上流の perplexity に有意な影響を与えるか？
• RQ2 再暖機中の最大学習率は、下流の改善と上流の忘却とのトレードオフにどのような影響を与えるか？
• RQ3 Pile に似た上流コーパスデータで再暖機して継続的に事前学習する場合と、全く異なる下流コーパス SlimPajama の場合の再暖機は有益か？
• RQ4 新規データでの継続的事前学習において、以前の事前学習チェックポイントは適応を妨げるか、それとも助けるか？
• RQ5 大規模な下流データセットに対して、再暖機を伴う継続的事前学習はスクラッチからの訓練を上回ることができるか？

主な発見

• 短いウォームアップ長さ（データの0–2%）は下流または上流の perplexity に有意な影響を与えない。
• より大きな最大学習率は下流の性能を改善するが上流データの忘却を増加させる；より小さな学習率は上流の性能を維持するが下流の適応を制限する。
• 減衰コサインスケジュールでの再暖機は、スクラッチや一定LRのベースラインより下流性能が良いが、ウォームアップ期間には初期の不安定性が生じる。
• 最新のチェックポイントを用いた継続的事前学習はスクラッチ訓練より優れており、以前のチェックポイントを使用するより下流の結果が良い。
• 同じデータ（Pile）での再暖機も安定性のギャップを生じさせるが、下流の利益につながる可能性があり、分布シフトを超えた最適化ダイナミクスが影響していることを示唆する。
• 実験を通じて、ウォームアップを用いてファインチューニングしたモデルはスクラッチから訓練したモデルより優れており、一定LRでの早期停止は新しい分布へ適応する際に合理的である。