提出PHD-Transformer,实现预训练阶段高效长文本扩展。
Efficient Pretraining Length Scaling
- 通过区分原始词元与隐式解码词元,优化KV缓存管理。
- 在不增加缓存大小前提下,实现长序列预训练效果提升。
- 适合需要长文本建模的LLM研究者和开发者。
近期大语言模型的研究表明,后训练阶段的长度扩展有效,但预训练阶段的潜力尚未充分探索。本文提出并行隐藏解码变换器(PHD-Transformer),一种可在预训练中实现高效长度扩展的新框架,同时保持推理效率。该方法通过创新的KV缓存管理策略,区分原始词元与隐藏解码词元:仅保留原始词元的KV缓存以维持长程依赖,而使用后立即丢弃隐藏解码词元的缓存,从而在不增加缓存规模的前提下实现有效的长度扩展。为进一步提升性能,提出两个优化变体:PHD-SWA采用滑动窗口注意力以保留局部依赖;PHD-CSWA引入分块滑动窗口注意力,消除预填充时间的线性增长。大量实验表明,该方法在多个基准上均取得持续改进。
原文摘要 · Abstract (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 (\textit{PHD}-Transformer), a novel framework that enables efficient length scaling during pre-training while maintaining inference efficiency. \textit{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: \textit{PHD-SWA} employs sliding window attention to preserve local dependencies, while \textit{PHD-CSWA} implements chunk-wise sliding window attention to eliminate linear growth in pre-filling time. Extensive experiments demonstrate consistent improvements across multiple benchmarks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。