《Transformer 關鍵要素:非所有注意力皆必要》論文解讀
摘要:隨著基于Transformer的大型語言模型(LLMs)在各種任務中展現出強大性能,其規模的擴大也帶來了冗余結構的問題。本文通過基于相似性的度量方法,研究了Transformer中不同模塊(包括塊、多層感知機(MLP)和注意力(Attention)層)的冗余情況。發現大量注意力層存在過高的相似性,可在不降低性能的情況下進行修剪,從而降低內存和計算成本。還提出了一種聯合丟棄注意力和MLP層的方法,提高了性能和丟棄比例。
一、研究背景和動機
(一)LLMs的發展與挑戰Transformer-basedLLMs在AI研究中取得了顯著進展,但模型的擴展導致冗余模塊增加,這不僅增加了部署成本和資源需求,還限制了其在邊緣設備上的應用。例如GPT-3的參數規模巨大,對內存要求過高。
(二)對模型結構冗余性的關注以往的研究雖有優化LLM效率的工作,但常忽略Transformer架構的獨特性。Transformer由多個堆疊的塊組成,每個塊包含MLP層和Attention層,不同模塊可能存在不同類型和程度的冗余,這激發了作者對這些不同模塊冗余性的研究。
二、研究方法
(一)基于相似性的度量方法通過計算模塊輸入和輸出之間的余弦相似性來確定模塊的重要性。相似性越高,模塊的重要性得分越低,越有可能被丟棄。公式為,其中為輸入,為輸出。
(二)不同模塊的丟棄方法塊丟棄(BlockDrop)計算每個塊的輸入和輸出的相似性得分,從淺層到深層迭代計算每個塊的重要性得分,丟棄得分最低的塊。層丟棄(LayerDrop)MLP丟棄(MLPDrop):考慮MLP層及其后續的層歸一化(LayerNorm)作為一個整體來計算相似性得分,同時丟棄不重要的MLP層和相關的LayerNorm層。注意力丟棄(AttentionDrop):類似MLPDrop,考慮注意力層和相關LayerNorm層的整體輸出,計算重要性得分。聯合層丟棄(JointLayerDrop)先分別獲取注意力層和MLP層的重要性得分和,然后將它們連接起來,根據綜合得分丟棄層,以優化丟棄過程。
三、實驗結果
(一)不同模塊丟棄的比較塊丟棄在Llama-2-13B上進行實驗,丟棄8個塊時,平均性能分別下降了7.5%,性能下降明顯,說明塊對于維持性能較為重要。MLP層丟棄同樣在Llama-2-13B和Mistral-7B上實驗,丟棄8個MLP層時,性能分別下降了6.3%和6.9%,也會導致性能大幅下降。注意力層丟棄與前兩者不同,在Llama-3-13B和Mistral-7B上丟棄8個注意力層后,模型仍能保持超過99%的原始性能。而且在不同的丟棄比例實驗中,當丟棄比例低于50%時,性能保持相對穩定。
(二)模塊丟棄的效率提升內存使用注意力層丟棄可以有效減少鍵值(KV)緩存,例如對于Llama-2-13B的一個輸入,可將約70GB的KV緩存減少一半,顯著降低內存需求。速度提升與MLP層相比,注意力層丟棄對速度提升的貢獻更大,如在Mistral-7B上丟棄8層時,注意力層丟棄的速度提升倍數為1.17×,而MLP層丟棄為1.09×。
(三)聯合層丟棄的效果聯合層丟棄的性能明顯高于單獨的注意力層丟棄或MLP層丟棄。它首先只丟棄注意力層,直到達到一定數量后再考慮丟棄MLP層,提高了丟棄比例和性能。
四、進一步分析
(一)注意力層冗余的一致性訓練過程中的相似性通過評估MAP-Neo-7B在不同訓練階段的檢查點,發現注意力層在所有訓練階段的重要性得分都極低,而MLP層和塊的重要性得分隨著訓練逐漸增加。更深層模塊的冗余性可視化不同模型的丟棄層或塊,發現無論是Llama-2-13B、Mistral-7B還是更大的Llama-2-70B,都傾向于先丟棄更深層的模塊,說明更深層模塊往往更冗余。
(二)塊丟棄和層丟棄的魯棒性對樣本數量的魯棒性隨著樣本數量增加(從4到1024),不同注意力層的特征相似性保持相對穩定,說明使用256個樣本計算相似性足以滿足實驗需求。對校準數據集的魯棒性改變校準數據集(從預訓練數據集到指令調整數據集),特征相似性沒有顯著變化,模塊的丟棄情況也基本相同,證明了塊丟棄和層丟棄對校準數據集變化的適應性。
(三)不同模型規模的影響更大模型的魯棒性在Llama-2-70B上進行實驗,發現它對塊丟棄和MLP層丟棄也敏感,但相比較小模型更具魯棒性。例如,丟棄8個MLP層時,Llama-2-70B平均性能僅下降0.6%,而Llama-2-13B丟棄4個MLP層時下降3.9%。注意力層丟棄在更大模型上的表現在Llama-2-70B上丟棄32個注意力層時,性能平均僅下降0.1%,而在Llama-2-13B上丟棄相同比例的注意力層會導致明顯的性能下降。五、研究貢獻開發了基于相似性的度量方法來衡量LLMs內的冗余,并研究了包括塊、MLP和注意力層在內的各種模塊的結構冗余。揭示了注意力層令人驚訝的冗余性,并證明了注意力層丟棄對準確性、內存和計算效率的好處。提出了聯合層丟棄等技術,進一步提高了模塊丟棄的有效性和效率。所提出的層丟棄方法能有效識別重要和冗余層,為現有模型評估提供了關鍵信息,也為網絡架構設計的未來研究提供了有價值的見解。六、研究局限所討論的丟棄技術雖在研究的模型上提高了效率,但還需進一步探索其在更廣泛模型(如視覺Transformer和視覺-語言模型)中的應用。主要關注訓練后的丟棄,未涉及重新訓練,而重新訓練可能進一步恢復或提高性能。更多內容歡迎來卡奧斯智能交互引擎搜索探討
