近日，字節(jié)跳動(dòng)的 Seed 團(tuán)隊(duì)在人工智能領(lǐng)域再傳佳音，推出了一種新型的 PHD-Transformer（Parallel Hidden Decoding Transformer），這項(xiàng)創(chuàng)新突破了預(yù)訓(xùn)練長(zhǎng)度的限制，有效解決了推理過(guò)程中的 KV 緩存膨脹問(wèn)題。隨著大型推理模型的迅速發(fā)展，研究人員在后訓(xùn)練階段嘗試通過(guò)強(qiáng)化學(xué)習(xí)方法來(lái)生成更長(zhǎng)的推理鏈，并在復(fù)雜的推理任務(wù)上取得了顯著成果。受到啟發(fā)，字節(jié) Seed 團(tuán)隊(duì)決定探索在預(yù)訓(xùn)練階段進(jìn)行長(zhǎng)度擴(kuò)展的可能性。

傳統(tǒng)的長(zhǎng)度擴(kuò)展方法常常涉及在序列中插入文本或潛在向量，這些方式往往導(dǎo)致 KV 緩存占用過(guò)多內(nèi)存，推理速度慢。而 PHD-Transformer 則采用了一種更為簡(jiǎn)便的策略 —— 直接重復(fù)輸入的 tokens。盡管這種方法在訓(xùn)練損失和模型性能上有明顯的提升，卻也帶來(lái)了 KV 緩存線性增長(zhǎng)、內(nèi)存壓力加大和解碼延遲等新問(wèn)題。

為了解決這些挑戰(zhàn)，PHD-Transformer 通過(guò)創(chuàng)新的 KV 緩存管理策略，保留了與原始 Transformer 相同的緩存大小。在推理時(shí)，PHD-Transformer 只保留由原始 tokens 生成的 KV 緩存，而對(duì)重復(fù)的 tokens 則在預(yù)測(cè)后立即丟棄，顯著加快了推理速度。此外，研究團(tuán)隊(duì)還引入了滑動(dòng)窗口注意力機(jī)制，稱之為 PHD-SWA，以保持局部滑動(dòng)窗口緩存的性能優(yōu)勢(shì)。為進(jìn)一步優(yōu)化預(yù)填充時(shí)間，研究者提出了逐塊滑動(dòng)窗口注意力機(jī)制 PHD-CSWA，限制了每個(gè)塊內(nèi)的順序依賴，從而大幅縮短了預(yù)填充時(shí)間。

在一系列實(shí)驗(yàn)中，PHD-CSWA 表現(xiàn)出色，在多個(gè)公開(kāi)基準(zhǔn)測(cè)試集上均實(shí)現(xiàn)了準(zhǔn)確率的提升。研究團(tuán)隊(duì)表示，PHD-CSWA 在保持原有效率的同時(shí)，為模型帶來(lái)了更大的性能提升，標(biāo)志著在大規(guī)模推理任務(wù)中，預(yù)訓(xùn)練長(zhǎng)度擴(kuò)展取得了實(shí)質(zhì)性進(jìn)展。