Efficient Pretraining Length Scaling
Authors: name: Bohong Wu; name: Shen Yan; name: Sijun Zhang; name: Jianqiao Lu; name: Yutao Zeng; name: Ya Wang; name: Xun Zhou
Published: 2025-04-21 09:41:26

Abstract: 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.