近日，字節(jié)跳動的 Seed 團隊在人工智能領(lǐng)域再傳佳音，推出了一種新型的 PHD-Transformer（Parallel Hidden Decoding Transformer），這項創(chuàng)新突破了預(yù)訓(xùn)練長度的限制，有效解決了推理過程中的 KV 緩存膨脹問題。隨著大型推理模型的迅速發(fā)展，研究人員在后訓(xùn)練階段嘗試通過強化學(xué)習(xí)方法來生成更長的推理鏈，并在復(fù)雜的推理任務(wù)上取得了顯著成果。受到啟發(fā)，字節(jié) Seed 團隊決定探索在預(yù)訓(xùn)練階段進行長度擴展的可能性。
傳統(tǒng)的長度擴展方法常常涉及在序列中插入文本或潛在向量，這些方式往往導(dǎo)致 KV 緩存占用過多內(nèi)存，推理速度慢。而 PHD-Transformer 則采用了一種更為簡便的策略 —— 直接重復(fù)輸入的 tokens。盡管這種方法在訓(xùn)練損失和模型性能上有明顯的提升，卻也帶來了 KV 緩存線性增長、內(nèi)存壓力加大和解碼延遲等新問題。
為了解決這些挑戰(zhàn)，PHD-Transformer 通過創(chuàng)新的 KV 緩存管理策略，保留了與原始 Transformer 相同的緩存大小。在推理時，PHD-Transformer 只保留由原始 tokens 生成的 KV 緩存，而對重復(fù)的 tokens 則在預(yù)測后立即丟棄，顯著加快了推理速度。此外，研究團隊還引入了滑動窗口注意力機制，稱之為 PHD-SWA，以保持局部滑動窗口緩存的性能優(yōu)勢。為進一步優(yōu)化預(yù)填充時間，研究者提出了逐塊滑動窗口注意力機制 PHD-CSWA，限制了每個塊內(nèi)的順序依賴，從而大幅縮短了預(yù)填充時間。
在一系列實驗中，PHD-CSWA 表現(xiàn)出色，在多個公開基準測試集上均實現(xiàn)了準確率的提升。研究團隊表示，PHD-CSWA 在保持原有效率的同時，為模型帶來了更大的性能提升，標志著在大規(guī)模推理任務(wù)中，預(yù)訓(xùn)練長度擴展取得了實質(zhì)性進展。