1兆トークンのデータストアを用いた検索ベース言語モデルのスケーリング

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

July 9, 2024

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

cs.AI

要旨

訓練データ量とパラメータ数に関するスケーリング法則を用いることで、異なる構成での言語モデル（LM）の事前学習におけるコストと利益のトレードオフを予測することが可能です。本論文では、スケーリングの別の次元として、推論時に利用可能なデータ量を考察します。具体的には、検索ベースのLMが使用するデータストアのサイズを増やすことで、言語モデリングや複数の下流タスクが単調に改善し、明らかな飽和が見られないことを発見しました。その結果、大規模なデータストアで補強された小規模モデルが、知識集約型タスクにおいて大規模なLMのみのモデルを上回ることを示しています。データストア、モデル、事前学習データのサイズを変えた計算最適なスケーリング曲線をプロットすることで、同じ訓練計算予算でより大きなデータストアを使用することでモデルの性能が大幅に向上することを示します。本研究では、1.4兆トークンのデータストア「MassiveDS」を構築し、これまでで最大かつ最も多様なオープンソースの検索ベースLM用データストアとして提供しました。また、計算的にアクセス可能な方法でデータストアのスケーリングを研究するための効率的なパイプラインを設計しました。最後に、リトリーバーの改善、データストアの品質フィルタリング、およびその他の設計選択が観測されたスケーリングトレンドに与える影響を分析します。全体として、データストアのサイズはLMの効率性と性能のトレードオフの重要な要素として考慮されるべきであることが示されました。今後の研究を促進するため、データストアとコードを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.