SkyLadder: Betere en Snellere Pretraining via Contextvensterplanning

SkyLadder: Better and Faster Pretraining via Context Window Scheduling

March 19, 2025

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

cs.AI

Samenvatting

Recente vooruitgang in het vooraf trainen van LLM's heeft steeds grotere contextvensters gekend om langere sequenties te verwerken. Uit onze pilotstudie blijkt echter dat modellen die zijn voorgetraind met kortere contextvensters consequent beter presteren dan hun tegenhangers met lange contextvensters, binnen een vast tokenbudget. Deze bevinding motiveert ons om een optimale strategie voor het plannen van contextvensters te onderzoeken, om zo een betere balans te vinden tussen de mogelijkheid om lange contexten te verwerken en de efficiëntie van het vooraf trainen. Hiertoe stellen we SkyLadder voor, een eenvoudige maar effectieve aanpak die een overgang van kort naar lang contextvenster implementeert. SkyLadder behoudt een sterke prestatie op standaard benchmarks, terwijl het de basislijnresultaten op taken met lange context evenaart of overtreft. Door middel van uitgebreide experimenten hebben we modellen met 1B parameters (tot 32K context) en 3B parameters (8K context) voorgetraind op 100B tokens, waarbij we aantonen dat SkyLadder consistente verbeteringen oplevert van tot 3,7% op gangbare benchmarks, terwijl het tot 22% snellere trainingssnelheden bereikt in vergelijking met de basislijnen. De code is beschikbaar op https://github.com/sail-sg/SkyLadder.

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.