近日，字节跳动的 Seed 团队在人工智能领域再传佳音，推出了一种新型的 PHD-Transformer（Parallel Hidden Decoding Transformer），这项创新突破了预训练长度的限制，有效解决了推理过程中的 KV 缓存膨胀问题。随着大型推理模型的迅速发展，研究人员在后训练阶段尝试通过强化学习方法来生成更长的推理链，并在复杂的推理任务上取得了显著成果。受到启发，字节 Seed 团队决定探索在预训练阶段进行长度扩展的可能性。
传统的长度扩展方法常常涉及在序列中插入文本或潜在向量，这些方式往往导致 KV 缓存占用过多内存，推理速度慢。而 PHD-Transformer 则采用了一种更为简便的策略 —— 直接重复输入的 tokens。尽管这种方法在训练损失和模型性能上有明显的提升，却也带来了 KV 缓存线性增长、内存压力加大和解码延迟等新问题。
为了解决这些挑战，PHD-Transformer 通过创新的 KV 缓存管理策略，保留了与原始 Transformer 相同的缓存大小。在推理时，PHD-Transformer 只保留由原始 tokens 生成的 KV 缓存，而对重复的 tokens 则在预测后立即丢弃，显著加快了推理速度。此外，研究团队还引入了滑动窗口注意力机制，称之为 PHD-SWA，以保持局部滑动窗口缓存的性能优势。为进一步优化预填充时间，研究者提出了逐块滑动窗口注意力机制 PHD-CSWA，限制了每个块内的顺序依赖，从而大幅缩短了预填充时间。
在一系列实验中，PHD-CSWA 表现出色，在多个公开基准测试集上均实现了准确率的提升。研究团队表示，PHD-CSWA 在保持原有效率的同时，为模型带来了更大的性能提升，标志着在大规模推理任务中，预训练长度扩展取得了实质性进展。