極客號(hào)（Daydx.com）12月7日 消息:田淵棟團(tuán)隊(duì)最新發(fā)表的論文解決了大型語言模型在實(shí)際部署中遇到的內(nèi)存和輸入長(zhǎng)度限制的問題，將推理系統(tǒng)的吞吐量提高了近30倍。論文提出了一種實(shí)現(xiàn)KV緩存的新方法，通過識(shí)別和保留重要的tokens，顯著減少了內(nèi)存占用，并在長(zhǎng)輸入序列的任務(wù)中表現(xiàn)良好。
這篇論文的研究對(duì)象是大型語言模型（LLM），該模型在實(shí)際部署中面臨著KV緩存成本昂貴和對(duì)長(zhǎng)序列的泛化能力差的問題。為了解決這些問題，田淵棟團(tuán)隊(duì)提出了一種基于重要tokens的KV緩存逐出策略，通過識(shí)別并保留對(duì)生成結(jié)果有重要貢獻(xiàn)的tokens，顯著減少了內(nèi)存占用，并提高了模型在處理長(zhǎng)輸入序列時(shí)的表現(xiàn)。
論文地址:https://arxiv.org/pdf/2306.14048.pdf
代碼地址:https://github.com/FMInference/H2O
在實(shí)驗(yàn)中，作者使用了OPT、LLaMA和GPT-NeoX等模型驗(yàn)證了他們提出的方法的準(zhǔn)確性和有效性。實(shí)驗(yàn)結(jié)果顯示，通過使用該方法，DeepSpeed Zero-Inference、Hugging Face Accelerate和FlexGen這三個(gè)推理系統(tǒng)的吞吐量分別提高了29倍、29倍和3倍，且在相同的批量大小下，延遲最多可以減少1.9倍。
通過研究發(fā)現(xiàn)，大部分注意力鍵和值嵌入在生成過程中貢獻(xiàn)較少的價(jià)值，只有一小部分tokens貢獻(xiàn)了大部分的價(jià)值。基于這個(gè)發(fā)現(xiàn)，作者提出了一種基于重要tokens的KV緩存逐出策略，動(dòng)態(tài)保持最近的tokens和重要tokens的平衡。通過使用這種策略，可以顯著減少KV緩存的大小，從而降低了內(nèi)存占用，并提高了模型的推理效率。
綜上所述，田淵棟團(tuán)隊(duì)的最新論文成功解決了大型語言模型在實(shí)際部署中的難題，通過優(yōu)化KV緩存的實(shí)現(xiàn)方法，將推理系統(tǒng)的吞吐量提高了近30倍。這一成果在NeurIPS’23上將進(jìn)行展示，對(duì)于大型語言模型的部署和應(yīng)用具有重要的意義。
