字節跳動旗下的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在保持原有高效性的基礎上，為模型帶來了更大的性能飛躍，這標志著在大規模推理任務中，預訓練長度擴展技術取得了實質性的進步。
