多頭注意力機(jī)制

多頭注意力機(jī)制是Transformer模型的關(guān)鍵部分，論文《Attention is All You Need — 2017 Google》提出了注意力機(jī)制，通過(guò)并行計(jì)算來(lái)捕獲序列中token之間的全局依賴(lài)關(guān)系。這種注意力機(jī)制是一種顛覆性的創(chuàng)新，打破了RNN和CNN模型的序列約束，直接推動(dòng)了語(yǔ)言模型的大規(guī)模演進(jìn)。

在多頭注意力機(jī)制中，每個(gè)頭的鍵、值和查詢都有一個(gè)專(zhuān)用的投影矩陣。例如，原始 Transformer 模型（在 Transformer 論文中）中有 8 個(gè)頭，應(yīng)該有 8×3（鍵、值和查詢）= 24 個(gè)獨(dú)立的投影矩陣。鍵、值和查詢將分別使用它們的投影矩陣投影到h 個(gè)不同的副本上。在上述具有 8 個(gè)頭的 Transformer 模型中，鍵將有 8 個(gè)投影副本，值將有 8 個(gè)投影副本，查詢也是如此。

多頭注意力機(jī)制

多查詢注意力機(jī)制

快速Transformer解碼：只需一個(gè)寫(xiě)入頭即可——2019谷歌這篇論文指出Transformer的一個(gè)主要挑戰(zhàn)是增量推理的速度，并且它受到重新加載編碼注意層狀態(tài)的大型“鍵”和“值”張量所需內(nèi)存帶寬的限制。本文引入了多查詢注意機(jī)制，讓所有頭共享相同的鍵和值。對(duì)于8個(gè)以上頭的注意情況，鍵只有1個(gè)副本，值只有1個(gè)副本，查詢?nèi)匀挥?個(gè)投影副本。采用多查詢注意機(jī)制后，解碼器增量推理速度提高了13.9倍（基線有8個(gè)頭）。

分組查詢注意力機(jī)制

雖然多查詢注意力機(jī)制可以顯著提高 Transformer 的增量推理速度，但Google 在 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — Dec 2023 中提到，多查詢注意力機(jī)制可能會(huì)導(dǎo)致質(zhì)量下降和訓(xùn)練不穩(wěn)定。為了克服這種質(zhì)量犧牲，本文引入了分組查詢注意力機(jī)制。

所有注意力頭并不共享鍵和值，而是一組注意力頭共享一個(gè)鍵和值的投影矩陣。本文在 T5（文本到文本傳輸）Transformer 模型中測(cè)試了 MQA，T5-Large 有 16 個(gè)注意力頭，T5-XXL 有 64 個(gè)注意力頭。本文為鍵/值選擇了 8 組（8 個(gè)注意力頭共享相同的鍵/值投影矩陣）。本文表明，與采用傳統(tǒng)多頭注意力的 T5 模型相比，采用分組查詢注意力的 T5 模型具有 5 到 6 倍的推理速度和相似的質(zhì)量（從下圖中，我們?nèi)匀豢梢钥吹劫|(zhì)量略有下降，即使它非常微小）。

多頭潛在注意力（MLA）

DeepSeek-V2：一種強(qiáng)大、經(jīng)濟(jì)、高效的混合專(zhuān)家語(yǔ)言模型為注意力模塊引入了多頭潛在注意力 (MLA)。“MLA 利用低秩鍵值聯(lián)合壓縮來(lái)消除推理時(shí)間鍵值緩存的瓶頸，從而支持高效推理。MLA 的核心是對(duì)鍵和值的低秩聯(lián)合壓縮，以減少 KV 緩存。”

低秩壓縮是一種用于減少深度學(xué)習(xí)模型的計(jì)算和內(nèi)存需求的技術(shù)。關(guān)于它的論文很多，例如：

神經(jīng)網(wǎng)絡(luò)的低秩壓縮：學(xué)習(xí)每一層的秩 UC Merced
使用低秩和低精度分解壓縮大型語(yǔ)言模型Stanford
LORC：采用漸進(jìn)壓縮策略的 LLM KV 緩存低秩壓縮 Gatech 和 Microsoft

DeepSeek V1 采用旋轉(zhuǎn)位置嵌入 (RoPE) 來(lái)攜帶 token 的位置信息（Transformer 采用正弦位置編碼，詳情請(qǐng)參考我的 Transformer 清晰解釋博客）。由于低秩 KV 壓縮與 RoPE 不兼容，DeepSeek V2 采用了解耦的 RoPE 策略，即使用額外的多頭查詢和共享密鑰來(lái)攜帶 RoPE。

論文表示MLA表現(xiàn)出比MHA更好的性能，并且MLA所需的KV緩存量（小型MoE模型為14%，大型MoE模型為4%）明顯小于MHA。

總結(jié)

上述測(cè)試報(bào)告中 MLA 的質(zhì)量超過(guò)了 MHA，這讓我很驚訝，我以為 MLA 的質(zhì)量與 MHA 差不多，甚至?xí)圆钜恍Ｒ驗(yàn)槊總€(gè)頭的 MHA 都有自己的投影鍵、值和查詢，這使得注意力模塊能夠輕松地學(xué)習(xí)關(guān)注不同的位置。