安定性ギャップの軽減による効率的な継続事前学習

Efficient Continual Pre-training by Mitigating the Stability Gap

June 21, 2024

著者: Yiduo Guo, Jie Fu, Huishuai Zhang, Dongyan Zhao, Yikang Shen

cs.AI

要旨

継続的プレトレーニングは、大規模言語モデル（LLM）を新しいドメインに適応させるための主要なアプローチとしてますます普及しています。このプロセスでは、事前学習済みのLLMを新しいドメインのコーパスで更新し、トレーニング分布をシフトさせます。このシフト中のLLMの挙動を研究するため、継続的プレトレーニングプロセス全体にわたってモデルのパフォーマンスを測定しました。その結果、最初に一時的なパフォーマンスの低下が観察され、その後回復フェーズが続く「安定性ギャップ」と呼ばれる現象が確認されました。これは、新しいクラスを分類する視覚モデルで以前に指摘された現象です。この問題に対処し、固定された計算予算内でLLMのパフォーマンスを向上させるために、以下の3つの効果的な戦略を提案します：（1）適切なサイズのサブセットでLLMを複数エポックにわたって継続的にプレトレーニングすることで、大規模なコーパスを単一エポックでプレトレーニングするよりも迅速にパフォーマンスを回復させる。（2）高品質なサブコーパスのみでLLMをプレトレーニングし、ドメインパフォーマンスを急速に向上させる。（3）プレトレーニングデータに類似したデータ混合を使用して分布ギャップを縮小する。これらの戦略の有効性を検証するため、Llamaファミリーモデルを用いて医療分野の継続的プレトレーニングと指示チューニングに関する様々な実験を実施しました。例えば、提案した戦略により、OpenLlama-3Bモデルの平均医療タスクパフォーマンスが36.2%から40.7%に向上し、元のトレーニング予算の40%のみで達成されました。さらに、一般タスクの平均パフォーマンスも向上し、忘却を引き起こすことなく改善されました。さらに、これらの戦略をLlama-3-8Bモデルに適用しました。その結果得られたモデル、Llama-3-Physicianは、現在のオープンソースモデルの中で最高の医療パフォーマンスを達成し、いくつかの医療ベンチマークではGPT-4と同等またはそれ以上の性能を示しました。モデルはhttps://huggingface.co/YiDuo1999/Llama-3-Physician-8B-Instructで公開しています。

English

Continual pre-training has increasingly become the predominant approach for adapting Large Language Models (LLMs) to new domains. This process involves updating the pre-trained LLM with a corpus from a new domain, resulting in a shift in the training distribution. To study the behavior of LLMs during this shift, we measured the model's performance throughout the continual pre-training process. we observed a temporary performance drop at the beginning, followed by a recovery phase, a phenomenon known as the "stability gap," previously noted in vision models classifying new classes. To address this issue and enhance LLM performance within a fixed compute budget, we propose three effective strategies: (1) Continually pre-training the LLM on a subset with a proper size for multiple epochs, resulting in faster performance recovery than pre-training the LLM on a large corpus in a single epoch; (2) Pre-training the LLM only on high-quality sub-corpus, which rapidly boosts domain performance; and (3) Using a data mixture similar to the pre-training data to reduce distribution gap. We conduct various experiments on Llama-family models to validate the effectiveness of our strategies in both medical continual pre-training and instruction tuning. For example, our strategies improve the average medical task performance of the OpenLlama-3B model from 36.2% to 40.7% with only 40% of the original training budget and enhance the average general task performance without causing forgetting. Furthermore, we apply our strategies to the Llama-3-8B model. The resulting model, Llama-3-Physician, achieves the best medical performance among current open-source models, and performs comparably to or even better than GPT-4 on several medical benchmarks. We release our models at https://huggingface.co/YiDuo1999/Llama-3-Physician-8B-Instruct.