實作注意力機制Coding attention mechanisms

整本書的三階段地圖裡，本章是 Stage 1 的第 2 步——也是整個架構裡最關鍵的零件。第 2 章的資料處理已完成，attention 做好之後，第 4 章把它裝進 GPT：

本章用四段進化把 attention 蓋出來，每一段都在前一段上加一個零件——會隨你往下讀自動打勾（也可以點擊跳到該節）：

先看 attention 出現之前的世界。以翻譯為例：逐字直譯行不通——德文和英文的語序、文法結構不同，某些字必須「看到句子其他位置」才能翻對。transformer 之前的主流解法是 encoder–decoder RNN：encoder 把整句讀完、壓成一顆 hidden state，decoder 再從這顆向量開始逐字生出翻譯。切換三種做法，看問題出在哪：

2014 年的 Bahdanau attention 對症下藥：讓 decoder 每生成一個字，都能直接回看輸入序列的全部位置，並用一組權重決定「現在該多看誰」——上一節第 ③ 個分頁演的就是這件事。三年後（2017）研究者發現更激進的事實：RNN 根本可以整個拿掉，只靠 attention 就能建出更強的架構——transformer 就此誕生。

全章的工作句子："Your journey starts with one step"——六個 token，每個已經是 3 維 embedding（書中為了排版選了迷你維度；真實 GPT 是 768 維起跳）。我們指定第 2 個 token journey（x²）當 query：「以 journey 的視角，該注意誰？」

第一步的核心：點積＝相似度

attention score 的算法出奇地簡單：query 和每個 token 的 embedding 做點積（對應元素相乘再相加）。點積衡量兩個向量的「對齊程度」——方向愈接近、值愈大。滑過任何一個 token，看 journey 對它的分數怎麼算出來：

三步組出 context vector（書中真實數值）

拿到分數之後還差兩步：把分數正規化成總和為 1 的權重，再用權重對所有 embedding 加權求和。按「下一步」逐步推導：

剛才只服務了 journey 一個 query。但每個 token 都需要自己的 context vector——把「query 對所有 token 點積」推廣成「所有 token 對所有 token」，就得到 6×6 的注意力矩陣。for 迴圈太慢，一個矩陣乘法 inputs @ inputs.T 全包。切換分數／權重，點左側 token 選 query 列，滑過格子看數值（底色深淺＝數值大小；書中真實輸出）：

# 整個 3.3 節，濃縮成三行（形狀： [6,3] @ [3,6] → [6,6] → [6,6] @ [6,3] → [6,3]）
attn_scores  = inputs @ inputs.T                    # 所有 pair 的點積
attn_weights = torch.softmax(attn_scores, dim=-1)   # dim=-1：對「每一列」正規化
all_context_vecs = attn_weights @ inputs            # 每列權重 × 所有 embedding

簡化版有個致命限制：「相似度」被 embedding 寫死了，模型沒有任何可以學的東西。升級的方式：點積前先讓向量過三個可訓練的投影矩陣 W_q / W_k / W_v——它們是注意力模組裡唯一的權重參數，訓練時逐步調整，讓模型「學會」該用什麼標準去衡量相關性。

照舊以 journey（x²）當主角，d_in=3、d_out=2（書中故意讓輸出維度不同，方便看清楚形狀；真實 GPT 裡 d_in=d_out）。按「下一步」逐步推導（書中真實數值，seed 123）：

為什麼要除以 √d_k？拖拖看

維度 d_k 愈大，點積是愈多項的加總，數值天然愈大；大分數進 softmax 會把分佈推成幾乎 one-hot——softmax 進入「飽和區」，梯度趨近 0，模型學不動。除以 √d_k 恰好把分數的變異數拉回原尺度。拖動 d_k（以 journey 的六個分數模擬放大，softmax 為即時計算）：

包成類別：SelfAttention_v1 → v2

class SelfAttention_v1(nn.Module):
    def __init__(self, d_in, d_out):
        super().__init__()
        self.W_query = nn.Parameter(torch.rand(d_in, d_out))   # W_key / W_value 同
    def forward(self, x):
        queries = x @ self.W_query ; keys = x @ self.W_key ; values = x @ self.W_value
        attn_scores  = queries @ keys.T
        attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
        return attn_weights @ values               # ← context vectors

GPT 的任務是預測下一個字。訓練時整句一起算，如果 journey 的 context vector 裡混進了 starts、step 的資訊，就等於作弊看到未來——預測「journey 的下一個字」時答案已經洩題。解法：把注意力矩陣對角線以上（每列的未來位置）遮掉。這叫 causal attention（也稱 masked attention）。

兩種遮法對照：為什麼 −inf 才是標準答案

直覺做法（A）：先照常 softmax，再把上三角歸零、每列重新除以列和——可行，但要三步。標準做法（B）：在分數階段就把上三角填 −inf，softmax 一步到位。切換做法、按步驟走（書中真實數值，權重來自上一節的 sa_v2）：

再加一層保險：dropout

訓練時隨機把一部分注意力權重歸零，逼模型不要過度依賴任何特定位置——這是對抗 overfitting 的經典手段。被留下的值要 ×1/(1−p) 補償，讓整列的期望值不變。拖動 dropout 率、按「重抽」換遮罩（在 causal 權重上模擬；瀏覽器端固定種子亂數）：

包成類別：CausalAttention（點一行看說明）

把 mask 和 dropout 收進類別，同時讓 forward 支援 batch 輸入——第 2 章 data loader 吐出來的可是 [batch, tokens, d_in] 的三維 tensor（書中用 batch = torch.stack((inputs, inputs)) → [2, 6, 3] 模擬）：

一組 W_q/W_k/W_v 只能學一種「相關性的標準」。multi-head attention 讓多組投影並行——有的 head 可能學語法關係、有的學指代關係——各自算出 context 後串接起來。最直觀的實作是疊 wrapper：

class MultiHeadAttentionWrapper(nn.Module):
    def __init__(self, d_in, d_out, context_length, dropout, num_heads, qkv_bias=False):
        super().__init__()
        self.heads = nn.ModuleList([CausalAttention(d_in, d_out, context_length, dropout, qkv_bias)
                                    for _ in range(num_heads)])   # num_heads 個獨立單頭
    def forward(self, x):
        return torch.cat([head(x) for head in self.heads], dim=-1) # 沿最後一維串接

num_heads=2、每頭 d_out=2 → 串接後每個 token 變成 2×2＝4 維（書中真實輸出，seed 123；顏色區分兩個 head）：

高效版：一次大投影，再「切」出多頭

wrapper 的缺點：for head in self.heads 是串行跑 num_heads 次矩陣乘法。MultiHeadAttention 類別改成：只做一次大投影（d_out 維），再用 view / transpose 把它「切」成 num_heads 份，批次矩陣乘法一次算完所有 head。點頭數，看形狀流與參數量（以 GPT-2 尺寸 d_out=768 示範）：

# MultiHeadAttention.forward 的核心（d_out 必須能被 num_heads 整除）
keys = self.W_key(x)                    # [b, T, d_out]——一次大投影
keys = keys.view(b, T, num_heads, head_dim)        # 「切」出 head 維（免費，只是重新解讀）
keys = keys.transpose(1, 2)                        # [b, heads, T, head_dim]——head 提前，方便批次運算
attn_scores = queries @ keys.transpose(2, 3)       # [b, heads, T, T]：每頭一張注意力矩陣，一次算完
attn_scores.masked_fill_(mask_bool, -torch.inf)    # causal mask 照舊
attn_weights = torch.softmax(attn_scores / head_dim**0.5, dim=-1)   # √d_k 用的是 head_dim
context = (attn_weights @ values).transpose(1, 2)  # 換回 [b, T, heads, head_dim]
context = context.contiguous().view(b, T, d_out)   # 串接＝把 head 維攤平
return self.out_proj(context)                      # 出場前再過一個 Linear 融合各頭（慣例配備）

• 為什麼要 attention：RNN 把整句壓進一顆固定 hidden state，長句必失真；attention 讓每個位置直接回看所有位置。
• 一條主線公式：分數（點積）→ softmax 正規化 → 對 value 加權求和 → context vector。全章四個版本都是它的加料版。
• 簡化版：分數＝embedding 兩兩點積（inputs @ inputs.T）；journey 的 z² = [0.4419, 0.6515, 0.5683]。
• 可訓練版（scaled dot-product）：x 先過 W_q/W_k/W_v 投影；分數除以 √d_k 防止 softmax 飽和、梯度消失；nn.Linear 版（v2）初始化更好。
• causal mask：分數上三角填 −inf 再 softmax（一步到位、數學上等價於「只在可見範圍 softmax」）——生成式 LLM 不准偷看未來。
• dropout：訓練時隨機丟權重、倖存值 ×1/(1−p) 補償；推論時關閉。
• multi-head：d_out 切成 num_heads × head_dim，view / transpose 批次算完；參數量與頭數無關。GPT-2 124M：12 頭 × 64 維。

// 七題，作答後立即看解析