PHD-Transformer 由字节Seed 团队发布，有效增幅了预训练模型长度并克服了 KV 缓存难题！

字节跳动Seed团队最近推出了一种名为PHD-Transformer（Parallel Hidden Decoding Transformer）的模型架构，成功解决了预训练长度扩展和KV缓存问题。

PHD-Transformer的核心创新

PHD-Transformer通过一种创新的KV缓存管理策略，实现了预训练长度的有效扩展，同时保持了与原始Transformer相同的KV缓存大小。具体来说：

• 重复输入tokens：研究者通过直接重复输入tokens（例如重复1到4次），而不是在中间层进行复杂处理，观察到训练损失和模型性能随重复倍数的扩展趋势。

• KV缓存优化：PHD-Transformer将第一个token视为原始token，其余重复的token视为解码token。它仅保留从原始token生成的KV缓存用于长距离依赖建模，并在解码token用于下一个token预测后丢弃其KV缓存。

• 滑动窗口注意力机制：为了进一步提升性能，研究者引入了PHD-SWA（滑动窗口注意力）和PHD-CSWA（逐块滑动窗口注意力）。PHD-SWA通过为隐藏解码token保留局部滑动窗口缓存，实现了显著的性能提升，同时仅需要少量额外的KV缓存内存。PHD-CSWA则通过限制每个块内的顺序依赖关系，显著缩短了预填充时间。

实验结果

PHD-Transformer及其变体在多个基准测试中表现出色：

• PHD-CSWA-2-16-32：token重复两次，保留16个token的局部窗口，块大小为32。在基准测试中平均准确率提升了1.5%，训练损失降低了0.025。

• PHD-CSWA-3-16-32：token重复三次，局部窗口和块大小与上述设置相同。平均准确率提升了2.0%，训练损失降低了0.034。

• 扩展性能：当扩展因子设置为5时，所有基准测试的平均准确率提高了1.8%，证实了该方法在更激进的扩展方面仍然有效。

优势与意义

PHD-Transformer的主要优势在于：

• 保持KV缓存大小不变：与简单重复tokens的方法相比，PHD-Transformer在推理过程中保持了与原始Transformer相同的KV缓存大小，避免了KV缓存规模线性增加的问题。

• 显著提升推理效率：通过优化KV缓存管理和注意力机制，PHD-Transformer在推理速度上实现了显著提升，同时降低了内存压力。

• 性能提升：在多个基准测试中，PHD-Transformer及其变体均实现了持续的性能提升，证明了其在预训练长度扩展方面的有效性。

PHD-Transformer的推出为Transformer架构的改进提供了新的思路，特别是在预训练长度扩展和KV缓存管理方面，为未来大模型的研究和应用提供了重要的参考。