SkyLadder: Pré-treinamento Melhor e Mais Rápido por meio de Agendamento de Janelas 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

Resumo

Os avanços recentes no pré-treinamento de LLMs têm apresentado janelas de contexto cada vez mais amplas para processar sequências mais longas. No entanto, nosso estudo piloto revela que modelos pré-treinados com janelas de contexto mais curtas consistentemente superam suas contrapartes de contexto longo sob um orçamento fixo de tokens. Essa descoberta nos motiva a explorar uma estratégia ideal de escalonamento de janelas de contexto para equilibrar melhor a capacidade de contexto longo com a eficiência do pré-treinamento. Para isso, propomos o SkyLadder, uma abordagem simples, porém eficaz, que implementa uma transição de janela de contexto curta para longa. O SkyLadder mantém um forte desempenho em benchmarks padrão, ao mesmo tempo em que iguala ou supera os resultados de baseline em tarefas de contexto longo. Por meio de extensos experimentos, pré-treinamos modelos de 1B parâmetros (até 32K de contexto) e modelos de 3B parâmetros (8K de contexto) em 100B de tokens, demonstrando que o SkyLadder proporciona ganhos consistentes de até 3,7% em benchmarks comuns, enquanto alcança velocidades de treinamento até 22% mais rápidas em comparação com as baselines. O código está disponível em 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.