近日，字節跳動的 Seed 團隊在人工智能領域再傳佳音，推出了一種新型的 PHD-Transformer（Parallel Hidden Decoding Transformer），這項創新突破了預訓練長度的限制，有效解決了推理過程中的 KV 緩存膨脹問題。隨着大型推理模型的迅速發展，研究人員在後訓練階段嘗試通過強化學習方法來生成更長的推理鏈，並在複雜的推理任務上取得了顯著成果。受到啓發，字節 Seed 團隊決定探索在預訓練階段進行長度擴展的可能性。

傳統的長度擴展方法常常涉及在序列中插入文本或潛在向量，這些方式往往導致 KV 緩存佔用過多內存，推理速度慢。而 PHD-Transformer 則採用了一種更爲簡便的策略 —— 直接重複輸入的 tokens。儘管這種方法在訓練損失和模型性能上有明顯的提升，卻也帶來了 KV 緩存線性增長、內存壓力加大和解碼延遲等新問題。

image.png

爲了解決這些挑戰，PHD-Transformer 通過創新的 KV 緩存管理策略，保留了與原始 Transformer 相同的緩存大小。在推理時，PHD-Transformer 只保留由原始 tokens 生成的 KV 緩存，而對重複的 tokens 則在預測後立即丟棄，顯著加快了推理速度。此外，研究團隊還引入了滑動窗口注意力機制，稱之爲 PHD-SWA，以保持局部滑動窗口緩存的性能優勢。爲進一步優化預填充時間，研究者提出了逐塊滑動窗口注意力機制 PHD-CSWA，限制了每個塊內的順序依賴，從而大幅縮短了預填充時間。

在一系列實驗中，PHD-CSWA 表現出色，在多個公開基準測試集上均實現了準確率的提升。研究團隊表示，PHD-CSWA 在保持原有效率的同時，爲模型帶來了更大的性能提升，標誌着在大規模推理任務中，預訓練長度擴展取得了實質性進展。