高效的预训练长度缩放

大型语言模型最近的进展表明了后训练阶段长度缩放的有效性，但其在预训练阶段的潜力仍未得到充分探索。我们提出了并行隐藏解码 Transformer (PHD-Transformer)，这是一种新颖的框架，可以在预训练期间实现高效的长度缩放，同时保持推理效率。PHD-Transformer 通过一种创新的 KV 缓存管理策略来实现这一点，该策略区分原始令牌和隐藏解码令牌。通过仅保留原始令牌的 KV 缓存用于长距离依赖关系，并在使用后立即丢弃隐藏解码令牌，我们的方法在实现有效长度缩放的同时保持了与普通 Transformer 相同的 KV 缓存大小。为了进一步提升性能，我们引入了两种优化变体：PHD-SWA 采用滑动窗口注意力以保留局部依赖关系，而 PHD-CSWA 实现块状滑动窗口注意力以消除预填充时间的线性增长。大量实验表明，在多个基准上持续取得了改进。