DeepSeek中的多頭潛在注意力(MLA)淺嘗 原創(chuàng)
MLA是MHA的變體,因此先來看看MHA。
MHA(多頭注意力)
MHA通過將輸入向量分割成多個并行的注意力“頭”,每個頭獨立地計算注意力權重并產(chǎn)生輸出,然后將這些輸出通過拼接和線性變換進行合并以生成最終的注意力表示。
Transformer 編碼器塊內(nèi)的縮放點積注意力機制和多頭注意力機制
MHA計算過程
MHA 能夠理解輸入不同部分之間的關系。然而,這種復雜性是有代價的——對內(nèi)存帶寬的需求很大,尤其是在解碼器推理期間。主要問題的關鍵在于內(nèi)存開銷。在自回歸模型中,每個解碼步驟都需要加載解碼器權重以及所有注意鍵和值。這個過程不僅計算量大,而且內(nèi)存帶寬也大。隨著模型規(guī)模的擴大,這種開銷也會增加,使得擴展變得越來越艱巨。
MLA(多頭潛在注意力)
概念:
- 多頭注意力機制:Transformer 的核心模塊,能夠通過多個注意力頭并行捕捉輸入序列中的多樣化特征。
- 潛在表示學習:通過將高維輸入映射到低維潛在空間,可以提取更抽象的語義特征,同時有效減少計算復雜度。
問題:
1.效率問題:傳統(tǒng)多頭注意力的計算復雜度為O(n2d),即隨著序列長度的增長,鍵值(Key-Value,KV)緩存的大小也會線性增加,這給模型帶來了顯著的內(nèi)存負擔。) 2.表達能力瓶頸:難以充分捕捉復雜全局依賴。MLA 通過在潛在空間中執(zhí)行注意力計算,有效降低復雜度,同時提升建模能力。
MLA 的提出:MLA 將多頭注意力機制 與 潛在表示學習 相結合,解決MHA在高計算成本和KV緩存方面的局限性。
MLA的具體做法(創(chuàng)新點): 采用低秩聯(lián)合壓縮鍵值技術,優(yōu)化了鍵值(KV)矩陣,顯著減少了內(nèi)存消耗并提高了推理效率。
如上圖,在MHA、GQA中大量存在于keys values中的KV緩存——帶陰影表示,到了MLA中時,只有一小部分的被壓縮Compressed的Latent KV了。
并且,在推理階段,MHA需要緩存獨立的鍵(Key)和值(Value)矩陣,這會增加內(nèi)存和計算開銷。而MLA通過低秩矩陣分解技術,顯著減小了存儲的KV(Key-Value)的維度,從而降低了內(nèi)存占用。
MLA的核心步驟:
- 輸入映射->潛在空間
- 潛在空間中的多頭注意力計算
- 映射回原始空間
將多頭注意力結果從潛在空間映射回原始空間:
參考文獻:https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
本文轉(zhuǎn)載自公眾號大模型自然語言處理 作者:余俊暉
原文鏈接:??https://mp.weixin.qq.com/s/CSK1t9o82un5QDphhOP6fQ??
