深度學習中的注意力機制革命：MHA、MQA、GQA至DeepSeek MLA的演變
在深度學習領域，注意力機制已然成為現代大模型的核心基石。從最初的多頭注意力（MHA，Multi-Head Attention）到如今的多查詢注意力（MQA，Multi-Query Attention）、分組查詢注意力（GQA，Grouped-Query Attention），再到DeepSeek提出的創新性多頭潛在注意力（MLA，Multi-Head Latent Attention）方法，這一演變歷程不僅是技術發展的脈絡，更是對效率與性能極致追求的生動寫照。本文將系統梳理這一發展脈絡，深入剖析MHA、MQA、GQA等變體的核心思路與實現方法。
一、Multi-Head Attention
多頭注意力（Multi-Head Attention，MHA）是Transformer模型架構中的一個核心組件，它允許模型在處理輸入序列時能夠同時關注來自不同位置的不同表示子空間的信息。
MHA通過將輸入向量分割成多個并行的注意力“頭”，每個頭獨立地計算注意力權重并產生輸出，然后將這些輸出通過拼接和線性變換進行合并以生成最終的注意力表示。
下面來看下計算公式：
1. 輸入變換：輸入序列首先通過三個不同的線性變換層，分別得到查詢（Query）、鍵（Key）和值（Value）矩陣。這些變換通常是通過全連接層實現的。
2. 分頭：將查詢、鍵和值矩陣分成多個頭（即多個子空間），每個頭具有不同的線性變換參數。
3. 注意力計算：對于每個頭，都執行一次縮放點積注意力（Scaled Dot-Product Attention）運算。具體來說，計算查詢和鍵的點積，經過縮放、加上偏置后，使用softmax函數得到注意力權重。這些權重用于加權值矩陣，生成加權和作為每個頭的輸出。
4. 拼接與融合：將所有頭的輸出拼接在一起，形成一個長向量。然后，對拼接后的向量進行一個最終的線性變換，以整合來自不同頭的信息，得到最終的多頭注意力輸出。
作為最早提出的注意力機制方法，多頭注意力機制存在的問題：
- 計算復雜度高：多頭注意力機制的計算復雜度與輸入序列長度的平方成正比（
），這使得在處理長序列時計算量顯著增加。例如，對于長度為1000的序列，計算復雜度將達到
，這在實際應用中可能導致訓練和推理速度變慢。 - 內存占用大：在多頭注意力機制中，每個頭都需要獨立存儲查詢（Query）、鍵（Key）和值（Value）矩陣，這導致內存消耗顯著增加。對于大規模模型，尤其是在長序列任務中，KV緩存的大小會線性增長，成為內存瓶頸。
- 特征冗余：多頭注意力機制中，不同頭可能學習到相似的特征，導致特征冗余。這種冗余不僅浪費計算資源，還可能降低模型的泛化能力。
- 模型解釋性差：多頭注意力機制的內部工作機制較為復雜，每個頭的具體功能難以直觀理解，降低了模型的可解釋性。盡管可以通過注意力權重可視化來理解模型關注的輸入信息，但這種解釋性仍然有限。
- 過擬合風險：由于多頭注意力機制增加了模型的參數量和復雜度，尤其是在數據量有限的情況下，模型可能會過度擬合訓練數據。
- 推理效率低：在自回歸模型中，每個解碼步驟都需要加載解碼器權重以及所有注意力的鍵和值，這不僅計算量大，還對內存帶寬要求高。隨著模型規模的擴大，這種開銷會進一步增加，使得模型擴展變得困難。
二、Multi-Query Attention
針對MHA存在的問題，Google提出了多查詢注意力（Multi-Query Attention，MQA）。MQA的設計初衷是為了在保持Transformer模型性能的同時，顯著提升計算效率和降低內存占用。
在MHA中，輸入分別經過
的變換之后，都切成了n份（n=頭數），維度也從
降到了
，分別進行attention計算再拼接。MQA的做法很簡單，在線性變換之后，只對Q進行切分（和MHA一樣），而K、V則直接在線性變換的時候把維度降到了
（而不是切分變小），然后這n個Query頭分別和同一個K、V進行attention計算，之后把結果拼接起來。
簡單來說，就是MHA中，每個注意力頭的K、V不一樣，而MQA中每個注意力頭的K、V一樣，值共享，其他步驟和MHA一樣。
簡單看下公式：
1. 查詢（Query）保持多頭設計：
2. 鍵（Key）和值（Value）共享一組矩陣：
3. 計算注意力輸出：
下圖是論文中MHA和MQA的對比結果，可以看到由于共享了多個頭的參數，限制了模型的表達能力，MQA雖然能好地支持推理加速，但是在效果上略比MHA差一點，但相比其他修改hidden size或者head num的做法效果都好。
MQA通過共享鍵（K）和值（V）矩陣的設計，顯著降低了計算復雜度和內存占用，同時保持了較好的性能表現。這種設計特別適合長序列任務、資源受限的設備以及需要快速推理的場景。但是MQA對于所有query全部共享同一個key、value可能會限制每個查詢頭捕捉不同特征的能力，進而影響模型的整體表達能力和靈活性。
三、Grouped-Query Attention
MQA對效果有點影響，MHA緩存又存不下，Google又繼續提出了一個折中的辦法組查詢注意力（Grouped-Query Attention，GQA），既能減少MQA效果的損失，又相比MHA需要更少的緩存。
簡單看下公式：
1. 將頭分為g組，每組有h/g個頭。對于每組i：
2. 計算每個組的注意力輸出并拼接：
來看下結果：
看表中2/3/4行對比，GQA的速度相比MHA有明顯提升，而效果上比MQA也好一些，能做到和MHA基本沒差距。文中提到，這里的MQA和GQA都是通過average pooling從MHA初始化而來，然后進行了少量的訓練得到的。
下面是Llama2技術報告中做的MHA、MQA、GQA效果對比，可以看到效果確實很不錯。
四、Multi-Head Latent Attention
在最新的DeepSeek論文中，為解決MHA在高計算成本和KV緩存方面的局限性，提出改進的多頭潛在注意力（Multi-Head Latent Attention，MLA），旨在提高Transformer模型在處理長序列時的效率和性能。
MLA的技術創新主要是采用低秩聯合壓縮鍵值技術，優化鍵值（KV）矩陣，顯著減少了內存消耗并提高了推理效率。
具體來說，MLA通過低秩聯合壓縮鍵值（Key-Value），將它們壓縮為一個潛在向量（latent vector），從而大幅減少所需的緩存容量，還降低了計算復雜度。在推理階段，MHA需要緩存獨立的鍵（Key）和值（Value）矩陣，這會增加內存和計算開銷。而MLA通過低秩矩陣分解技術，顯著減小了存儲的KV（Key-Value）的維度，從而降低了內存占用。
MLA利用低秩壓縮技術，使得DeepSeek的KV緩存減少了93.3%。來看下公式：
1. KV聯合低秩壓縮
MLA模型通過低秩壓縮對鍵（keys）和值（values）進行聯合壓縮，以減少KV緩存的大小。其核心公式為：
是key和value的壓縮潛在向量，
表示KV壓縮維度；
是降維投影矩陣，
分別是key和value的升維投影矩陣。
2. Q的低秩壓縮
為了降低訓練時的激活內存占用，MLA對query（Q）也進行了低秩壓縮。其核心公式為：
是query的壓縮潛在向量，
表示query壓縮的維度；
是query的降維投影矩陣，
是query的升維投影矩陣。
3. RoPE 與低秩 KV 壓縮不兼容問題-解耦 RoPE 策略
RoPE 對keys和queries都是位置敏感的。如果對鍵
應用 RoPE，則會有一個與位置相關的 RoPE 矩陣。 這種情況下，
在推理過程中不能再被吸收進
，因為一個與當前生成的 token 相關的 RoPE 矩陣會存在于
之間，矩陣乘法不遵循交換律。因此，必須在推理過程中重新計算所有前綴 token 的鍵，這將顯著阻礙推理效率。 為了解決這個問題，論文提出了解耦 RoPE 策略，該策略使用額外的多查詢注意力（MQA）
和共享的鍵
來攜帶 RoPE，其中
表示解耦query和key的head_dim。
在解耦 RoPE 策略下，MLA 執行以下計算：
- 其中
和
是分別生成解耦query和key的矩陣；
下面附上代碼：
import torch
import torch.nn as nn
import math
class MLA(nn.Module):
def __init__(self, d_model=512, down_dim=128, up_dim=256, num_heads=8, rope_head_dim=26, dropout_prob=0.1):
super(MLA, self).__init__()
self.d_model = d_model
self.down_dim = down_dim
self.up_dim = up_dim
self.num_heads = num_heads
self.head_dim = d_model // num_heads
self.rope_head_dim = rope_head_dim
self.v_head_dim = up_dim // num_heads
# 初始化kv聯合以及q對應的dow,up projection
self.down_proj_kv = nn.Linear(d_model, down_dim) # W^{DKV}
self.up_proj_k = nn.Linear(down_dim, up_dim)# W^{UK}
self.up_proj_v = nn.Linear(down_dim, up_dim) # W^{UV}
self.down_proj_q = nn.Linear(d_model, down_dim) #W^{DQ}
self.up_proj_q = nn.Linear(down_dim, up_dim) # W^{UQ}
# 初始化解耦的q,k進行MQA計算的映射矩陣
self.proj_qr = nn.Linear(down_dim, rope_head_dim * num_heads)
self.proj_kr = nn.Linear(d_model, rope_head_dim*1)
#初始化解耦的q,k對應的rope類，因為頭的數量不同，初始化2個實例
self.rope_q = RotaryEmbedding(rope_head_dim * num_heads, num_heads)
self.rope_k = RotaryEmbedding(rope_head_dim, 1)
# Dropout and final linear layer
self.dropout = nn.Dropout(dropout_prob)
self.fc = nn.Linear(num_heads * self.v_head_dim, d_model)
self.res_dropout = nn.Dropout(dropout_prob)
def forward(self, h, mask=None):
bs, seq_len, _ = h.size()
# setp1 :低秩轉換
c_t_kv = self.down_proj_kv(h)
k_t_c = self.up_proj_k(c_t_kv)
v_t_c = self.up_proj_v(c_t_kv)
c_t_q = self.down_proj_q(h)
q_t_c = self.up_proj_q(c_t_q)
#step2:解耦的q,k進行MQA計算，同時引入ROPE
#q_t_r,k_t_r施加rope時均擴展了n_h_r維度->[bs,n_h_r,seq_len,rope_head_dim]
q_t_r = self.rope_q(self.proj_qr(c_t_q))
k_t_r = self.rope_k(self.proj_kr(h))
#step3:拼接step1，step2得到的q,k,進行sdpa計算
#q_t_c擴展出num_heads為4維，以便于和q_t_r拼接
q_t_c = q_t_c.reshape(bs, seq_len, self.num_heads, -1).transpose(1, 2)
#head_dim,rope_head_dim拼接
q = torch.cat([q_t_c, q_t_r], dim=-1)
#k_t_c擴展出num_heads為4維，以便于和k_t_r拼接
k_t_c = k_t_c.reshape(bs, seq_len, self.num_heads, -1).transpose(1, 2)
#k_t_r為MQA,n_h_k_r=1,為了和q_t_r計算，需要在n_h_k_r維度復制
#k_t_r:[bs,n_h_r_k,seq_len,rope_head_dim]->[bs,num_heads,seq_len,rope_head_dim]
k_t_r=k_t_r.repeat(1,self.num_heads,1,1)
#head_dim,rope_head_dim拼接
k = torch.cat([k_t_c, k_t_r], dim=-1)
# 注意力計算,[bs,num_heads,seq_len,seq_len]
scores = torch.matmul(q, k.transpose(-1, -2))
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
scores = torch.softmax(scores / (math.sqrt(self.head_dim) + math.sqrt(self.rope_head_dim)), dim=-1)
scores = self.dropout(scores)
#v_t_c和scores計算，擴展出num_heads維度
v_t_c = v_t_c.reshape(bs, seq_len, self.num_heads, self.v_head_dim).transpose(1, 2)
output = torch.matmul(scores, v_t_c)
#壓縮num_head,送入最終統一映射層
output = output.transpose(1, 2).reshape(bs, seq_len, -1)
output = self.fc(output)
output = self.res_dropout(output)
return output
bs, seq_len, d_model = 4, 10, 512
h = torch.randn(bs, seq_len, d_model)
mla = MLA(d_model=d_model)
output = mla(h)DeepSeek沒有給出MLA與其他幾個注意力機制對比的實驗結果，但是結果導向來看，MLA的KV緩存大幅減少，大幅提高模型推理速度，在減少資源消耗的同時，保持甚至提升模型性能。
五、總結
從MHA到MQA、GQA，再到MLA，注意力機制的演變展示了在效率與性能之間不斷優化的軌跡。MLA通過創新的KV緩存壓縮和恢復機制，實現了在資源消耗、推理速度和模型性能之間的最佳平衡，為大語言模型的高效部署和應用提供了新的可能性。
[1]MHA: Attention Is All You Need(https://arxiv.org/pdf/1706.03762)
[2]MQA: Fast Transformer Decoding: One Write-Head is All You Need(https://arxiv.org/pdf/1911.02150)
[3]GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints(https://arxiv.org/pdf/2305.13245)
[4]MLA:https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf
