Transformer 需要三個投影嗎？QKV 變體分析

主要結論：投影共享降低記憶體開銷

Transformer 可以在保持效能的同時，透過共享權重投影大幅降低記憶體需求。具體而言，將鍵（Key）與值（Value）投影共享 (K=V) 可在結合 Multi-Query Attention (MQA) 時，使 KV 快取減少高達 96.9%，且模型困惑度（perplexity）僅有輕微下降。這暗示標準的三投影（Query、Key、Value）架構對許多任務而言可能過於複雜。

QKV 變體的系統性評估

研究人員評估了注意力機制中權重投影的三項主要限制，以判斷是否真的需要全部三個投影：

• Shared Key-Value (K=V)： 鍵（Key）與值（Value）投影相同。
• Shared Query-Key (Q=K)： 查詢（Query）與鍵（Key）投影相同，產生對稱的注意力圖。
• Single Projection (Q=K=V)： 三個投影共用相同的權重。

為了減輕 Q=K 與 Q=K=V 所帶來的對稱性問題，研究使用了 2D 位置編碼，以在注意力機制中引入非對稱性。

各領域的效能表現

實驗在合成任務、視覺任務（MNIST、CIFAR、TinyImageNet 以及異常偵測）以及語言建模上進行。結果顯示，投影共享的 Transformer 與標準 QKV Transformer 表現相當，甚至在某些情況下更佳。

語言建模結果

在使用 300M 與 1.2B 參數模型、以 10B token 訓練的語言建模測試中，K=V 變體展現了顯著的效率提升：

• KV 快取減少： K=V 投影共享可將 KV 快取減少 50%。
• 困惑度影響： 此減少僅導致困惑度下降 3.1%。

與頭部共享的協同效應（GQA/MQA）

投影共享與現有的頭部共享技術，如 Grouped-Query Attention (GQA) 與 Multi-Query Attention (MQA)，相輔相成。將這些方法結合可帶來巨大的記憶體節省，適用於邊緣部署：

• K=V + GQA-4： KV 快取減少 87.5%。
• K=V + MQA： KV 快取減少 96.9%。

理論洞見：為何 K=V 有效

研究指出，K=V 能保持品質是因為鍵與值常常位於相似的表徵空間。此外，注意力通常在低秩（low-rank）情況下運作，標準 QKV 公式中的冗餘往往未被充分利用。相反地，Q=K 變體往往失效，因為它破壞了注意力機制的基本方向性。

社群批評與技術討論

雖然結果對於裝置端推論相當有前景，但機器學習社群對此研究的普適性與方法論提出了多項質疑：

訓練規模與過度訓練

有評論指出，1.2B 參數模型僅在 10B token 上訓練，遠低於 Chinchilla 最佳計算預算。擔心的是，對注意力的簡化在訓練不足的情況下可能顯得有效，但在現代大型語言模型（LLM）以「過度訓練」的規模（兆級 token）時，標準注意力的完整表達能力變得更為關鍵，簡化方法可能無法擴展。

數學符號

部分讀者指出論文符號不夠清晰，尤其是使用「Q-K=V」來表示共享的鍵值投影 (K=V)，而非數學上的減法運算。

歸納偏差與架構

從實務者的討論可見，標準的 QKV 公式之所以被廣泛採用，未必是因為它是最佳的數學方法，而是因為它在 GPU 上計算效率高，且提供了一個安全且易於理解的擴展基線。