Mise à l'échelle des modèles de langage basés sur la récupération avec un magasin de données d'un trillion de tokens

Scaling Retrieval-Based Language Models with a Trillion-Token Datastore

July 9, 2024

Auteurs: Rulin Shao, Jacqueline He, Akari Asai, Weijia Shi, Tim Dettmers, Sewon Min, Luke Zettlemoyer, Pang Wei Koh

cs.AI

Résumé

Les lois d'échelle concernant la quantité de données d'entraînement et le nombre de paramètres nous permettent de prédire les compromis coût-bénéfice du pré-entraînement des modèles de langage (LM) dans différentes configurations. Dans cet article, nous examinons une autre dimension de l'échelle : la quantité de données disponibles au moment de l'inférence. Plus précisément, nous constatons que l'augmentation de la taille du magasin de données utilisé par un LM basé sur la récupération améliore de manière monotone la modélisation du langage et plusieurs tâches en aval sans saturation évidente, de sorte qu'un modèle plus petit augmenté d'un grand magasin de données surpasse un modèle LM plus grand sur des tâches nécessitant une connaissance approfondie. En traçant des courbes d'échelle optimales en termes de calcul avec des tailles variées de magasin de données, de modèle et de données de pré-entraînement, nous montrons que l'utilisation de magasins de données plus volumineux peut considérablement améliorer les performances du modèle pour un même budget de calcul d'entraînement. Nous menons notre étude en construisant un magasin de données de 1,4 trillion de tokens nommé MassiveDS, qui est le magasin de données le plus grand et le plus diversifié en open source pour les LM basés sur la récupération à ce jour, et en concevant un pipeline efficace pour étudier l'échelle des magasins de données de manière accessible sur le plan informatique. Enfin, nous analysons l'effet de l'amélioration du récupérateur, du filtrage de qualité du magasin de données et d'autres choix de conception sur les tendances d'échelle observées. Globalement, nos résultats montrent que la taille du magasin de données devrait être considérée comme une partie intégrante des compromis d'efficacité et de performance des LM. Pour faciliter les recherches futures, nous mettons en open source notre magasin de données et notre code à l'adresse https://github.com/RulinShao/retrieval-scaling.

English

Scaling laws with respect to the amount of training data and the number of parameters allow us to predict the cost-benefit trade-offs of pretraining language models (LMs) in different configurations. In this paper, we consider another dimension of scaling: the amount of data available at inference time. Specifically, we find that increasing the size of the datastore used by a retrieval-based LM monotonically improves language modeling and several downstream tasks without obvious saturation, such that a smaller model augmented with a large datastore outperforms a larger LM-only model on knowledge-intensive tasks. By plotting compute-optimal scaling curves with varied datastore, model, and pretraining data sizes, we show that using larger datastores can significantly improve model performance for the same training compute budget. We carry out our study by constructing a 1.4 trillion-token datastore named MassiveDS, which is the largest and the most diverse open-sourced datastore for retrieval-based LMs to date, and designing an efficient pipeline for studying datastore scaling in a computationally accessible manner. Finally, we analyze the effect of improving the retriever, datastore quality filtering, and other design choices on our observed scaling trends. Overall, our results show that datastore size should be considered as an integral part of LM efficiency and performance trade-offs. To facilitate future research, we open-source our datastore and code at https://github.com/RulinShao/retrieval-scaling.