寬度跨頭部分割
單一頭部看到一種模式
活動 67 涵蓋了縮放點積注意力:一個查詢向量 Q、一個鍵向量 K、一個值向量 V;計算 Q·Kᵀ/√d_k,遮罩,softmax,加權 V。一個頭學習一種關係模式:也許是主詞-動詞一致,也許是標點配對,也許沒什麼用處。
多頭注意力在平行中運行相同的操作數次,對一個 token 表示的不同切片。十二個平行頭部。十二種可能的關係模式。頭部僅透過訓練壓力專門化;沒有架構師告訴頭部 4 要關注動詞時態。
分割關係
ANDREA-120M 設定 d_model = 768 & n_head = 12。多頭注意力將 768 分割成 12 個每個 64 的區塊:
head_dim = d_model / n_head
64 = 768 / 12
每個頭操作 64 維向量。分割應用得很乾淨:d_model 必須能被 n_head 整除,餘數為零。違反此規則的配置會在配置驗證時失敗,而不是運行時。
三個模型,三種分割
| Variant | d_model | n_head | head_dim |
|---|---|---|---|
| ANDREA-12M | 384 | 12 | 32 |
| ANDREA-120M | 768 | 12 | 64 |
| ANDREA-480M | 1536 | 24 | 64 |
注意:ANDREA-12M 和 ANDREA-120M 保持 n_head=12 不變;僅 d_model 因此 head_dim 會擴展。ANDREA-480M 將頭數加倍至 24,保持 head_dim=64 與 ANDREA-120M 匹配。
計算 head_dim
每個頭三個矩陣,或一個大矩陣
每個頭的視圖
每個頭需要自己的查詢投影、鍵投影和值投影。對於頭 h:
Q_h = X · W_Q^h 其中 W_Q^h 的形狀為 [d_model, head_dim]
K_h = X · W_K^h 其中 W_K^h 的形狀為 [d_model, head_dim]
V_h = X · W_V^h 其中 W_V^h 的形狀為 [d_model, head_dim]
X 承載輸入形狀 [batch, seq_len, d_model]。投影後,Q_h、K_h、V_h 各自承載形狀 [batch, seq_len, head_dim]。
融合視圖
每個頭的矩陣在記憶體中並排存放。單一融合矩陣 W_Q 形狀為 [d_model, d_model],一次產生所有頭:
Q_fused = X · W_Q # [batch, seq_len, d_model]
Q_per_head = reshape(Q_fused) # [batch, n_head, seq_len, head_dim]
融合矩陣乘法只需發送一個 BLAS 呼叫,而不是 12 個。CUDA tensor cores 在這種大小的矩陣乘法上達到峰值吞吐量;每個頭部的矩陣乘法會低度使用硬體。
參數數量
三個融合矩陣 W_Q、W_K、W_V,每個 d_model × d_model。加上輸出投影 W_O,也是 d_model × d_model。對於 ANDREA-120M:
每層注意力參數 = 4 × 768² = 2,359,296 ≈ 2.36M
12 層總參數 = 12 × 2.36M ≈ 28.3M
ANDREA-120M 總參數的大約四分之一位於注意力投影中。剩餘四分之三位於 MLP 子層與嵌入層。
命名投影矩陣
十二個向量變成一個
每個頭計算後
每個頭會產生一個形狀為 [batch, seq_len, head_dim] 的輸出張量。十二個頭會產生十二個這樣的張量。沿著特徵維度進行串聯,將它們重新堆疊在一起:
concat_output = concat(head_1, head_2, ..., head_12)
shape = [batch, seq_len, n_head × head_dim]
= [batch, seq_len, 768] # for ANDREA-120M
Concat 會逆轉 split。總特徵維度回到 d_model。維度上沒有資訊損失;差異在於每個區塊包含的內容:head 1 的區塊反映 head 1 學習到的注意力模式。
輸出投影 W_O
單純的 Concat 讓頭部保持隔離:head 4 的輸出緊鄰 head 7 的輸出,彼此毫無認知。形狀為 [d_model, d_model] 的輸出投影 W_O 會將它們混合:
attention_output = concat_output · W_O
shape = [batch, seq_len, d_model]
在 W_O 之後,每個輸出維度都攜帶所有十二個頭的學習線性組合。資訊透過這個單一矩陣乘法在頭之間自由流動。
為什麼頭會專門化
架構中沒有任何東西強迫頭 4 學習動詞時態或頭 9 學習配對標點。專門化從梯度壓力中出現:在訓練期間,貢獻冗餘的頭接收比貢獻獨特的頭更小的梯度。經過數千步,每個頭會安頓到一個最有效地減少總損失的利基。
經驗上,訓練好的 transformer 顯示出處理以下的頭:位置模式(頭看前一個 token)、語法模式(頭看匹配的結束括號)、語義模式(頭看最近的名稱實體)。沒有標籤訓練這種專門化。僅訓練信號,透過 W_O 傳播,就會排序頭。
為什麼是十二個頭,而不是一個更寬的頭
CUDA 如何儲存頭部
單一張量,重新塑形
ANDREA 的訓練引擎 microgpt_cuda.cu 並未為十二個頭部分配十二個獨立的緩衝區。它分配一個融合張量,並將頭部維度視為步幅模式:
// 在 Q = X · W_Q (一次矩陣乘法,跨頭部融合) 之後
// Q 的形狀為 [batch, seq_len, d_model]
// 重塑為 [batch, seq_len, n_head, head_dim]
// 轉置為 [batch, n_head, seq_len, head_dim]
// 每個頭現在在內部兩個維度中是連續的
轉置將 n_head 移到 seq_len 前面。為什麼?因為下一個操作 (Q_h · K_h^T) 需要每個頭的 seq_len × head_dim 切片在記憶體中連續。CUDA 矩陣乘法在連續張量上運行更快。
一個核心,多個頭
單一注意力 CUDA 核心會平行處理所有頭部。每個執行緒區塊處理一個 (batch, head) 對;區塊內的執行緒會合作處理 seq_len × head_dim 圖塊。核心從不知道它在處理多個頭部;啟動網格負責平行處理。
配置反映硬體
ANDREA-120M 選擇 n_head=12, head_dim=64 與 RTX 4090 張量核心相符,後者偏好 matmul 圖塊為 16 的倍數。head_dim=64 = 4 × 16 完全符合圖塊形狀。head_dim=32 (ANDREA-12M) 也符合但未充分利用圖塊。head_dim=72 不符合且會強制回退核心。
最終圖示
| 步驟 | 操作 | 輸出形狀 |
|---|---|---|
| 1. 專案 | Q = X · W_Q (同樣適用於 K, V) | [batch, seq, d_model] |
| 2. 重塑與轉置 | 將 d_model 分割 → (n_head, head_dim) | [batch, n_head, seq, head_dim] |
| 3. 每個頭的注意力 | 每個頭上的縮放點積 | [batch, n_head, seq, head_dim] |
| 4. 轉置與重塑 | 合併 (n_head, head_dim) → d_model | [batch, seq, d_model] |
| 5. 輸出投影 | output = concat · W_O | [batch, seq, d_model] |
五個步驟。三個矩陣乘法觸及輸入 (Q, K, V 投影)。一個矩陣乘法觸及串聯的頭部 (W_O)。一個注意力核心並行處理所有頭部。ANDREA-120M 每層執行所有五個步驟,十二層深,每個前向傳遞一次。