Effiziente Skalierung der Vorausbildungslänge

Efficient Pretraining Length Scaling

April 21, 2025

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

cs.AI

Zusammenfassung

Jüngste Fortschritte bei großen Sprachmodellen haben die Wirksamkeit der Längenskalierung während des Post-Trainings aufgezeigt, doch ihr Potenzial im Pre-Training bleibt weitgehend unerforscht. Wir stellen den Parallel Hidden Decoding Transformer (PHD-Transformer) vor, ein neuartiges Framework, das eine effiziente Längenskalierung während des Pre-Trainings ermöglicht, während die Inferenzeffizienz erhalten bleibt. Der PHD-Transformer erreicht dies durch eine innovative KV-Cache-Management-Strategie, die zwischen ursprünglichen Tokens und verborgenen Decoding-Tokens unterscheidet. Indem nur der KV-Cache der ursprünglichen Tokens für langreichweitige Abhängigkeiten beibehalten wird, während verborgene Decoding-Tokens nach der Verwendung sofort verworfen werden, behält unser Ansatz die gleiche KV-Cache-Größe wie der herkömmliche Transformer bei und ermöglicht gleichzeitig eine effektive Längenskalierung. Um die Leistung weiter zu steigern, führen wir zwei optimierte Varianten ein: PHD-SWA verwendet Sliding-Window-Attention, um lokale Abhängigkeiten zu bewahren, während PHD-CSWA eine chunkweise Sliding-Window-Attention implementiert, um das lineare Wachstum der Vorbereitungszeit zu eliminieren. Umfangreiche Experimente zeigen konsistente Verbesserungen über mehrere Benchmarks hinweg.

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.