SkyLadder: Besseres und schnelleres Pretraining durch Kontextfensterplanung

SkyLadder: Better and Faster Pretraining via Context Window Scheduling

March 19, 2025

Autoren: Tongyao Zhu, Qian Liu, Haonan Wang, Shiqi Chen, Xiangming Gu, Tianyu Pang, Min-Yen Kan

cs.AI

Zusammenfassung

Jüngste Fortschritte im Pre-Training von LLMs (Large Language Models) waren durch immer größere Kontextfenster geprägt, um längere Sequenzen verarbeiten zu können. Unsere Pilotstudie zeigt jedoch, dass Modelle, die mit kürzeren Kontextfenstern vortrainiert wurden, bei einem festen Token-Budget durchweg besser abschneiden als ihre Pendants mit langen Kontextfenstern. Diese Erkenntnis motiviert uns, eine optimale Strategie für die Planung von Kontextfenstern zu erforschen, um die Fähigkeit zur Verarbeitung langer Kontexte besser mit der Effizienz des Pre-Trainings in Einklang zu bringen. Zu diesem Zweck schlagen wir SkyLadder vor, einen einfachen, aber effektiven Ansatz, der einen Übergang von kurzen zu langen Kontextfenstern implementiert. SkyLadder bewahrt eine starke Leistung auf Standard-Benchmarks, während es die Ergebnisse der Baselines bei Aufgaben mit langen Kontexten erreicht oder übertrifft. Durch umfangreiche Experimente haben wir 1B-Parameter-Modelle (bis zu 32K Kontext) und 3B-Parameter-Modelle (8K Kontext) auf 100B Tokens vortrainiert und gezeigt, dass SkyLadder auf gängigen Benchmarks konsistente Gewinne von bis zu 3,7 % erzielt, während es im Vergleich zu den Baselines eine bis zu 22 % schnellere Trainingsgeschwindigkeit erreicht. Der Code ist unter https://github.com/sail-sg/SkyLadder verfügbar.

English

Recent advancements in LLM pretraining have featured ever-expanding context windows to process longer sequences. However, our pilot study reveals that models pretrained with shorter context windows consistently outperform their long-context counterparts under a fixed token budget. This finding motivates us to explore an optimal context window scheduling strategy to better balance long-context capability with pretraining efficiency. To this end, we propose SkyLadder, a simple yet effective approach that implements a short-to-long context window transition. SkyLadder preserves strong standard benchmark performance, while matching or exceeding baseline results on long context tasks. Through extensive experiments, we pre-train 1B-parameter models (up to 32K context) and 3B-parameter models (8K context) on 100B tokens, demonstrating that SkyLadder yields consistent gains of up to 3.7% on common benchmarks, while achieving up to 22% faster training speeds compared to baselines. The code is at https://github.com/sail-sg/SkyLadder.