字节跳动旗下的Seed团队近期在人工智能领域取得了新的突破，他们研发出一种名为PHD-Transformer（Parallel Hidden Decoding Transformer）的新型模型，该模型在预训练长度扩展方面展现出了非凡的能力，有效解决了推理过程中KV缓存膨胀的难题。随着大型推理模型的快速发展，如何在后训练阶段生成更长的推理链，并在复杂任务上取得优异表现，成为了研究人员关注的焦点。Seed团队受此启发，决定在预训练阶段就进行长度扩展的探索。
传统方法在处理长度扩展时，通常会在序列中插入额外的文本或潜在向量，但这种方式会导致KV缓存占用大量内存，从而影响推理速度。而PHD-Transformer则采用了更为高效的策略，即直接重复输入的tokens。然而，这种方法虽然提升了训练损失和模型性能，但也带来了新的问题，如KV缓存的线性增长、内存压力的增加以及解码延迟的加剧。
为了克服这些挑战，PHD-Transformer引入了创新的KV缓存管理策略。在推理过程中，该模型只保留由原始tokens生成的KV缓存，对于重复的tokens，则在预测完成后立即丢弃。这一策略显著提升了推理速度。研究团队还开发了一种称为PHD-SWA（Sliding Window Attention）的滑动窗口注意力机制，以保持局部滑动窗口缓存的性能优势。为了进一步优化预填充时间，他们提出了PHD-CSWA（Chunk-wise Sliding Window Attention）机制，通过限制每个块内的顺序依赖，大幅缩短了预填充时间。
在一系列严格的实验中，PHD-CSWA展现出了卓越的性能。在多个公开的基准测试集上，该模型均实现了准确率的提升。Seed团队表示，PHD-CSWA在保持原有高效性的基础上，为模型带来了更大的性能飞跃，这标志着在大规模推理任务中，预训练长度扩展技术取得了实质性的进步。
