Datenengineering für die Skalierung von Sprachmodellen auf 128K Kontext

Data Engineering for Scaling Language Models to 128K Context

February 15, 2024

Autoren: Yao Fu, Rameswar Panda, Xinyao Niu, Xiang Yue, Hannaneh Hajishirzi, Yoon Kim, Hao Peng

cs.AI

Zusammenfassung

Wir untersuchen das Rezept für kontinuierliches Vortraining, um die Kontextlängen von Sprachmodellen auf 128K zu skalieren, mit einem Fokus auf Data Engineering. Wir stellen die Hypothese auf, dass die Modellierung langer Kontexte, insbesondere die Fähigkeit, Informationen an beliebigen Eingabepositionen zu nutzen, eine Fähigkeit ist, die größtenteils bereits durch groß angelegtes Vortraining erworben wurde und dass diese Fähigkeit durch leichtgewichtiges kontinuierliches Vortraining mit geeigneten Datenmischungen problemlos auf deutlich längere Kontexte (z.B. von 4K auf 128K) erweitert werden kann. Wir untersuchen die Menge und Qualität der Daten für das kontinuierliche Vortraining: (1) Für die Menge zeigen wir, dass 500 Millionen bis 5 Milliarden Token ausreichen, um dem Modell zu ermöglichen, Informationen überall innerhalb des 128K-Kontexts abzurufen; (2) Für die Qualität betonen unsere Ergebnisse gleichermaßen die Domänenbalance und die Längen-Upsampling. Konkret stellen wir fest, dass das naive Upsampling längerer Daten in bestimmten Domänen wie Büchern, eine gängige Praxis bestehender Arbeiten, suboptimale Leistung liefert und dass eine ausgewogene Domänenmischung wichtig ist. Wir demonstrieren, dass das kontinuierliche Vortraining des gesamten Modells mit 1 bis 5 Milliarden Token solcher Daten eine effektive und erschwingliche Strategie ist, um die Kontextlänge von Sprachmodellen auf 128K zu skalieren. Unser Rezept übertrifft starke Open-Source-Modelle mit langen Kontexten und schließt die Lücke zu Spitzenmodellen wie GPT-4 128K.

English

We study the continual pretraining recipe for scaling language models' context lengths to 128K, with a focus on data engineering. We hypothesize that long context modeling, in particular the ability to utilize information at arbitrary input locations, is a capability that is mostly already acquired through large-scale pretraining, and that this capability can be readily extended to contexts substantially longer than seen during training~(e.g., 4K to 128K) through lightweight continual pretraining on appropriate data mixture. We investigate the quantity and quality of the data for continual pretraining: (1) for quantity, we show that 500 million to 5 billion tokens are enough to enable the model to retrieve information anywhere within the 128K context; (2) for quality, our results equally emphasize domain balance and length upsampling. Concretely, we find that naively upsampling longer data on certain domains like books, a common practice of existing work, gives suboptimal performance, and that a balanced domain mixture is important. We demonstrate that continual pretraining of the full model on 1B-5B tokens of such data is an effective and affordable strategy for scaling the context length of language models to 128K. Our recipe outperforms strong open-source long-context models and closes the gap to frontier models like GPT-4 128K.