SkyLadder : Un prétraitement meilleur et plus rapide grâce à la planification des fenêtres contextuelles

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

Résumé

Les récents progrès dans le pré-entraînement des LLM ont mis en avant des fenêtres de contexte toujours plus larges pour traiter des séquences plus longues. Cependant, notre étude pilote révèle que les modèles pré-entraînés avec des fenêtres de contexte plus courtes surpassent systématiquement leurs homologues à contexte long pour un budget fixe de tokens. Cette découverte nous motive à explorer une stratégie optimale de planification des fenêtres de contexte pour mieux équilibrer la capacité à gérer des contextes longs avec l'efficacité du pré-entraînement. À cette fin, nous proposons SkyLadder, une approche simple mais efficace qui met en œuvre une transition progressive de fenêtres de contexte courtes à longues. SkyLadder préserve de solides performances sur les benchmarks standards, tout en égalant ou dépassant les résultats de référence sur les tâches à contexte long. À travers des expériences approfondies, nous pré-entraînons des modèles de 1 milliard de paramètres (jusqu'à 32K de contexte) et de 3 milliards de paramètres (8K de contexte) sur 100 milliards de tokens, démontrant que SkyLadder apporte des gains constants allant jusqu'à 3,7 % sur les benchmarks courants, tout en atteignant des vitesses d'entraînement jusqu'à 22 % plus rapides par rapport aux références. Le code est disponible à l'adresse suivante : 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.