Stratégies simples et évolutives pour le pré-entraînement continu des grands modèles de langage

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

March 13, 2024

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

cs.AI

Résumé

Les grands modèles de langage (LLM) sont généralement pré-entraînés sur des milliards de tokens, pour ensuite recommencer le processus dès que de nouvelles données deviennent disponibles. Une solution bien plus efficace consiste à pré-entraîner ces modèles de manière continue, ce qui permet d'économiser des ressources de calcul considérables par rapport à un ré-entraînement complet. Cependant, le décalage de distribution induit par les nouvelles données entraîne généralement une dégradation des performances sur les données précédentes ou une mauvaise adaptation aux nouvelles données. Dans ce travail, nous montrons qu'une combinaison simple et évolutive de réchauffement du taux d'apprentissage (LR), de réduction progressive du LR et de réutilisation des données précédentes suffit à égaler les performances d'un ré-entraînement complet sur toutes les données disponibles, mesurées par la perte finale et les benchmarks d'évaluation des modèles de langage (LM). Nous démontrons cela pour un décalage de distribution faible mais réaliste entre deux ensembles de données couramment utilisés pour le pré-entraînement des LLM (Anglais→Anglais) et un décalage plus marqué (Anglais→Allemand) à l'échelle d'un modèle de 405M de paramètres avec des ensembles de données volumineux (centaines de milliards de tokens). En choisissant le décalage faible mais réaliste pour des expériences à plus grande échelle, nous constatons également que nos stratégies d'apprentissage continu égalent la base de référence du ré-entraînement pour un LLM de 10B de paramètres. Nos résultats montrent que les LLM peuvent être mis à jour avec succès grâce à des stratégies d'apprentissage continu simples et évolutives, égalant la base de référence du ré-entraînement en utilisant seulement une fraction des ressources de calcul. Enfin, inspirés par des travaux antérieurs, nous proposons des alternatives au planning de taux d'apprentissage cosinus qui aident à contourner l'oubli induit par le réchauffement du LR et qui ne sont pas liées à un budget fixe de tokens.

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.