「DeepSeek-V3 技術(shù)解析」：多頭潛在注意力機(jī)制（MLA） 原創(chuàng)
編者按： 模型在生成長序列文本時，鍵值緩存占用的顯存是否讓 GPU 不堪重負(fù)？如何在保持模型性能的同時有效降低大語言模型推理階段的內(nèi)存占用？
每一次模型推理，大量重復(fù)計算和龐大的內(nèi)存占用不僅大幅增加了運(yùn)營成本，還嚴(yán)重限制了模型的批處理能力和響應(yīng)速度。多查詢注意力機(jī)制（MQA）和分組查詢注意力機(jī)制（GQA）雖能部分緩解問題，但往往以犧牲模型性能為代價。
DeepSeek 團(tuán)隊在模型優(yōu)化過程中開發(fā)的多頭潛在注意力機(jī)制（MLA）為資源受限場景提供了可行性方案。通過巧妙的低維潛在向量壓縮和解耦式 RoPE 設(shè)計， MLA 能在顯存效率與建模能力間實(shí)現(xiàn)更優(yōu)平衡。
作者 | Shirley Li
編譯 | 岳揚(yáng)
本文是「DeepSeek-V3 技術(shù)解析」系列文章的第一篇，我們將深入解讀深度求索公司最新開源的 DeepSeek-V3 模型[1,2]。
在本系列文章中，我們計劃涵蓋兩大主題：
- DeepSeek-V3 的主要架構(gòu)創(chuàng)新，包括：多頭潛在注意力（Multi-head Latent Attention，MLA）[3]、DeepSeekMoE[4]、無輔助損失的負(fù)載均衡技術(shù)（auxiliary-loss-free load balancing）[5]、多詞元預(yù)測訓(xùn)練方法（multi-token prediction training）。
- DeepSeek-V3 的訓(xùn)練過程，包括：預(yù)訓(xùn)練階段、微調(diào)階段、強(qiáng)化學(xué)習(xí)對齊階段。
本文重點(diǎn)解析首次在 DeepSeek-V2 中提出、并延續(xù)應(yīng)用于 DeepSeek-V3 的多頭潛在注意力（MLA）架構(gòu)。
內(nèi)容大綱：
- 技術(shù)背景：從標(biāo)準(zhǔn)多頭注意力（MHA）機(jī)制出發(fā)，解析推理階段 Key-Value 緩存的必要性，探討 MQA（多查詢注意力機(jī)制）和 GQA（分組查詢注意力機(jī)制）的優(yōu)化思路，并闡述“旋轉(zhuǎn)”位置編碼（RoPE）的工作原理。
- 多頭潛在注意力機(jī)制（MLA）：深入解析 MLA 的創(chuàng)新動機(jī)，闡釋解耦式 RoPE 的設(shè)計原理，并展示其性能優(yōu)勢。
- 參考文獻(xiàn)
01 技術(shù)背景
為了便于讀者更好地理解 MLA 架構(gòu)并保持本文的完整性，在深入探討 MLA 的技術(shù)細(xì)節(jié)之前，我們將首先回顧幾個相關(guān)技術(shù)概念。
1.1 Decoder-only Transformers 中的多頭注意力機(jī)制
有一點(diǎn)需要特別說明，MLA 架構(gòu)的研發(fā)初衷是加快自回歸文本生成的推理效率，因此本文討論的 MHA（多頭注意力機(jī)制）特指 Decoder-only Transformer 架構(gòu)。
圖 1 對比了三種可用于解碼的 Transformer 架構(gòu)：(a) 展示的是《Attention is All You Need》論文提出的編碼器-解碼器架構(gòu)。其解碼器部分隨后被文獻(xiàn) [6] 簡化，形成如 (b) 所示的 decoder-only Transformer 架構(gòu)，該架構(gòu)被 GPT 等生成模型廣泛采用[8]。
(c) 則是當(dāng)前大語言模型更常采用的優(yōu)化版本，對輸入而不是輸出進(jìn)行歸一化。并采用 RMS Norm 替代 LayerNorm，本文將以 (c) 所示架構(gòu)作為 baseline 架構(gòu)展開討論。
圖 1. Transformer 架構(gòu)。(a) 文獻(xiàn) [6] 中提出的編碼器-解碼器結(jié)構(gòu)；(b) 文獻(xiàn) [7] 提出的 Decoder-only Transformer，并被 GPT [8] 采用；(c)是(b) 的優(yōu)化版本，在注意力層前加入 RMS Norm [3]。
在這種情況下， MHA 計算大致遵循文獻(xiàn) [6] 中的流程，如下圖所示：
圖 2. 縮放點(diǎn)積注意力與多頭注意力架構(gòu)對比（圖源文獻(xiàn)[6]）
假設(shè)模型包含 n_h 個注意力頭，每個注意力頭維度為 d_h，則拼接后總維度為 n_h·d_h。
對于具有 l 層的模型，若設(shè)某層第 t 個 token 的輸入表示為維度 d 的向量 h_t，則需要通過線性映射矩陣（linear mapping matrices）將 h_t 的維度從 d 映射到（h_n - d_h）。
更形式化地表述如下（公式引用自文獻(xiàn)[3]）：
其中，WQ、WK 和 W^V 是線性映射矩陣：
完成映射后，q_t、k_t 和 v_t 將被分割為 n_h 個注意力頭進(jìn)行縮放點(diǎn)積注意力計算：
最終通過投影矩陣 W^O 將維度從 (h_n - d_h) 反向映射到 d：
需要強(qiáng)調(diào)的是，上述公式 (1)-(8) 描述的是單 token 處理流程。在推理過程中，每個新生成的 token 都需要重復(fù)這一計算流程，這產(chǎn)生了大量重復(fù)計算 —— 這正是催生鍵值緩存（Key-Value Cache）技術(shù)的關(guān)鍵動因。
1.2 Key-Value Cache
正如其名稱所示，鍵值緩存（Key-Value Cache）是一種通過緩存和復(fù)用先前計算的鍵值對（Keys and Values）來加速自回歸過程的技術(shù)，避免了在每個解碼步驟重新計算這些參數(shù)的需求。
需要注意的是，鍵值緩存通常僅在推理階段使用，因?yàn)樵谟?xùn)練階段，我們?nèi)匀恍枰⑿刑幚碚麄€輸入序列。
鍵值緩存通常以滾動緩沖區(qū)（rolling buffer）的形式實(shí)現(xiàn)。在每個解碼步驟中，僅計算新的查詢向量 Q（Query），而緩存中存儲的 K（Keys）和 V（Values）會被復(fù)用，注意力機(jī)制將通過新計算的 Q 與復(fù)用的 K、V 進(jìn)行運(yùn)算。與此同時，新生成的 token 對應(yīng)的 K 和 V 也會被追加到緩存中以供后續(xù)使用。
然而，鍵值緩存帶來的加速是以內(nèi)存消耗為代價的。由于鍵值緩存的大小通常與批處理量大?。╞atch size）× 序列長度（sequence length）× 隱藏層維度（hidden size）× 注意力頭數(shù)（number of heads）成正比，當(dāng)出現(xiàn)更大批處理量或更長序列時，極易形成內(nèi)存瓶頸。
這種限制進(jìn)一步催生了兩項旨在突破該瓶頸的技術(shù)：多查詢注意力機(jī)制（Multi-Query Attention）和分組查詢注意力機(jī)制（Grouped-Query Attention）。
1.3 多查詢注意力機(jī)制（MQA） vs 分組查詢注意力機(jī)制（GQA）
下圖展示了原始的多頭注意力機(jī)制（MHA）、分組查詢注意力機(jī)制（GQA）[10] 和多查詢注意力機(jī)制（MQA）[9] 的對比。
圖 3. MHA [6]、GQA [10] 與 MQA [9]（圖源文獻(xiàn)[10]）
MQA 的核心思想是讓所有查詢頭（query heads）共享一個單獨(dú)的鍵頭（key head）和值頭（value head）。 這種做法能顯著降低內(nèi)存占用，但也會對注意力的計算精度產(chǎn)生負(fù)面影響。
GQA 可視為 MHA 與 MQA 的一種折中方案。 在這種方法中，一組查詢頭僅共享一對鍵頭和值頭，而非所有查詢頭共享同一對。盡管如此，相較于原始的 MHA，其效果仍會有所遜色。
在后續(xù)章節(jié)中，我們將探討 MLA 如何在內(nèi)存效率與建模精度之間實(shí)現(xiàn)平衡。
1.4 旋轉(zhuǎn)位置編碼（RoPE）
最后需要提及的相關(guān)背景知識是旋轉(zhuǎn)位置編碼（RoPE）[11]。該方法通過在多頭注意力機(jī)制中對查詢向量（Query）和鍵向量（Key）施加基于正弦函數(shù)的旋轉(zhuǎn)操作，將位置信息直接編碼到注意力計算中。
具體而言，RoPE 會對每個詞元（token）的查詢向量和鍵向量應(yīng)用一個與位置相關(guān)的旋轉(zhuǎn)矩陣。該矩陣以正弦和余弦函數(shù)為基礎(chǔ)，但以一種獨(dú)特的方式應(yīng)用它們來實(shí)現(xiàn)旋轉(zhuǎn)。
為理解其“位置相關(guān)（position-dependent）”特性，我們以一個 4 維的嵌入向量（x?, x?, x?, x?）為例。
要應(yīng)用 RoPE，首先需要將連續(xù)的維度兩兩配對：
- (x?, x?) → 位置1
- (x?, x?) → 位置2
然后，對其中的每一對應(yīng)用旋轉(zhuǎn)矩陣：
圖 4. 應(yīng)用到詞元對的旋轉(zhuǎn)矩陣示意圖。由原文作者供圖。
其中旋轉(zhuǎn)角度 θ = θ(p) = p ? θ?（θ?為基頻參數(shù)）。在此示例中，(x?, x?) 會被旋轉(zhuǎn) θ? 角度，而 (x?, x?) 則被旋轉(zhuǎn) 2 · θ? 角度。
因此，我們稱這種旋轉(zhuǎn)矩陣為“位置相關(guān)（position-dependent）”：在每個位置（或每對），將應(yīng)用不同的旋轉(zhuǎn)矩陣，其中旋轉(zhuǎn)角度由位置決定。
RoPE 因其在長序列編碼中的高效性被現(xiàn)代大語言模型廣泛采用。然而，從上述公式可見，該方法對查詢向量和鍵向量的位置敏感性，使其在部分場景下與 MLA 機(jī)制不兼容。
02 Multi-head Latent Attention
最后我們終于可以進(jìn)入 MLA 部分的講解。在這一章節(jié)中，我們將首先闡述 MLA 的核心設(shè)計思想，然后深入探討其為何需要修改 RoPE，最后展示 MLA 的具體算法實(shí)現(xiàn)及其性能表現(xiàn)。
2.1 MLA：核心設(shè)計思想
MLA 的核心設(shè)計思想是將注意力機(jī)制的輸入 h_t 壓縮為一個低維的潛在向量（維度 d_c，且 d_c 遠(yuǎn)小于原始維度（h_n·d_h））。 當(dāng)需要計算注意力時，我們可以將這個潛在向量重新映射回高維空間，以還原鍵向量和值向量。通過這種方式，只需存儲潛在向量即可，從而實(shí)現(xiàn)顯存占用的大幅降低。
這個過程可以通過以下公式更正式地進(jìn)行描述。其中 c^{KV}_t 表示潛在向量，W^{DKV} 是壓縮矩陣（上標(biāo) D 代表"下投影"，即降維操作），負(fù)責(zé)將 h_t 的維度從（h_n·d_h）壓縮到d_c；而 W^{UK} 和 W^{UV} 則是上投影矩陣，負(fù)責(zé)將共享的潛在向量映射回高維空間。
類似地，我們也可以將查詢向量映射到一個潛在的低維向量，然后再將其映射回原始的高維空間：
2.2 為何需要解耦式 RoPE
正如前文所述，RoPE 是訓(xùn)練生成模型處理長序列的常用位置編碼方案。如果我們直接應(yīng)用上述 MLA 策略，將會與 RoPE 不兼容。
為了更清晰地理解這一點(diǎn)，來看看使用公式 (7) 計算注意力時的情形：當(dāng)我們對 q 進(jìn)行轉(zhuǎn)置并與 k 相乘時，矩陣 W^Q 和 W^{UK} 會出現(xiàn)在中間環(huán)節(jié)，它們的組合等效于從將 d_c（輸入向量的維度）映射到目標(biāo)維度 d。
在文獻(xiàn) [3] 中，作者將此現(xiàn)象描述為 W^{UK} 可以被 W^Q “吸收”，因此我們無需在緩存中存儲 W^{UK}，從而進(jìn)一步降低了內(nèi)存使用量。
然而，當(dāng)我們將圖 (4) 中的旋轉(zhuǎn)矩陣考慮在內(nèi)時，情況就不是這樣了 —— RoPE 會在 W^{UK} 左側(cè)施加一個旋轉(zhuǎn)矩陣，而該旋轉(zhuǎn)矩陣最終會夾在轉(zhuǎn)置后的 W^Q 和 W^{UK} 之間。
正如技術(shù)背景部分所述，這個旋轉(zhuǎn)矩陣是位置相關(guān)（position-dependent）的，即每個位置對應(yīng)的旋轉(zhuǎn)矩陣都不同。因此，W^{UK} 無法再被 W^Q 吸收。
為解決這一矛盾，作者提出了“解耦式RoPE”方案：通過引入額外的查詢向量和一個共享的鍵向量，并僅在 RoPE 過程中使用這些新增向量，同時保持原始鍵向量與旋轉(zhuǎn)矩陣的隔離。
完整的 MLA 流程可總結(jié)如下（公式編號沿用文獻(xiàn) [3] 附錄 C 的編排）：
圖 5. MLA 處理流程（作者根據(jù)文獻(xiàn) [3] 中的公式所編輯的示意圖）
其中：
- 公式 (37)-(40) 描述查詢向量的處理流程
- 公式 (41)-(42) 描述鍵向量的處理流程
- 公式 (43)-(44) 說明如何將新增的共享鍵向量用于 RoPE，需注意公式 (42) 的輸出不參與 RoPE 計算
- 公式 (45) 描述值向量的處理流程
在此過程中，僅需緩存藍(lán)色變量。該流程可通過下圖更直觀地展示：
圖 6. MLA 流程圖（來源：文獻(xiàn)[3]）
2.3 MLA 的性能表現(xiàn)
下表對比了 MHA、GQA、MQA 與 MLA 在 KV 緩存（每個 token）所需元素數(shù)量及建模能力上的差異，證明 MLA 確實(shí)能在顯存效率與建模能力間實(shí)現(xiàn)更優(yōu)平衡。
有趣的是，MLA 的建模能力甚至超越了原始 MHA 架構(gòu)。
文獻(xiàn) [3] 中的表 1
下表展示了 7B 參數(shù)規(guī)模下 MHA、GQA 與 MQA 的性能表現(xiàn)，其中 MHA 顯著優(yōu)于 MQA 和 GQA。
文獻(xiàn) [3] 中的表 8
文獻(xiàn) [3] 的作者還對 MHA 與 MLA 進(jìn)行了對比分析，結(jié)果如下表所示。數(shù)據(jù)顯示 MLA 在整體指標(biāo)上表現(xiàn)更優(yōu)。
文獻(xiàn) [3] 中的表 9
03 References
[1] DeepSeek（??https://www.deepseek.com/）??
[2] DeepSeek-V3 Technical Report（??https://github.com/deepseek-ai/DeepSeek-V3/blob/main/DeepSeek_V3.pdf）??
[3] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model（??https://arxiv.org/abs/2405.04434）??
[4] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models（??https://arxiv.org/abs/2401.06066）??
[5] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts（??https://arxiv.org/abs/2408.15664）??
[6] Attention Is All You Need（??https://arxiv.org/abs/1706.03762）??
[7] Generating Wikipedia by Summarizing Long Sequences（??https://arxiv.org/pdf/1801.10198）??
[8] Improving Language Understanding by Generative Pre-Training（??https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf）??
[9] Fast Transformer Decoding: One Write-Head is All You Need（??https://arxiv.org/pdf/1911.02150）??
[10] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints（??https://arxiv.org/abs/2305.13245）??
[11] RoFormer: Enhanced Transformer with Rotary Position Embedding（??https://arxiv.org/abs/2104.09864??
Thanks for reading!
Hope you have enjoyed and learned new things from this blog!
About the author
Shirley Li
I am a Machine Learning Engineer working on building multi-modality models to solve real-world problems.
END
本期互動內(nèi)容 ??
?你最看重的是模型性能還是推理效率？為什么？歡迎在評論區(qū)談?wù)勀愕目捶ǎ?/strong>
原文鏈接：
