麻省理工提出“跨層注意力”，極大優(yōu)化Transformer緩存
Transformer在大模型領(lǐng)域的影響力不言而喻，ChatGPT、Sora、Midjourney、Suno、Llama、Stable Difusion等幾乎所有知名開閉源模型，皆基于該架構(gòu)開發(fā)而成。
但隨著大模型參數(shù)呈指數(shù)級增長，小的幾百億大的上千億甚至萬億，這使得Transformer在解碼時所需的KV（鍵值）緩存急劇增加，會導(dǎo)致內(nèi)存占用過大造成部署、推理方面的瓶頸。
所以，麻省理工的研究人員提出了全新的跨層注意力（Cross-Layer Attention, 簡稱“CLA”），通過在不同解碼層間共享KV來顯著降低對內(nèi)存的使用，從而提升大模型在處理長序列和大批次推理任務(wù)時的效率以及準(zhǔn)確率。
論文地址：https://arxiv.org/abs/2405.12981
目前，解決Transformer緩存瓶頸的方法是多查詢注意力和分組查詢注意力， CLA則是在這兩者之上做了進(jìn)一步優(yōu)化。
在傳統(tǒng)的Transformer架構(gòu)中，每一層都包含獨立的KV投影，用于生成鍵值對；查詢投影則負(fù)責(zé)生成查詢向量；注意力模塊可根據(jù)查詢和KV對計算注意力分布；
輸出投影進(jìn)行整合注意力輸出；殘差連接和層歸一化，主要用來確保學(xué)習(xí)穩(wěn)定性和表現(xiàn)力。
?
CLA的核心思想是通過在相鄰層之間共享KV來降低內(nèi)存占用，而不是每層獨立計算和存儲。
在CLA中不同層之間的通信也是通過共享KV激活來完成的，這極大減少了模型必須維護(hù)的獨立KV集合的數(shù)量，而其他層則通過層間連接重用這些激活。
此外，CLA提供了靈活的配置允許開發(fā)者根據(jù)需要調(diào)整共享因子，以平衡內(nèi)存使用和模型性能。共享因子定義了每個KV投影被多少層共享。例如，在CLA2配置中，每對相鄰層共享一個KV緩存。
為了測試CLA的性能，研究人員在10億和30億參數(shù)上訓(xùn)練了多款CLA和非CLA模型，目的是在固定內(nèi)存預(yù)算下找到CLA的最佳準(zhǔn)確性。這些模型的變化涵蓋了從MHA到GQA再到MQA的范圍，KV緩存的內(nèi)存需求也由此跨越了兩個數(shù)量級。
為了確保結(jié)果的穩(wěn)健性，研究人員對選定的幾個模型進(jìn)行了學(xué)習(xí)率調(diào)整實驗，以確認(rèn)CLA在與經(jīng)過良好調(diào)整學(xué)習(xí)率的基線模型相比時是否具有優(yōu)勢。
實驗結(jié)果顯示，CLA在減少KV緩存大小的同時，能夠?qū)崿F(xiàn)與非CLA模型相當(dāng)?shù)耐评頊?zhǔn)確性。在10億參數(shù)規(guī)模的實驗中，CLA模型成功地在保持準(zhǔn)確度的基礎(chǔ)上顯著減少了KV緩存的大小。
特別是MQA-CLA2配置在減少KV緩存大小方面表現(xiàn)非常出色，相比基線MQA模型，在相同的KV緩存內(nèi)存下實現(xiàn)了更低的驗證困惑度。
在30億參數(shù)模型的實驗中，CLA也展現(xiàn)了其高效性。CLA2配置在減少KV緩存大小的同時，與相同頭維度的MQA基線模型相比，實現(xiàn)了更低的驗證困惑度，進(jìn)一步證明了CLA在大規(guī)模模型中的有效性。
本文轉(zhuǎn)自 AIGC開放社區(qū) ，作者： AIGC開放社區(qū)
