注意力機制

對應原題（Common ‧ LLM Internals）

• Explain scaled dot-product attention and why the 1/√d_k scaling factor matters.
• What are Multi-Query Attention (MQA) and Grouped-Query Attention (GQA), and what do they trade away? — Asked at: Meta, Mistral AI
• What is Multi-head Latent Attention (MLA) and why did DeepSeek introduce it? — Asked at: DeepSeek, Moonshot AI
• Implement scaled dot-product attention with a causal mask, from scratch. — Asked at: Anthropic, Google DeepMind, Amazon

🎯 什麼情境該想到我

當你「要理解 Transformer 到底在算什麼、或被要求手刻一個 attention」的時候。這是所有 LLM 內部題的地基——KV Cache、FlashAttention、RoPE 全都長在這上面。

⚙️ 怎麼用（公式與推導）

Q、K、V 是什麼

一句話：每個 token 拿自己的「問題」去比對所有 token 的「索引」，依相似度加權取回它們的「內容」。

輸入 $X\in {\mathbb{R}}^{N\times d_{model}}$（N 個 token），用三個學來的矩陣投影：

名稱	算式	直覺
Query	$Q=X{W}_Q$	我想找什麼
Key	$K=X{W}_K$	我能被什麼找到
Value	$V=X{W}_V$	被找到時我交出什麼

核心公式

$\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^{\top }}{\sqrt{d_k}}\right)V$

四步：投影 → 打分（$Q{K}^{\top }$）→ 縮放後 softmax 成權重 → 加權平均 V。

★ 為什麼要除 $\sqrt{d_k}$（這題一定會追問推導）

$q\cdot k={\sum }_{i=1}^{d_k}{q}_i{k}_i$ 是 $d_k$ 項相加。假設 $q_i,k_i$ 獨立、均值 0、變異數 1：

• $\mathrm{E}[q_i{k}_i]=0$，$\mathrm{Var}(q_i{k}_i)=1$
• 獨立項相加，變異數相加 → $\mathrm{Var}(q\cdot k)=d_k$，標準差 $\sqrt{d_k}$

所以點積的尺度會隨維度以 $\sqrt{d_k}$ 成長。$d_k=128$ 時 logits 的典型量級就到 ±11 了。

後果：softmax 的輸入一大，輸出就飽和成接近 one-hot；而 softmax 飽和時梯度趨近 0（$\frac{\partial }{\partial z_i}\text{softmax}=p_i(1-p_i)$，$p_i\to 1$ 或 $0$ 都讓它 $\to 0$）。訓練初期就梯度消失，學不動。

除以 $\sqrt{d_k}$ 把變異數拉回 1，讓 logits 不隨維度膨脹。

追問準備：「為什麼不是除 $d_k$？」——因為要抵銷的是標準差不是變異數。除 $d_k$ 會過度壓縮，讓分布太平坦、attention 失去選擇性。

Causal Mask（自迴歸的關鍵）

decoder-only 模型不能看到未來。做法是在 softmax 之前把上三角（$j>i$）設成負無限大：

import torch, math
 
def sdpa(q, k, v, causal=True):
    # q,k,v: (B, H, N, d_head)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)   # (B,H,N,N)
    if causal:
        N = scores.size(-1)
        mask = torch.ones(N, N, dtype=torch.bool, device=scores.device).triu(1)
        # ★ 用 dtype 的最小值，不要寫死 -1e9（fp16 上限 65504，會溢位成 -inf 或 NaN）
        scores = scores.masked_fill(mask, torch.finfo(scores.dtype).min)
    attn = scores.softmax(dim=-1)
    return attn @ v, attn

為什麼是 $-\infty$ 不是 0：softmax 後才變 0。若直接把權重設 0，其他位置的權重和不會是 1。

多頭注意力 MHA

把 $d_{model}$ 切成 $h$ 份，每份 $d_{head}=d_{model}/h$，各自跑一次 attention，concat 後過 $W_O$：

$\text{MHA}(X)=\text{Concat}({\text{head}}_1,\dots ,{\text{head}}_h)W_O$

為什麼要多頭：單頭的 softmax 只能表達「一種」相似度關係；多頭讓不同子空間學不同的關係（語法依存、指代、位置鄰近…）。總 FLOPs 與單頭近似，因為維度被切開了。

複雜度（所有推論最佳化的起點）

項目	複雜度	意義
時間	$O(N^{2}d)$	序列翻倍，計算量四倍
注意力矩陣記憶體	$O(N^2)$	這是 FlashAttention 要消滅的東西
參數量	$O(d^2)$	與 N 無關

★ MHA → MQA → GQA → MLA：省的是 KV，不是 FLOPs

推論時瓶頸不在算力而在記憶體頻寬——每生成一個 token 都要把整個 KV cache 讀一遍。於是演化出「少存一點 K/V」的路線：

變體	KV 頭數	KV cache 相對大小	代價	代表
MHA	$h$	$1\times$	基準	GPT-2、Llama 1
MQA（Shazeer, 2019）	1	$1/h$	品質明顯下滑、訓練不穩	PaLM
GQA（Ainslie et al., 2023）	$g$ 組（$1<g<h$）	$g/h$	品質接近 MHA	Llama 2 70B（$g=8$）、Mistral
MLA（DeepSeek-V2, 2024）	低秩 latent	比 GQA 更小	多一次升維運算	DeepSeek-V2/V3

• MQA：所有 Query 頭共用同一組 K/V。cache 直接除以 $h$，但表達力損失大。
• GQA：折衷。把 $h$ 個 Q 頭分成 $g$ 組，組內共用一組 K/V。$g=1$ 退化成 MQA，$g=h$ 就是 MHA。論文的關鍵貢獻是能從既有 MHA checkpoint mean-pool 出 GQA 權重再少量續訓，不用重練。
• MLA：不是減頭數，是壓維度。把 K/V 投影到一個低維 latent $c$，cache 只存 $c$；用的時候再升回來。DeepSeek 的動機是長上下文下 KV cache 比模型權重還大。

一句話回答面試官：MQA/GQA/MLA 省的是 KV cache 佔用與記憶體頻寬，不是注意力的 FLOPs。它們讓你在同樣顯存下塞更大的 batch，進而提高吞吐。

🧪 我實際套用的紀錄

• 2026-09-20：從面試題庫展開，尚未實際套用。

⚠️ 注意 / 什麼時候不適用

• mask 值別寫死。-1e9 在 fp32 沒事，在 fp16 會溢位（fp16 最大約 65504）。一律用 torch.finfo(dtype).min。
• softmax 要減 max 才數值穩定。PyTorch 內建已處理，自己刻的話別漏。
• $\sqrt{d_k}$ 的前提是「$q,k$ 分量均值 0、變異數 1」。這靠初始化與 LayerNorm 維持；如果你改了初始化尺度，這個縮放未必最佳（muP 就是在重新設計這類尺度）。
• MQA/GQA 不會讓 prefill 變快。prefill 是 compute-bound，瓶頸在算力；它們解的是 decode 階段的頻寬問題。見 Prefill 與 Decode。
• 實務上不要自己刻。手刻是面試題；上線用 torch.nn.functional.scaled_dot_product_attention（會自動選 FlashAttention 等後端）。

🔗 相關

• KV Cache —— 這裡省下的 K/V，就是那裡要存的東西
• FlashAttention —— 同一個公式，改成不具體化 $N\times N$ 矩陣的算法
• 位置編碼與 RoPE —— attention 本身沒有位置概念，位置是外掛進來的
• Prefill 與 Decode —— 為什麼 decode 階段卡在記憶體頻寬
• 工具-LLM推論最佳化 —— 決策層：什麼時候該動哪一招