跳至主要內容

    Self-Attention:Q、K、V 究竟在算什麼?

    從一個 query 的手算,走到多頭注意力的輸出投影

    sunny 桑尼· · 31 MIN READ · 213 次瀏覽

    Self-Attention:Q、K、V 究竟在算什麼?
    本文目錄

    如果一個句子有三個 token,self-attention 會不會先替每個 token 算一個「重要度」?這個說法不太準確。更貼近計算的描述是:每個位置先產生一個 query,拿它和各位置的 key 比分,再把分數變成權重,用權重混合各位置的 value。每個 query 都有自己的權重列,所以不同位置可以讀取不同的內容。

    本文先用圖分清連接範圍,再拆出 Q、K、V 的角色;接著用一組固定小數手算第一個 query,最後把流程延伸成兩個 head,核對每個張量的形狀、拼接與輸出投影。數字來自一段可重跑的 CPU 機制原型,權重是手動指定的,不是語料訓練結果,也不是翻譯或分類品質測試。

    Self-attention 的「全序列」是連接方式

    在 CNN 中,某個位置通常先和附近位置互動;RNN 則沿序列一步一步傳遞狀態。Self-attention 的每個位置可以直接對序列中的其他位置計算關聯。這說的是允許的資訊連接,不代表每個位置實際上都得到一樣大的權重。

    圖 1|完整圖先比較三種連接方式;手機長圖接著依原圖順序放大 CNN、RNN、self-attention 三個既有區塊,跨區關係請回看上方全圖。虛線也是原圖的連接,不能讀成權重大小。來源:D2L〈Self-Attention and Positional Encoding〉,Aston Zhang、Zachary C. Lipton、Mu Li、Alexander J. Smola,CC BY-SA 4.0。本文將固定版 SVG 轉為白底 PNG;手機版重用原圖完整視圖與原有三區裁切,只加白色間隔,沒有新增圖元。改作圖同採 CC BY-SA 4.0。

    「可以互相連接」和「模型已經知道詞序」是兩件事。如果只有 self-attention,沒有位置表示或位置遮罩,同一組 token 一起換順序,輸出也會跟著換位置;這稱為置換等變。它不是輸出完全不變,也不等於把整句平均成一個 bag-of-words 向量。本文的原型沒有實跑換序實驗;這裡說的是在相同位置共享投影、且無位置資訊與 mask 條件下,從逐位置計算可推得的性質。位置資訊與因果遮罩留到下一篇處理。

    Q、K、V 是同一輸入的三種投影

    把一段長度為 T 的輸入向量排成矩陣 X。Self-attention 讓同一個 X 分別經過三組不同的可學線性投影:

    Q = X W_Q + b_Q
    K = X W_K + b_K
    V = X W_V + b_V
    

    同一個輸入位置會因此得到 query、key 和 value 三種向量,但它們通常不是相同數值,因為投影矩陣與偏置不同。本文公式把權重寫成 [輸入特徵, 輸出特徵];程式為配合 NumPy 與 PyTorch 的 [輸出特徵, 輸入特徵] 儲存方式,實際計算 X @ W.T + b。可以把 query 想成「這個位置現在要找什麼」,key 是各位置供比較的特徵,value 是最後要依權重取回的內容。這是協助記憶的比喻;它們不是資料庫欄位,也不保證某個 head 會固定學會某一種人類可命名的語法角色。

    圖 2|此圖展示一般 attention pooling:一個 query 對多個 key 評分,再用權重混合對應 values。它不是產生 Q/K/V 線性投影的圖,也不是 multi-head attention。來源:D2L〈Queries, Keys, and Values〉,Aston Zhang、Zachary C. Lipton、Mu Li、Alexander J. Smola,CC BY-SA 4.0。本文由固定版原 SVG 直接轉為白底 PNG,圖元未改;封面亦將此完整原圖等比置入白底 1600×900 畫布,沒有裁切或新增圖元。改作圖同採 CC BY-SA 4.0。

    這個區分很重要:先由 X 投影出 Q/K/V,再由 Q 和 K 算相容分數,最後才用分數加權 V。圖 2 的 α(q, k) 已是配對權重的通用概念;它沒有畫出 Transformer 的投影矩陣,也沒有顯示本文原型的數值。

    從分數走到加權和

    本文的數值例聚焦於常見的 scaled dot-product attention、softmax 與 dropout=0;attention 的評分函數還有其他設計。這個版本先計算 query 和 key 的點積,再除以 key 維度平方根,最後沿 keys 做 softmax:

    scores   = Q Kᵀ / √d_k
    weights  = softmax(scores, along keys)
    context  = weights V
    

    若有 T 個 token、單一 head,Q 和 K 的形狀都是 T × d_k,所以 QKᵀ 是 T × T。每一列代表一個 query token,每一欄代表一個 key token;softmax 對每列的 key 位置正規化。接著權重矩陣乘上 V,每個 query 位置得到一個加權 context。

    圖 3|完整原圖顯示 query 對多列 key 打分,經 softmax 後分別乘上 values,再加總為 output。圖中五組 key/value 是同一流程的多個位置,不是五個 head,也沒有本站實跑的分數或權重。來源:D2L〈Attention Scoring Functions〉,Aston Zhang、Zachary C. Lipton、Mu Li、Alexander J. Smola,CC BY-SA 4.0。本文將固定版原 SVG 直接轉為白底 PNG,保留原線路與間距;改作圖同採 CC BY-SA 4.0。

    來看原型中的第一個 query。輸入第一列是 x₁ = [0.2, −0.4, 1.0, 0.5]。三組投影產生:

    q₁ = [ 0.63, −0.32, −0.16, 0.95]
    k₁ = [−0.47, −0.05,  0.57, 0.69]
    k₂ = [ 0.38, −0.09,  0.90, 0.07]
    k₃ = [ 0.11,  0.84, −0.75, 0.55]
    

    兩個 head 各取向量中的兩個特徵。先看 head 0:q₁⁽⁰⁾ = [0.63, −0.32],三個 key 的 head 0 部分分別是 [−0.47, −0.05]、[0.38, −0.09]、[0.11, 0.84]。點積、縮放和權重如下:

    未縮放點積: [−0.2801,  0.2682, −0.1995]
    除以 √dₕ:  [−0.1981,  0.1896, −0.1411]   (dₕ = 2)
    softmax:    [ 0.2831,  0.4172,  0.2997]
    

    第一個權重不是最大值,第二個 key 在這個手設案例中拿到較高比重。Head 0 對應的 values 為 [0.07, 1.02]、[0.78, 0.10]、[−0.09, −0.35],因此第一個 query 的 context 是:

    0.2831 × [ 0.07,  1.02]
    + 0.4172 × [ 0.78,  0.10]
    + 0.2997 × [−0.09, −0.35]
    ≈ [0.3182, 0.2256]
    

    這是一組完全指定的向量與投影參數;它只能驗算「點積→縮放→softmax→加權 values」這條路徑。不能把 0.4172 說成模型學會第二個 token 比較重要,也不能據此推論語意或任務表現。Softmax 在某些分數分布下會飽和,梯度也可能很小;注意力計算本身不保證梯度永遠不消失。

    把 token-by-token 計算改成矩陣

    若有三個 query 和三個 key,一個 head 會一次形成 3 × 3 分數矩陣。沿每一列的 key 維度做 softmax 後,得到同樣大小的權重矩陣 A;它乘上 3 × dₕ 的 V,輸出 3 × dₕ 的 context。這樣寫能清楚標出每個輸出位置各自讀取哪些 value。

    圖 4|TensorFlow 教材圖畫出五個 query 與五組 key/value 位置可互相連接。原圖 query 沿水平軸、key/value 沿垂直軸,與常見 QKᵀ 矩陣的 query 列方向相反;藍格表示連接關係,不是均勻權重或本次實跑值,也不是五個 head。來源:TensorFlow〈Neural machine translation with a Transformer and Keras〉,TensorFlow Authors,CC BY 4.0。本文保留整張原圖,按比例放到白底畫布,未改圖元。

    把多個 query 寫在矩陣裡時,A[i, j] 就是第 i 個 query 對第 j 個 key 的權重。每一列各自加總為 1(在沒有 dropout、且該列有有效 key 的 softmax 設定下);這不表示整個矩陣所有元素加起來是 1。也不要將教材示意圖的格子顏色或可連接線誤當成這些數值。

    Multi-head:切分特徵、分別注意、再投影

    Multi-head attention 對每個 head 使用不同的 query、key、value 投影。若模型寬度 d_model = 4 且 head 數 h = 2,這個小例子把每個投影拆成兩個寬度 dₕ = d_model / h = 2 的 head。每個 head 都會對三個位置各算一列權重,產生三個二維 context;兩組 context 沿特徵維度接起來,回到四維,再通過輸出投影 W_O。

    圖 5|不同 head 各自接收 Queries、Keys、Values 的投影結果,分別計算 attention;輸出 concat 後,還要經上方 FC(W_O)投影。圖中兩個 head 和省略號是多頭結構示意,不代表模型只允許兩個 head。來源:D2L〈Multi-Head Attention〉,Aston Zhang、Zachary C. Lipton、Mu Li、Alexander J. Smola,CC BY-SA 4.0。本文將固定版原 SVG 直接轉為白底 PNG,沒有改動圖元;改作圖同採 CC BY-SA 4.0。

    這次原型的張量形狀可整理如下:

    階段形狀表示什麼
    輸入 X 與完整 Q/K/V3 × 4三個 token,每個四個特徵
    切成兩個 head2 × 3 × 2[head, token, head 特徵]
    每個 head 的分數與權重2 × 3 × 3[head, query 位置, key 位置]
    每個 head 的 context2 × 3 × 2每個 query 位置各有一個二維輸出
    拼接後與 W_O 輸出3 × 4兩個二維 head 接回四維,再投影

    原型計算的首 token,head 0 的 context 是 [0.3182, 0.2256];head 1 是 [0.1852, 0.0810]。依 head 次序拼接後:

    concat = [0.3182, 0.2256, 0.1852, 0.0810]
    W_O(concat) + b_O
           = [0.2705, 0.0593, 0.2071, −0.0055]
    

    所以 multi-head 的結果不只是把兩份向量接起來;W_O 會把拼接後的特徵再映射回模型寬度。對本例而言,三個 Q/K/V 線性層加一個輸出線性層,各自都是 4 × 4 權重和長度 4 的 bias,參數總數為 4 × (16 + 4) = 80。這只是本例配置的參數算術,不含完整 Transformer 區塊,也不含訓練、位置編碼或殘差層。

    若 d_model 固定,增加 head 通常會減少每個 head 分到的特徵寬度;並不是複製完整的 d_model 寬度後再無限制堆疊。多個投影提供多種可學的特徵子空間,但圖和這個小例子都不能證明每個 head 會自動分工成某種固定、可命名的語言規則。

    可重跑的 CPU 機制原型

    下面是完整程式:三個手設 token、d_model=4、兩個 head;以 NumPy 逐步計算投影、兩個 head 的 score、softmax、context、concat 和 W_O,再和 PyTorch nn.MultiheadAttention 的輸出及每-head 權重對照。PyTorch 的 Q/K/V 是 packed parameters,程式完整覆寫該參數及 W_O,沒有使用隨機初始化值。torch_scores 是用同一組 packed Q/K 另外以 F.linear/matmul 算出的參照分數,不是從 PyTorch attention 內部 hook 擷取的 score。

    實際記錄環境是 Python 3.10.12、NumPy 1.26.3、PyTorch 2.3.1+cu118,運算只跑 CPU float64。若要重跑,先準備已有 NumPy 1.26.3 與 PyTorch 2.3.1 的環境,把下方完整程式存成 multihead_attention_mechanics.py,再執行 python3 multihead_attention_mechanics.py。新虛擬環境的安裝、pip 來源與 GPU 執行都未在本次驗證。PyTorch 官方文件的輸入形狀、batch_first 和 per-head 權重選項,可參考 PyTorch 2.3 MultiheadAttention。

    """Deterministic CPU/float64 self-attention and 2-head MHA comparison.
    
    No training, corpus, positional encoding, mask, benchmark, or performance claim.
    """
    from __future__ import annotations
    
    import json
    import platform
    import numpy as np
    import torch
    from torch import nn
    
    
    def softmax_last(scores: np.ndarray) -> np.ndarray:
        """Stable softmax over the final (key) axis."""
        shifted = scores - np.max(scores, axis=-1, keepdims=True)
        exp_scores = np.exp(shifted)
        return exp_scores / np.sum(exp_scores, axis=-1, keepdims=True)
    
    
    def rounded(array, digits: int = 12):
        return np.round(np.asarray(array, dtype=np.float64), digits).tolist()
    
    
    def shape(array):
        return list(np.asarray(array).shape)
    
    
    def main() -> None:
        # Three token vectors, d_model=4. The same X is Q/K/V input: self-attention.
        x = np.array([
            [0.2, -0.4,  1.0,  0.5],
            [1.1,  0.3, -0.2,  0.8],
            [-0.7, 1.2,  0.4, -0.1],
        ], dtype=np.float64)
    
        # Dense, non-identity/non-diagonal matrices, stored as [out_features, in_features].
        w_q = np.array([
            [ 0.50, -0.20,  0.30,  0.10],
            [-0.40,  0.60,  0.20, -0.30],
            [ 0.10,  0.20, -0.50,  0.40],
            [ 0.30, -0.10,  0.60,  0.20],
        ], dtype=np.float64)
        w_k = np.array([
            [ 0.20,  0.40, -0.30,  0.10],
            [-0.10,  0.50,  0.20, -0.20],
            [ 0.60, -0.30,  0.10,  0.40],
            [-0.20,  0.10,  0.50,  0.30],
        ], dtype=np.float64)
        w_v = np.array([
            [ 0.40,  0.20, -0.10,  0.30],
            [ 0.10, -0.50,  0.60,  0.20],
            [-0.20,  0.30,  0.40, -0.10],
            [ 0.50,  0.10,  0.20, -0.40],
        ], dtype=np.float64)
        b_q = np.array([0.10, -0.05, 0.20, 0.15], dtype=np.float64)
        b_k = np.array([-0.10, 0.07, 0.03, 0.12], dtype=np.float64)
        b_v = np.array([0.02, 0.10, -0.08, 0.05], dtype=np.float64)
        w_o = np.array([
            [ 0.70, -0.20,  0.10,  0.30],
            [-0.30,  0.40,  0.50, -0.10],
            [ 0.20,  0.60, -0.40,  0.15],
            [ 0.10, -0.35,  0.20,  0.55],
        ], dtype=np.float64)
        b_o = np.array([0.05, -0.02, 0.07, -0.04], dtype=np.float64)
    
        seq_len, d_model = x.shape
        num_heads = 2
        head_dim = d_model // num_heads
        assert d_model % num_heads == 0
    
        # Linear projections: each token x_t maps independently to q_t, k_t, and v_t.
        q = x @ w_q.T + b_q
        k = x @ w_k.T + b_k
        v = x @ w_v.T + b_v
        # Split final feature dimension into heads: [T, H, Dh] -> [H, T, Dh].
        qh = q.reshape(seq_len, num_heads, head_dim).transpose(1, 0, 2)
        kh = k.reshape(seq_len, num_heads, head_dim).transpose(1, 0, 2)
        vh = v.reshape(seq_len, num_heads, head_dim).transpose(1, 0, 2)
    
        # Per-head scaled dot product: [H, T, Dh] @ [H, Dh, T] -> [H, T, T].
        raw_scores = qh @ kh.transpose(0, 2, 1)
        scores = raw_scores / np.sqrt(head_dim)
        weights = softmax_last(scores)
        contexts = weights @ vh                         # [H, T, Dh]
        concatenated = contexts.transpose(1, 0, 2).reshape(seq_len, d_model)
        output = concatenated @ w_o.T + b_o            # output projection W_O
    
        # PyTorch expects [batch, sequence, feature] when batch_first=True.
        torch.set_num_threads(1)
        tx = torch.tensor(x[None, :, :], dtype=torch.float64, device="cpu")
        mha = nn.MultiheadAttention(
            embed_dim=d_model,
            num_heads=num_heads,
            dropout=0.0,
            bias=True,
            batch_first=True,
            device="cpu",
            dtype=torch.float64,
        )
        with torch.no_grad():
            # PyTorch packs Q/K/V projection weights and biases as [3*E, E] / [3*E].
            mha.in_proj_weight.copy_(torch.tensor(np.concatenate([w_q, w_k, w_v], axis=0)))
            mha.in_proj_bias.copy_(torch.tensor(np.concatenate([b_q, b_k, b_v], axis=0)))
            mha.out_proj.weight.copy_(torch.tensor(w_o))
            mha.out_proj.bias.copy_(torch.tensor(b_o))
        mha.eval()
        with torch.no_grad():
            torch_output, torch_weights = mha(
                tx, tx, tx,
                need_weights=True,
                average_attn_weights=False,
            )
            projected = torch.nn.functional.linear(tx, mha.in_proj_weight, mha.in_proj_bias)
            tq, tk, tv = projected.chunk(3, dim=-1)
            tqh = tq.reshape(1, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
            tkh = tk.reshape(1, seq_len, num_heads, head_dim).permute(0, 2, 1, 3)
            torch_scores = torch.matmul(tqh, tkh.transpose(-2, -1)) / np.sqrt(head_dim)
            torch_head_weights = torch.softmax(torch_scores, dim=-1)
    
        # Independent numeric checks: this is a single deterministic hand-set case.
        max_score_diff = float(np.max(np.abs(scores - torch_scores[0].cpu().numpy())))
        max_weight_diff = float(np.max(np.abs(weights - torch_weights[0].cpu().numpy())))
        max_projected_q_diff = float(np.max(np.abs(q - tq[0].cpu().numpy())))
        max_projected_k_diff = float(np.max(np.abs(k - tk[0].cpu().numpy())))
        max_projected_v_diff = float(np.max(np.abs(v - tv[0].cpu().numpy())))
        max_output_diff = float(np.max(np.abs(output - torch_output[0].cpu().numpy())))
        tolerance = 1e-12
        assert max_score_diff < tolerance
        assert max_weight_diff < tolerance
        assert max_projected_q_diff < tolerance
        assert max_projected_k_diff < tolerance
        assert max_projected_v_diff < tolerance
        assert max_output_diff < tolerance
        assert shape(q) == [3, 4] and shape(qh) == [2, 3, 2]
        assert shape(scores) == [2, 3, 3] and shape(contexts) == [2, 3, 2]
        assert shape(concatenated) == [3, 4] and shape(output) == [3, 4]
        assert list(torch_weights.shape) == [1, 2, 3, 3]
    
        result = {
            "scope": "Three hand-set token vectors; self-attention with 2 heads; NumPy equations compared to torch.nn.MultiheadAttention. No training, corpus, positional encoding, mask, task metric, speed test, or benchmark.",
            "environment": {
                "python": platform.python_version(),
                "numpy": np.__version__,
                "torch": torch.__version__,
                "device": "cpu",
                "dtype": "float64",
                "batch_first": True,
                "dropout": 0.0,
                "module_eval": True,
                "need_weights": True,
                "average_attn_weights": False,
            },
            "dimensions": {
                "batch": 1,
                "sequence_length": seq_len,
                "d_model": d_model,
                "num_heads": num_heads,
                "head_dim": head_dim,
                "shapes": {
                    "X": shape(x), "Q": shape(q), "K": shape(k), "V": shape(v),
                    "Q_heads_K_heads_V_heads": shape(qh),
                    "scores_per_head": shape(scores), "weights_per_head": shape(weights),
                    "contexts_per_head": shape(contexts), "concat_heads": shape(concatenated),
                    "output_after_WO": shape(output),
                    "torch_output": list(torch_output.shape),
                    "torch_weights_per_head": list(torch_weights.shape),
                },
                "meaning": "The first dimension of projected tensors is token position; the head axis holds 2 separate 2D subspaces. Each score/weight matrix is [head, query-token, key-token].",
            },
            "parameters_row_vectors_use_X_times_W_transpose": {
                "X": rounded(x),
                "W_Q_out_by_in": rounded(w_q), "b_Q": rounded(b_q),
                "W_K_out_by_in": rounded(w_k), "b_K": rounded(b_k),
                "W_V_out_by_in": rounded(w_v), "b_V": rounded(b_v),
                "W_O_out_by_in": rounded(w_o), "b_O": rounded(b_o),
                "note": "All Q/K/V and output matrices are dense, explicitly overwritten, and not identity or diagonal matrices. These are illustrative fixed parameters, not learned weights.",
            },
            "projected_Q_K_V": {"Q": rounded(q), "K": rounded(k), "V": rounded(v)},
            "per_head": [
                {
                    "head": head,
                    "Q": rounded(qh[head]),
                    "K": rounded(kh[head]),
                    "V": rounded(vh[head]),
                    "raw_QK_transpose_scores": rounded(raw_scores[head]),
                    "scaled_scores_divide_sqrt_head_dim": rounded(scores[head]),
                    "softmax_over_key_token_axis": rounded(weights[head]),
                    "weight_row_sums": rounded(weights[head].sum(axis=-1)),
                    "weighted_value_context": rounded(contexts[head]),
                }
                for head in range(num_heads)
            ],
            "head_concat_and_output_projection": {
                "concatenated_heads_before_WO": rounded(concatenated),
                "W_O_times_concat_plus_b_O": rounded(output),
                "torch_mha_output": rounded(torch_output[0].cpu().numpy()),
                "torch_per_head_weights": rounded(torch_weights[0].cpu().numpy()),
                "torch_manual_scores_from_same_packed_QK": rounded(torch_scores[0].cpu().numpy()),
            },
            "max_abs_difference_numpy_vs_torch": {
                "Q": max_projected_q_diff,
                "K": max_projected_k_diff,
                "V": max_projected_v_diff,
                "scaled_scores": max_score_diff,
                "per_head_weights": max_weight_diff,
                "final_output_after_WO": max_output_diff,
                "assertion_tolerance": tolerance,
            },
            "checks": "PASS: dense hand-set Q/K/V and W_O projections, per-head scores/weights/contexts, concatenation and output projection match PyTorch MultiheadAttention in the single specified CPU/float64 case.",
        }
        print(json.dumps(result, ensure_ascii=False, indent=2, allow_nan=False))
    
    
    if __name__ == "__main__":
        main()
    

    程式輸出的第一列權重和每列權重總和都可直接檢查;NumPy 計算與 PyTorch module 的 Q/K/V、縮放分數、每-head 權重和 W_O 後輸出最大絕對差均小於 1e-12。這只說明此組手設值在指定 CPU/float64 執行環境下兩種計算一致,不是任務表現、訓練收斂、head 專業化或速度比較。完整逐步輸出與程式執行紀錄保存在本機研究附件,不用附件中的數字冒充學到的模型結果。

    本文圖像來源

    D2L 四張圖均為固定版教材的原圖轉檔或原圖等比置入封面;四作者為 Aston Zhang、Zachary C. Lipton、Mu Li、Alexander J. Smola,圖片依 CC BY-SA 4.0 使用,並對改作圖維持相同授權。TensorFlow 圖是該教材頁自有示意,作者 TensorFlow Authors,依 CC BY 4.0 使用;兩種授權分別適用,不混為一種授權。

    留言0

    登入後就能留言。

    登入
    sunny 桑尼
    sunny 桑尼

    AI 資料科學 · 產品經理 · 蒙奇 D 桑尼

    專注於 AI 工具應用與資料科學教育,致力於讓每個人都能用 AI 提升工作效率。曾任職於科技業,現為全職教育創作者。

    讀更多文章 →想自己動手試?領免費 Starter Kit →