高效的预训练长度缩放
Bohong Wu, Shen Yan, Sijun Zhang, Jianqiao Lu, Yutao Zeng, Ya Wang, Xun Zhou
发布日期: 4/23/2025

摘要
近期在大型语言模型方面的进展展示了在后训练阶段使用长度缩放的有效性,然而其在预训练阶段的潜力尚未得到充分探索。我们提出了平行隐藏解码变换器(Parallel Hidden Decoding Transformer, PHD-Transformer),这是一种新颖的框架,能够在预训练阶段高效地实现长度缩放,同时保持推理效率。PHD-Transformer 通过一种创新的 KV 缓存管理策略实现了这一点,该策略区分了原始标记和隐藏解码标记。通过仅保留原始标记的 KV 缓存以维持长距离依赖关系,并在使用后立即丢弃隐藏解码标记,我们的方法保持了与普通变换器相同的 KV 缓存大小,同时实现了有效的长度缩放。为了进一步提升性能,我们引入了两种优化变体:PHD-SWA 采用滑动窗口注意力机制来保留局部依赖关系,而 PHD-CSWA 则实施分块滑动窗口注意力机制以消除预填充时间的线性增长。广泛的实验表明,在多个基准测试中均取得了持续的性能改进。