SkyLadder: Улучшенное и ускоренное предварительное обучение с помощью планирования контекстного окна

SkyLadder: Better and Faster Pretraining via Context Window Scheduling

March 19, 2025

Авторы: Tongyao Zhu, Qian Liu, Haonan Wang, Shiqi Chen, Xiangming Gu, Tianyu Pang, Min-Yen Kan

cs.AI

Аннотация

Последние достижения в предварительном обучении языковых моделей (LLM) характеризуются постоянно расширяющимися окнами контекста для обработки более длинных последовательностей. Однако наше пилотное исследование показывает, что модели, предварительно обученные с более короткими окнами контекста, стабильно превосходят свои аналоги с длинным контекстом при фиксированном бюджете токенов. Это открытие побуждает нас изучить оптимальную стратегию планирования окна контекста, чтобы лучше сбалансировать способность к обработке длинного контекста с эффективностью предварительного обучения. С этой целью мы предлагаем SkyLadder — простой, но эффективный подход, который реализует переход от короткого к длинному окну контекста. SkyLadder сохраняет высокую производительность на стандартных бенчмарках, одновременно соответствуя или превосходя базовые результаты на задачах с длинным контекстом. В ходе обширных экспериментов мы предварительно обучаем модели с 1 миллиардом параметров (до 32K контекста) и модели с 3 миллиардами параметров (8K контекста) на 100 миллиардах токенов, демонстрируя, что SkyLadder обеспечивает стабильный прирост до 3,7% на общих бенчмарках, одновременно достигая до 22% более высокой скорости обучения по сравнению с базовыми подходами. Код доступен по адресу 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.