Эффективное масштабирование длины предварительного обучения

Efficient Pretraining Length Scaling

April 21, 2025

Авторы: Bohong Wu, Shen Yan, Sijun Zhang, Jianqiao Lu, Yutao Zeng, Ya Wang, Xun Zhou

cs.AI

Аннотация

Недавние достижения в области крупных языковых моделей продемонстрировали эффективность масштабирования длины в процессе посттренировки, однако его потенциал на этапе предтренировки остается недостаточно изученным. Мы представляем Parallel Hidden Decoding Transformer (PHD-Transformer) — новую архитектуру, которая обеспечивает эффективное масштабирование длины на этапе предтренировки, сохраняя при этом эффективность на этапе вывода. PHD-Transformer достигает этого за счет инновационной стратегии управления кэшем ключей и значений (KV), которая различает исходные токены и токены скрытого декодирования. Сохраняя кэш KV только для исходных токенов для учета долгосрочных зависимостей и немедленно удаляя токены скрытого декодирования после их использования, наш подход сохраняет тот же размер кэша KV, что и в классическом трансформере, обеспечивая при этом эффективное масштабирование длины. Для дальнейшего повышения производительности мы представляем две оптимизированные версии: PHD-SWA использует скользящее окно внимания для сохранения локальных зависимостей, а PHD-CSWA реализует чанковое скользящее окно внимания, чтобы устранить линейный рост времени предзаполнения. Многочисленные эксперименты демонстрируют стабильные улучшения на множестве бенчмарков.

English

Recent advances in large language models have demonstrated the effectiveness of length scaling during post-training, yet its potential in pre-training remains underexplored. We present the Parallel Hidden Decoding Transformer (PHD-Transformer), a novel framework that enables efficient length scaling during pre-training while maintaining inference efficiency. PHD-Transformer achieves this through an innovative KV cache management strategy that distinguishes between original tokens and hidden decoding tokens. By retaining only the KV cache of original tokens for long-range dependencies while immediately discarding hidden decoding tokens after use, our approach maintains the same KV cache size as the vanilla transformer while enabling effective length scaling. To further enhance performance, we introduce two optimized variants: PHD-SWA employs sliding window attention to preserve local dependencies, while PHD-CSWA implements chunk-wise sliding window attention to eliminate linear growth in pre-filling time. Extensive experiments demonstrate consistent improvements across multiple benchmarks.