Scalabilità dei Modelli Linguistici Basati su Recupero con un Archivio Dati da Mille Miliardi di Token

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

July 9, 2024

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

cs.AI

Abstract

Le leggi di scala rispetto alla quantità di dati di addestramento e al numero di parametri ci permettono di prevedere i compromessi costo-beneficio del pre-addestramento di modelli linguistici (LM) in diverse configurazioni. In questo articolo, consideriamo un'altra dimensione della scalabilità: la quantità di dati disponibili al momento dell'inferenza. Nello specifico, scopriamo che aumentare le dimensioni del datastore utilizzato da un LM basato su retrieval migliora in modo monotono la modellazione del linguaggio e diverse attività downstream senza evidenti saturazioni, tanto che un modello più piccolo potenziato con un datastore di grandi dimensioni supera un modello LM più grande su compiti ad alta intensità di conoscenza. Tracciando curve di scala ottimali in termini di calcolo con dimensioni variabili del datastore, del modello e dei dati di pre-addestramento, dimostriamo che l'uso di datastore più grandi può migliorare significativamente le prestazioni del modello per lo stesso budget di calcolo di addestramento. Portiamo avanti il nostro studio costruendo un datastore da 1,4 trilioni di token chiamato MassiveDS, che è il datastore open-source più grande e diversificato per LM basati su retrieval fino ad oggi, e progettando una pipeline efficiente per studiare la scalabilità del datastore in modo computazionalmente accessibile. Infine, analizziamo l'effetto del miglioramento del retriever, del filtraggio della qualità del datastore e di altre scelte progettuali sulle tendenze di scala osservate. Nel complesso, i nostri risultati dimostrano che le dimensioni del datastore dovrebbero essere considerate come parte integrante dei compromessi tra efficienza e prestazioni dei LM. Per facilitare la ricerca futura, rendiamo open-source il nostro datastore e il codice all'indirizzo 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.