「DeepSeek-V3 技術解析」：多頭潛在注意力機制（MLA） 原創
編者按： 模型在生成長序列文本時，鍵值緩存占用的顯存是否讓 GPU 不堪重負？如何在保持模型性能的同時有效降低大語言模型推理階段的內存占用？
每一次模型推理，大量重復計算和龐大的內存占用不僅大幅增加了運營成本，還嚴重限制了模型的批處理能力和響應速度。多查詢注意力機制（MQA）和分組查詢注意力機制（GQA）雖能部分緩解問題，但往往以犧牲模型性能為代價。
DeepSeek 團隊在模型優化過程中開發的多頭潛在注意力機制（MLA）為資源受限場景提供了可行性方案。通過巧妙的低維潛在向量壓縮和解耦式 RoPE 設計， MLA 能在顯存效率與建模能力間實現更優平衡。
作者 | Shirley Li
編譯 | 岳揚
本文是「DeepSeek-V3 技術解析」系列文章的第一篇，我們將深入解讀深度求索公司最新開源的 DeepSeek-V3 模型[1,2]。
在本系列文章中，我們計劃涵蓋兩大主題：
- DeepSeek-V3 的主要架構創新，包括：多頭潛在注意力（Multi-head Latent Attention，MLA）[3]、DeepSeekMoE[4]、無輔助損失的負載均衡技術（auxiliary-loss-free load balancing）[5]、多詞元預測訓練方法（multi-token prediction training）。
- DeepSeek-V3 的訓練過程，包括：預訓練階段、微調階段、強化學習對齊階段。
本文重點解析首次在 DeepSeek-V2 中提出、并延續應用于 DeepSeek-V3 的多頭潛在注意力（MLA）架構。
內容大綱：
- 技術背景：從標準多頭注意力（MHA）機制出發，解析推理階段 Key-Value 緩存的必要性，探討 MQA（多查詢注意力機制）和 GQA（分組查詢注意力機制）的優化思路，并闡述“旋轉”位置編碼（RoPE）的工作原理。
- 多頭潛在注意力機制（MLA）：深入解析 MLA 的創新動機，闡釋解耦式 RoPE 的設計原理，并展示其性能優勢。
- 參考文獻
01 技術背景
為了便于讀者更好地理解 MLA 架構并保持本文的完整性，在深入探討 MLA 的技術細節之前，我們將首先回顧幾個相關技術概念。
1.1 Decoder-only Transformers 中的多頭注意力機制
有一點需要特別說明，MLA 架構的研發初衷是加快自回歸文本生成的推理效率，因此本文討論的 MHA（多頭注意力機制）特指 Decoder-only Transformer 架構。
圖 1 對比了三種可用于解碼的 Transformer 架構：(a) 展示的是《Attention is All You Need》論文提出的編碼器-解碼器架構。其解碼器部分隨后被文獻 [6] 簡化，形成如 (b) 所示的 decoder-only Transformer 架構，該架構被 GPT 等生成模型廣泛采用[8]。
(c) 則是當前大語言模型更常采用的優化版本，對輸入而不是輸出進行歸一化。并采用 RMS Norm 替代 LayerNorm，本文將以 (c) 所示架構作為 baseline 架構展開討論。
圖 1. Transformer 架構。(a) 文獻 [6] 中提出的編碼器-解碼器結構；(b) 文獻 [7] 提出的 Decoder-only Transformer，并被 GPT [8] 采用；(c)是(b) 的優化版本，在注意力層前加入 RMS Norm [3]。
在這種情況下， MHA 計算大致遵循文獻 [6] 中的流程，如下圖所示：
圖 2. 縮放點積注意力與多頭注意力架構對比（圖源文獻[6]）
假設模型包含 n_h 個注意力頭，每個注意力頭維度為 d_h，則拼接后總維度為 n_h·d_h。
對于具有 l 層的模型，若設某層第 t 個 token 的輸入表示為維度 d 的向量 h_t，則需要通過線性映射矩陣（linear mapping matrices）將 h_t 的維度從 d 映射到（h_n - d_h）。
更形式化地表述如下（公式引用自文獻[3]）：
其中，WQ、WK 和 W^V 是線性映射矩陣：
完成映射后，q_t、k_t 和 v_t 將被分割為 n_h 個注意力頭進行縮放點積注意力計算：
最終通過投影矩陣 W^O 將維度從 (h_n - d_h) 反向映射到 d：
需要強調的是，上述公式 (1)-(8) 描述的是單 token 處理流程。在推理過程中，每個新生成的 token 都需要重復這一計算流程，這產生了大量重復計算 —— 這正是催生鍵值緩存（Key-Value Cache）技術的關鍵動因。
1.2 Key-Value Cache
正如其名稱所示，鍵值緩存（Key-Value Cache）是一種通過緩存和復用先前計算的鍵值對（Keys and Values）來加速自回歸過程的技術，避免了在每個解碼步驟重新計算這些參數的需求。
需要注意的是，鍵值緩存通常僅在推理階段使用，因為在訓練階段，我們仍然需要并行處理整個輸入序列。
鍵值緩存通常以滾動緩沖區（rolling buffer）的形式實現。在每個解碼步驟中，僅計算新的查詢向量 Q（Query），而緩存中存儲的 K（Keys）和 V（Values）會被復用，注意力機制將通過新計算的 Q 與復用的 K、V 進行運算。與此同時，新生成的 token 對應的 K 和 V 也會被追加到緩存中以供后續使用。
然而，鍵值緩存帶來的加速是以內存消耗為代價的。由于鍵值緩存的大小通常與批處理量大小（batch size）× 序列長度（sequence length）× 隱藏層維度（hidden size）× 注意力頭數（number of heads）成正比，當出現更大批處理量或更長序列時，極易形成內存瓶頸。
這種限制進一步催生了兩項旨在突破該瓶頸的技術：多查詢注意力機制（Multi-Query Attention）和分組查詢注意力機制（Grouped-Query Attention）。
1.3 多查詢注意力機制（MQA） vs 分組查詢注意力機制（GQA）
下圖展示了原始的多頭注意力機制（MHA）、分組查詢注意力機制（GQA）[10] 和多查詢注意力機制（MQA）[9] 的對比。
圖 3. MHA [6]、GQA [10] 與 MQA [9]（圖源文獻[10]）
MQA 的核心思想是讓所有查詢頭（query heads）共享一個單獨的鍵頭（key head）和值頭（value head）。 這種做法能顯著降低內存占用，但也會對注意力的計算精度產生負面影響。
GQA 可視為 MHA 與 MQA 的一種折中方案。 在這種方法中，一組查詢頭僅共享一對鍵頭和值頭，而非所有查詢頭共享同一對。盡管如此，相較于原始的 MHA，其效果仍會有所遜色。
在后續章節中，我們將探討 MLA 如何在內存效率與建模精度之間實現平衡。
1.4 旋轉位置編碼（RoPE）
最后需要提及的相關背景知識是旋轉位置編碼（RoPE）[11]。該方法通過在多頭注意力機制中對查詢向量（Query）和鍵向量（Key）施加基于正弦函數的旋轉操作，將位置信息直接編碼到注意力計算中。
具體而言，RoPE 會對每個詞元（token）的查詢向量和鍵向量應用一個與位置相關的旋轉矩陣。該矩陣以正弦和余弦函數為基礎，但以一種獨特的方式應用它們來實現旋轉。
為理解其“位置相關（position-dependent）”特性，我們以一個 4 維的嵌入向量（x?, x?, x?, x?）為例。
要應用 RoPE，首先需要將連續的維度兩兩配對：
- (x?, x?) → 位置1
- (x?, x?) → 位置2
然后，對其中的每一對應用旋轉矩陣：
圖 4. 應用到詞元對的旋轉矩陣示意圖。由原文作者供圖。
其中旋轉角度 θ = θ(p) = p ? θ?（θ?為基頻參數）。在此示例中，(x?, x?) 會被旋轉 θ? 角度，而 (x?, x?) 則被旋轉 2 · θ? 角度。
因此，我們稱這種旋轉矩陣為“位置相關（position-dependent）”：在每個位置（或每對），將應用不同的旋轉矩陣，其中旋轉角度由位置決定。
RoPE 因其在長序列編碼中的高效性被現代大語言模型廣泛采用。然而，從上述公式可見，該方法對查詢向量和鍵向量的位置敏感性，使其在部分場景下與 MLA 機制不兼容。
02 Multi-head Latent Attention
最后我們終于可以進入 MLA 部分的講解。在這一章節中，我們將首先闡述 MLA 的核心設計思想，然后深入探討其為何需要修改 RoPE，最后展示 MLA 的具體算法實現及其性能表現。
2.1 MLA：核心設計思想
MLA 的核心設計思想是將注意力機制的輸入 h_t 壓縮為一個低維的潛在向量（維度 d_c，且 d_c 遠小于原始維度（h_n·d_h））。 當需要計算注意力時，我們可以將這個潛在向量重新映射回高維空間，以還原鍵向量和值向量。通過這種方式，只需存儲潛在向量即可，從而實現顯存占用的大幅降低。
這個過程可以通過以下公式更正式地進行描述。其中 c^{KV}_t 表示潛在向量，W^{DKV} 是壓縮矩陣（上標 D 代表"下投影"，即降維操作），負責將 h_t 的維度從（h_n·d_h）壓縮到d_c；而 W^{UK} 和 W^{UV} 則是上投影矩陣，負責將共享的潛在向量映射回高維空間。
類似地，我們也可以將查詢向量映射到一個潛在的低維向量，然后再將其映射回原始的高維空間：
2.2 為何需要解耦式 RoPE
正如前文所述，RoPE 是訓練生成模型處理長序列的常用位置編碼方案。如果我們直接應用上述 MLA 策略，將會與 RoPE 不兼容。
為了更清晰地理解這一點，來看看使用公式 (7) 計算注意力時的情形：當我們對 q 進行轉置并與 k 相乘時，矩陣 W^Q 和 W^{UK} 會出現在中間環節，它們的組合等效于從將 d_c（輸入向量的維度）映射到目標維度 d。
在文獻 [3] 中，作者將此現象描述為 W^{UK} 可以被 W^Q “吸收”，因此我們無需在緩存中存儲 W^{UK}，從而進一步降低了內存使用量。
然而，當我們將圖 (4) 中的旋轉矩陣考慮在內時，情況就不是這樣了 —— RoPE 會在 W^{UK} 左側施加一個旋轉矩陣，而該旋轉矩陣最終會夾在轉置后的 W^Q 和 W^{UK} 之間。
正如技術背景部分所述，這個旋轉矩陣是位置相關（position-dependent）的，即每個位置對應的旋轉矩陣都不同。因此，W^{UK} 無法再被 W^Q 吸收。
為解決這一矛盾，作者提出了“解耦式RoPE”方案：通過引入額外的查詢向量和一個共享的鍵向量，并僅在 RoPE 過程中使用這些新增向量，同時保持原始鍵向量與旋轉矩陣的隔離。
完整的 MLA 流程可總結如下（公式編號沿用文獻 [3] 附錄 C 的編排）：
圖 5. MLA 處理流程（作者根據文獻 [3] 中的公式所編輯的示意圖）
其中：
- 公式 (37)-(40) 描述查詢向量的處理流程
- 公式 (41)-(42) 描述鍵向量的處理流程
- 公式 (43)-(44) 說明如何將新增的共享鍵向量用于 RoPE，需注意公式 (42) 的輸出不參與 RoPE 計算
- 公式 (45) 描述值向量的處理流程
在此過程中，僅需緩存藍色變量。該流程可通過下圖更直觀地展示：
圖 6. MLA 流程圖（來源：文獻[3]）
2.3 MLA 的性能表現
下表對比了 MHA、GQA、MQA 與 MLA 在 KV 緩存（每個 token）所需元素數量及建模能力上的差異，證明 MLA 確實能在顯存效率與建模能力間實現更優平衡。
有趣的是，MLA 的建模能力甚至超越了原始 MHA 架構。
文獻 [3] 中的表 1
下表展示了 7B 參數規模下 MHA、GQA 與 MQA 的性能表現，其中 MHA 顯著優于 MQA 和 GQA。
文獻 [3] 中的表 8
文獻 [3] 的作者還對 MHA 與 MLA 進行了對比分析，結果如下表所示。數據顯示 MLA 在整體指標上表現更優。
文獻 [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
本期互動內容 ??
?你最看重的是模型性能還是推理效率？為什么？歡迎在評論區談談你的看法！
原文鏈接：
