English· Español· Deutsch· Nederlands· Français· 日本語· ქართული· 繁體中文· 简体中文· Português· Русский· العربية· हिन्दी· Italiano· 한국어· Polski· Svenska· Türkçe· Українська· Tiếng Việt· Bahasa Indonesia

un

guest
1 / ?
back to lessons

寬度跨頭部分割

單一頭部看到一種模式

活動 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 整除,餘數為零。違反此規則的配置會在配置驗證時失敗,而不是運行時。


三個模型,三種分割


Variantd_modeln_headhead_dim
ANDREA-12M3841232
ANDREA-120M7681264
ANDREA-480M15362464

注意:ANDREA-12M 和 ANDREA-120M 保持 n_head=12 不變;僅 d_model 因此 head_dim 會擴展。ANDREA-480M 將頭數加倍至 24,保持 head_dim=64 與 ANDREA-120M 匹配。

計算 head_dim

假設一個假設的 ANDREA-240M 變體使用 `d_model=1024` 和 `n_head=16`。計算 `head_dim`。然後用一句話解釋如果一位同事提出 `d_model=1024` 和 `n_head=15` 時,在配置驗證中會發生什麼。

每個頭三個矩陣,或一個大矩陣

每個頭的視圖

每個頭需要自己的查詢投影、鍵投影和值投影。對於頭 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_QW_KW_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 子層與嵌入層。

命名投影矩陣

ANDREA-120M 具有 `d_model=768` & `n_head=12`。請命名一個 transformer 區塊注意力子層中的四個投影矩陣,並給出每個矩陣的形狀(行 × 列)。然後說明這四個矩陣中哪個使用輸入 X,以及哪個使用串聯的頭輸出。

十二個向量變成一個

Multi-Head Attention Flow


每個頭計算後

每個頭會產生一個形狀為 [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 傳播,就會排序頭。

為什麼是十二個頭,而不是一個更寬的頭

想像一個替代方案:用 1 個 768 維的頭取代每個 64 維的 12 個頭。投影中的總參數保持相同。給出兩個不同的理由,解釋為什麼 12 頭配置在語言建模任務上優於 1 頭配置。在你的答案中參考注意力模式或訓練動態。

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 每層執行所有五個步驟,十二層深,每個前向傳遞一次。

閱讀管道

逐步說明在 ANDREA-120M 的一個多頭注意力子層中,單一 token 的表示 X[i](一個 768 維向量)流經時會發生什麼。提及:它經歷的投影、如何分割、注意力對它的作用,以及如何返回 768 維向量。在每個步驟顯示維度。