如果一個句子有三個 token,self-attention 會不會先替每個 token 算一個「重要度」?這個說法不太準確。更貼近計算的描述是:每個位置先產生一個 query,拿它和各位置的 key 比分,再把分數變成權重,用權重混合各位置的 value。每個 query 都有自己的權重列,所以不同位置可以讀取不同的內容。
本文先用圖分清連接範圍,再拆出 Q、K、V 的角色;接著用一組固定小數手算第一個 query,最後把流程延伸成兩個 head,核對每個張量的形狀、拼接與輸出投影。數字來自一段可重跑的 CPU 機制原型,權重是手動指定的,不是語料訓練結果,也不是翻譯或分類品質測試。
Self-attention 的「全序列」是連接方式
在 CNN 中,某個位置通常先和附近位置互動;RNN 則沿序列一步一步傳遞狀態。Self-attention 的每個位置可以直接對序列中的其他位置計算關聯。這說的是允許的資訊連接,不代表每個位置實際上都得到一樣大的權重。
「可以互相連接」和「模型已經知道詞序」是兩件事。如果只有 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 會固定學會某一種人類可命名的語法角色。
這個區分很重要:先由 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。
來看原型中的第一個 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。
把多個 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。
這次原型的張量形狀可整理如下:
| 階段 | 形狀 | 表示什麼 |
|---|---|---|
輸入 X 與完整 Q/K/V | 3 × 4 | 三個 token,每個四個特徵 |
| 切成兩個 head | 2 × 3 × 2 | [head, token, head 特徵] |
| 每個 head 的分數與權重 | 2 × 3 × 3 | [head, query 位置, key 位置] |
| 每個 head 的 context | 2 × 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
登入後就能留言。
登入