MHA、MQA、GQA 都是深度學習中注意力機制的相關概念,其中 MHA 是標準的多頭注意力機制,MQA 和 GQA 則是其優化變體,以下是它們的區別、優缺點介紹:
區別
- MHA(多頭注意力):是 Transformer 架構的核心組件,會將輸入的查詢(Query)、鍵(Key)和值(Value)分別投影到多個不同的頭中,每個頭都擁有專屬的 Q、K、V 映射參數,它們獨立計算注意力,然后將多個頭的結果拼接后投影回原始維度。
- MQA(多查詢注意力):屬于 MHA 的優化方案,其所有的查詢頭共享同一組 K、V 參數,每個頭依然保持獨立的 Q,以此在保持查詢頭多樣性的同時,大幅減少內存占用。
- GQA(分組查詢注意力):是 MHA 和 MQA 的折中方案,它先把查詢頭劃分成若干個組,每組配備獨立的 K、V 參數,而組內所有查詢頭共享該組的 K、V 來計算注意力,最后將頭部輸出拼接后經線性層整合。
優缺點
- MHA
- 優點:能讓模型在不同的子空間中學習不同的注意力模式,提高了模型的表達能力,可使模型關注輸入序列的不同部分,并且支持并行計算,能在 GPU 上高效執行。
- 缺點:每個查詢頭都需要與多個鍵計算注意力,計算量較大,并且還需要存儲多個 Q、K、V 頭,大模型中顯存占用較多。
- MQA
- 優點:由于所有查詢頭共享 K、V,極大減少了計算成本,尤其適用于推理階段,也降低了內存占用,對處理超長文本較為友好。
- 缺點:僅有一個 K、V 組,可能會讓模型的表達能力受到影響,難以很好區分不同查詢頭的語義信息,上下文信息捕捉能力有所下降。
- GQA
- 優點:相比 MHA,減少了鍵值計算的開銷,計算效率有所提升,同時在視覺 Transformer 等計算機視覺任務里也能表現不錯,可降低視覺 Transformer 在圖像數據上的計算復雜度。
- 缺點:因查詢頭共享 K、V,會讓模型在靈活性上有一定損失,相對標準 MHA,其表達能力可能會降低。
備注:
在 MQA(多查詢注意力)中,K、V 的維度不會除以頭數?H?。
在標準的 MHA(多頭注意力)里,假設模型隱藏層維度是?dmodel?,注意力頭數為?H,通常會將?dmodel??均勻劃分到各個頭中,所以每個頭的 K、V 維度是?dmodel?/H,K、V 的張量形狀是 (batch_size, seq_len,?H,?dmodel?/H) 。但 MQA 中所有查詢頭共享同一組 K、V 參數,其 Key 與 Value 線性層輸出維度為 (batch_size, seq_len,?dmodel?),也就是不再需要按頭數對維度進行切分,K、V 本身的維度就是原始的模型維度,后續會通過廣播機制將它們擴展到和查詢頭數量相匹配的形狀以完成注意力計算 。
MQA減少計算量體現在兩方面:
1、減少多余投影的計算次數。
2、? 標準的MHA每個頭都擁有專屬的 Q、K、V 映射參數?每個頭的?Ki??是獨立存儲的 MQA 通過?K 的共享和廣播,減少了內存訪問的離散性,提高了硬件利用率。
3、減少對額外的每個頭的K V投影矩陣的梯度計算。