SkyLadder: Mejor y más rápido preentrenamiento mediante programación de ventanas de contexto

SkyLadder: Better and Faster Pretraining via Context Window Scheduling

March 19, 2025

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

cs.AI

Resumen

Los avances recientes en el preentrenamiento de LLM han destacado ventanas de contexto en constante expansión para procesar secuencias más largas. Sin embargo, nuestro estudio piloto revela que los modelos preentrenados con ventanas de contexto más cortas superan consistentemente a sus contrapartes de contexto largo bajo un presupuesto fijo de tokens. Este hallazgo nos motiva a explorar una estrategia óptima de programación de ventanas de contexto para equilibrar mejor la capacidad de contexto largo con la eficiencia del preentrenamiento. Con este fin, proponemos SkyLadder, un enfoque simple pero efectivo que implementa una transición de ventanas de contexto cortas a largas. SkyLadder preserva un rendimiento sólido en los puntos de referencia estándar, mientras iguala o supera los resultados de referencia en tareas de contexto largo. A través de extensos experimentos, preentrenamos modelos de 1B parámetros (hasta 32K de contexto) y modelos de 3B parámetros (8K de contexto) en 100B tokens, demostrando que SkyLadder produce ganancias consistentes de hasta un 3.7% en puntos de referencia comunes, mientras logra velocidades de entrenamiento hasta un 22% más rápidas en comparación con las líneas de base. El código está disponible en 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.