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