在深度學(xué)習(xí)領(lǐng)域，注意力機(jī)制已然成為現(xiàn)代大模型的核心基石。從最初的多頭注意力（MHA，Multi-Head Attention）到如今的多查詢注意力（MQA，Multi-Query Attention）、分組查詢注意力（GQA，Grouped-Query Attention），再到DeepSeek提出的創(chuàng)新性多頭潛在注意力（MLA，Multi-Head Latent Attention）方法，這一演變歷程不僅是技術(shù)發(fā)展的脈絡(luò)，更是對(duì)效率與性能極致追求的生動(dòng)寫照。本文將系統(tǒng)梳理這一發(fā)展脈絡(luò)，深入剖析MHA、MQA、GQA等變體的核心思路與實(shí)現(xiàn)方法。

一、Multi-Head Attention

多頭注意力（Multi-Head Attention，MHA）是Transformer模型架構(gòu)中的一個(gè)核心組件，它允許模型在處理輸入序列時(shí)能夠同時(shí)關(guān)注來自不同位置的不同表示子空間的信息。

MHA通過將輸入向量分割成多個(gè)并行的注意力“頭”，每個(gè)頭獨(dú)立地計(jì)算注意力權(quán)重并產(chǎn)生輸出，然后將這些輸出通過拼接和線性變換進(jìn)行合并以生成最終的注意力表示。

下面來看下計(jì)算公式：

1. 輸入變換：輸入序列首先通過三個(gè)不同的線性變換層，分別得到查詢（Query）、鍵（Key）和值（Value）矩陣。這些變換通常是通過全連接層實(shí)現(xiàn)的。

2. 分頭：將查詢、鍵和值矩陣分成多個(gè)頭（即多個(gè)子空間），每個(gè)頭具有不同的線性變換參數(shù)。

3. 注意力計(jì)算：對(duì)于每個(gè)頭，都執(zhí)行一次縮放點(diǎn)積注意力（Scaled Dot-Product Attention）運(yùn)算。具體來說，計(jì)算查詢和鍵的點(diǎn)積，經(jīng)過縮放、加上偏置后，使用softmax函數(shù)得到注意力權(quán)重。這些權(quán)重用于加權(quán)值矩陣，生成加權(quán)和作為每個(gè)頭的輸出。

4. 拼接與融合：將所有頭的輸出拼接在一起，形成一個(gè)長(zhǎng)向量。然后，對(duì)拼接后的向量進(jìn)行一個(gè)最終的線性變換，以整合來自不同頭的信息，得到最終的多頭注意力輸出。

作為最早提出的注意力機(jī)制方法，多頭注意力機(jī)制存在的問題：

1. 計(jì)算復(fù)雜度高：多頭注意力機(jī)制的計(jì)算復(fù)雜度與輸入序列長(zhǎng)度的平方成正比（），這使得在處理長(zhǎng)序列時(shí)計(jì)算量顯著增加。例如，對(duì)于長(zhǎng)度為1000的序列，計(jì)算復(fù)雜度將達(dá)到，這在實(shí)際應(yīng)用中可能導(dǎo)致訓(xùn)練和推理速度變慢。
2. 內(nèi)存占用大：在多頭注意力機(jī)制中，每個(gè)頭都需要獨(dú)立存儲(chǔ)查詢（Query）、鍵（Key）和值（Value）矩陣，這導(dǎo)致內(nèi)存消耗顯著增加。對(duì)于大規(guī)模模型，尤其是在長(zhǎng)序列任務(wù)中，KV緩存的大小會(huì)線性增長(zhǎng)，成為內(nèi)存瓶頸。
3. 特征冗余：多頭注意力機(jī)制中，不同頭可能學(xué)習(xí)到相似的特征，導(dǎo)致特征冗余。這種冗余不僅浪費(fèi)計(jì)算資源，還可能降低模型的泛化能力。
4. 模型解釋性差：多頭注意力機(jī)制的內(nèi)部工作機(jī)制較為復(fù)雜，每個(gè)頭的具體功能難以直觀理解，降低了模型的可解釋性。盡管可以通過注意力權(quán)重可視化來理解模型關(guān)注的輸入信息，但這種解釋性仍然有限。
5. 過擬合風(fēng)險(xiǎn)：由于多頭注意力機(jī)制增加了模型的參數(shù)量和復(fù)雜度，尤其是在數(shù)據(jù)量有限的情況下，模型可能會(huì)過度擬合訓(xùn)練數(shù)據(jù)。
6. 推理效率低：在自回歸模型中，每個(gè)解碼步驟都需要加載解碼器權(quán)重以及所有注意力的鍵和值，這不僅計(jì)算量大，還對(duì)內(nèi)存帶寬要求高。隨著模型規(guī)模的擴(kuò)大，這種開銷會(huì)進(jìn)一步增加，使得模型擴(kuò)展變得困難。

二、Multi-Query Attention

針對(duì)MHA存在的問題，Google提出了多查詢注意力（Multi-Query Attention，MQA）。MQA的設(shè)計(jì)初衷是為了在保持Transformer模型性能的同時(shí)，顯著提升計(jì)算效率和降低內(nèi)存占用。

在MHA中，輸入分別經(jīng)過的變換之后，都切成了n份（n=頭數(shù)），維度也從降到了，分別進(jìn)行attention計(jì)算再拼接。MQA的做法很簡(jiǎn)單，在線性變換之后，只對(duì)Q進(jìn)行切分（和MHA一樣），而K、V則直接在線性變換的時(shí)候把維度降到了（而不是切分變?。?，然后這n個(gè)Query頭分別和同一個(gè)K、V進(jìn)行attention計(jì)算，之后把結(jié)果拼接起來。

簡(jiǎn)單來說，就是MHA中，每個(gè)注意力頭的K、V不一樣，而MQA中每個(gè)注意力頭的K、V一樣，值共享，其他步驟和MHA一樣。

簡(jiǎn)單看下公式：

1. 查詢（Query）保持多頭設(shè)計(jì)：

2. 鍵（Key）和值（Value）共享一組矩陣：

3. 計(jì)算注意力輸出：

下圖是論文中MHA和MQA的對(duì)比結(jié)果，可以看到由于共享了多個(gè)頭的參數(shù)，限制了模型的表達(dá)能力，MQA雖然能好地支持推理加速，但是在效果上略比MHA差一點(diǎn)，但相比其他修改hidden size或者h(yuǎn)ead num的做法效果都好。

MQA通過共享鍵（K）和值（V）矩陣的設(shè)計(jì)，顯著降低了計(jì)算復(fù)雜度和內(nèi)存占用，同時(shí)保持了較好的性能表現(xiàn)。這種設(shè)計(jì)特別適合長(zhǎng)序列任務(wù)、資源受限的設(shè)備以及需要快速推理的場(chǎng)景。但是MQA對(duì)于所有query全部共享同一個(gè)key、value可能會(huì)限制每個(gè)查詢頭捕捉不同特征的能力，進(jìn)而影響模型的整體表達(dá)能力和靈活性。

三、Grouped-Query Attention

MQA對(duì)效果有點(diǎn)影響，MHA緩存又存不下，Google又繼續(xù)提出了一個(gè)折中的辦法組查詢注意力（Grouped-Query Attention，GQA），既能減少M(fèi)QA效果的損失，又相比MHA需要更少的緩存。

簡(jiǎn)單看下公式：

1. 將頭分為g組，每組有h/g個(gè)頭。對(duì)于每組i：

2. 計(jì)算每個(gè)組的注意力輸出并拼接：

來看下結(jié)果：

看表中2/3/4行對(duì)比，GQA的速度相比MHA有明顯提升，而效果上比MQA也好一些，能做到和MHA基本沒差距。文中提到，這里的MQA和GQA都是通過average pooling從MHA初始化而來，然后進(jìn)行了少量的訓(xùn)練得到的。

下面是Llama2技術(shù)報(bào)告中做的MHA、MQA、GQA效果對(duì)比，可以看到效果確實(shí)很不錯(cuò)。

四、Multi-Head Latent Attention

在最新的DeepSeek論文中，為解決MHA在高計(jì)算成本和KV緩存方面的局限性，提出改進(jìn)的多頭潛在注意力（Multi-Head Latent Attention，MLA），旨在提高Transformer模型在處理長(zhǎng)序列時(shí)的效率和性能。

MLA的技術(shù)創(chuàng)新主要是采用低秩聯(lián)合壓縮鍵值技術(shù)，優(yōu)化鍵值（KV）矩陣，顯著減少了內(nèi)存消耗并提高了推理效率。

具體來說，MLA通過低秩聯(lián)合壓縮鍵值（Key-Value），將它們壓縮為一個(gè)潛在向量（latent vector），從而大幅減少所需的緩存容量，還降低了計(jì)算復(fù)雜度。在推理階段，MHA需要緩存獨(dú)立的鍵（Key）和值（Value）矩陣，這會(huì)增加內(nèi)存和計(jì)算開銷。而MLA通過低秩矩陣分解技術(shù)，顯著減小了存儲(chǔ)的KV（Key-Value）的維度，從而降低了內(nèi)存占用。

MLA利用低秩壓縮技術(shù)，使得DeepSeek的KV緩存減少了93.3%。來看下公式：

1. KV聯(lián)合低秩壓縮

MLA模型通過低秩壓縮對(duì)鍵（keys）和值（values）進(jìn)行聯(lián)合壓縮，以減少KV緩存的大小。其核心公式為：

• 是key和value的壓縮潛在向量，表示KV壓縮維度；
• 是降維投影矩陣，分別是key和value的升維投影矩陣。

2. Q的低秩壓縮

為了降低訓(xùn)練時(shí)的激活內(nèi)存占用，MLA對(duì)query（Q）也進(jìn)行了低秩壓縮。其核心公式為：

• 是query的壓縮潛在向量，表示query壓縮的維度；
• 是query的降維投影矩陣，是query的升維投影矩陣。

3. RoPE 與低秩 KV 壓縮不兼容問題-解耦 RoPE 策略

RoPE 對(duì)keys和queries都是位置敏感的。如果對(duì)鍵應(yīng)用 RoPE，則會(huì)有一個(gè)與位置相關(guān)的 RoPE 矩陣。 這種情況下，在推理過程中不能再被吸收進(jìn)，因?yàn)橐粋€(gè)與當(dāng)前生成的 token 相關(guān)的 RoPE 矩陣會(huì)存在于之間，矩陣乘法不遵循交換律。因此，必須在推理過程中重新計(jì)算所有前綴 token 的鍵，這將顯著阻礙推理效率。 為了解決這個(gè)問題，論文提出了解耦 RoPE 策略，該策略使用額外的多查詢注意力（MQA）和共享的鍵來攜帶 RoPE，其中表示解耦query和key的head_dim。

在解耦 RoPE 策略下，MLA 執(zhí)行以下計(jì)算：

• 其中和是分別生成解耦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聯(lián)合以及q對(duì)應(yīng)的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進(jìn)行MQA計(jì)算的映射矩陣
        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對(duì)應(yīng)的rope類，因?yàn)轭^的數(shù)量不同，初始化2個(gè)實(shí)例
        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 :低秩轉(zhuǎn)換
        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進(jìn)行MQA計(jì)算，同時(shí)引入ROPE
        #q_t_r,k_t_r施加rope時(shí)均擴(kuò)展了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,進(jìn)行sdpa計(jì)算
        #q_t_c擴(kuò)展出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擴(kuò)展出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計(jì)算，需要在n_h_k_r維度復(fù)制
        #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)  
        # 注意力計(jì)算,[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計(jì)算，擴(kuò)展出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,送入最終統(tǒng)一映射層
        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與其他幾個(gè)注意力機(jī)制對(duì)比的實(shí)驗(yàn)結(jié)果，但是結(jié)果導(dǎo)向來看，MLA的KV緩存大幅減少，大幅提高模型推理速度，在減少資源消耗的同時(shí)，保持甚至提升模型性能。

五、總結(jié)

從MHA到MQA、GQA，再到MLA，注意力機(jī)制的演變展示了在效率與性能之間不斷優(yōu)化的軌跡。MLA通過創(chuàng)新的KV緩存壓縮和恢復(fù)機(jī)制，實(shí)現(xiàn)了在資源消耗、推理速度和模型性能之間的最佳平衡，為大語言模型的高效部署和應(yīng)用提供了新的可能性。

[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