Efficiënte schaalbaarheid van voorafgaande training

Efficient Pretraining Length Scaling

April 21, 2025

Auteurs: Bohong Wu, Shen Yan, Sijun Zhang, Jianqiao Lu, Yutao Zeng, Ya Wang, Xun Zhou

cs.AI

Samenvatting

Recente vooruitgang in grote taalmodel(len) heeft de effectiviteit van lengteschaling tijdens post-training aangetoond, maar het potentieel ervan in pre-training blijft onderbelicht. Wij presenteren de Parallel Hidden Decoding Transformer (PHD-Transformer), een nieuw framework dat efficiënte lengteschaling tijdens pre-training mogelijk maakt terwijl de inferentie-efficiëntie behouden blijft. De PHD-Transformer bereikt dit door een innovatieve KV-cachebeheerstrategie die onderscheid maakt tussen originele tokens en verborgen decodeertokens. Door alleen de KV-cache van originele tokens te behouden voor langeafstandsafhankelijkheden en verborgen decodeertokens direct na gebruik te verwijderen, behoudt onze aanpak dezelfde KV-cachegrootte als de standaard transformer terwijl effectieve lengteschaling mogelijk wordt gemaakt. Om de prestaties verder te verbeteren, introduceren we twee geoptimaliseerde varianten: PHD-SWA maakt gebruik van sliding window attention om lokale afhankelijkheden te behouden, terwijl PHD-CSWA chunk-wise sliding window attention implementeert om lineaire groei in pre-fillingtijd te elimineren. Uitgebreide experimenten tonen consistente verbeteringen aan over meerdere benchmarks.

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.