字節(jié)跳動(dòng)旗下的Seed團(tuán)隊(duì)近期在人工智能領(lǐng)域取得了新的突破，他們研發(fā)出一種名為PHD-Transformer（Parallel Hidden Decoding Transformer）的新型模型，該模型在預(yù)訓(xùn)練長度擴(kuò)展方面展現(xiàn)出了非凡的能力，有效解決了推理過程中KV緩存膨脹的難題。隨著大型推理模型的快速發(fā)展，如何在后訓(xùn)練階段生成更長的推理鏈，并在復(fù)雜任務(wù)上取得優(yōu)異表現(xiàn)，成為了研究人員關(guān)注的焦點(diǎn)。Seed團(tuán)隊(duì)受此啟發(fā)，決定在預(yù)訓(xùn)練階段就進(jìn)行長度擴(kuò)展的探索。
傳統(tǒng)方法在處理長度擴(kuò)展時(shí)，通常會(huì)在序列中插入額外的文本或潛在向量，但這種方式會(huì)導(dǎo)致KV緩存占用大量內(nèi)存，從而影響推理速度。而PHD-Transformer則采用了更為高效的策略，即直接重復(fù)輸入的tokens。然而，這種方法雖然提升了訓(xùn)練損失和模型性能，但也帶來了新的問題，如KV緩存的線性增長、內(nèi)存壓力的增加以及解碼延遲的加劇。
為了克服這些挑戰(zhàn)，PHD-Transformer引入了創(chuàng)新的KV緩存管理策略。在推理過程中，該模型只保留由原始tokens生成的KV緩存，對(duì)于重復(fù)的tokens，則在預(yù)測完成后立即丟棄。這一策略顯著提升了推理速度。研究團(tuán)隊(duì)還開發(fā)了一種稱為PHD-SWA（Sliding Window Attention）的滑動(dòng)窗口注意力機(jī)制，以保持局部滑動(dòng)窗口緩存的性能優(yōu)勢。為了進(jìn)一步優(yōu)化預(yù)填充時(shí)間，他們提出了PHD-CSWA（Chunk-wise Sliding Window Attention）機(jī)制，通過限制每個(gè)塊內(nèi)的順序依賴，大幅縮短了預(yù)填充時(shí)間。
在一系列嚴(yán)格的實(shí)驗(yàn)中，PHD-CSWA展現(xiàn)出了卓越的性能。在多個(gè)公開的基準(zhǔn)測試集上，該模型均實(shí)現(xiàn)了準(zhǔn)確率的提升。Seed團(tuán)隊(duì)表示，PHD-CSWA在保持原有高效性的基礎(chǔ)上，為模型帶來了更大的性能飛躍，這標(biāo)志著在大規(guī)模推理任務(wù)中，預(yù)訓(xùn)練長度擴(kuò)展技術(shù)取得了實(shí)質(zhì)性的進(jìn)步。
