MicroGPT 運作圖解

MicroGPT 不依賴任何框架，從零實作計算圖（Computation Graph）與反向傳播（Backpropagation）。

Value 節點結構

class Value:
    __slots__ = ('data', 'grad', '_children', '_local_grads')

    def __init__(self, data, children=(), local_grads=()):
        self.data        = data          # 正向傳播的純量值
        self.grad        = 0            # ∂Loss/∂self (反向傳播填入)
        self._children   = children     # 計算圖的子節點
        self._local_grads = local_grads # 對子節點的局部偏導數

基礎運算的局部梯度

Topological Sort 反向傳播

先對計算圖做拓撲排序，再從 Loss 節點反向遍歷，用 Chain Rule 累積梯度。

def backward(self):
    topo = []
    visited = set()
    def build_topo(v):                    # DFS 建立拓撲序
        if v not in visited:
            visited.add(v)
            for child in v._children: build_topo(child)
            topo.append(v)
    build_topo(self)
    self.grad = 1                          # dLoss/dLoss = 1
    for v in reversed(topo):              # 從 Loss 往輸入反向
        for child, local_grad in zip(v._children, v._local_grads):
            child.grad += local_grad * v.grad  # Chain Rule: ∂L/∂child += local * ∂L/∂v

互動：Chain Rule 計算圖展示

以 loss = -log(softmax(z)[y]) 為例，點擊節點查看梯度傳播：

MicroGPT 實作

for h in range(n_head):
    hs = h * head_dim
    q_h = q[hs:hs+head_dim]                         # Query 切片
    k_h = [ki[hs:hs+head_dim] for ki in keys[li]]  # 所有 K
    v_h = [vi[hs:hs+head_dim] for vi in values[li]] # 所有 V

    # Scaled dot-product: Q·Kᵀ / √d_k
    attn_logits = [
        sum(q_h[j] * k_h[t][j] for j in range(head_dim)) / head_dim**0.5
        for t in range(len(k_h))
    ]
    attn_weights = softmax(attn_logits)              # 各 token 的注意力分數
    head_out = [
        sum(attn_weights[t] * v_h[t][j] for t in range(len(v_h)))
        for j in range(head_dim)
    ]                                                 # 加權聚合 V

為什麼要除以 √d_k？

注意力矩陣互動視覺化

輸入一段文字，查看 causal attention mask（下三角）與模擬注意力熱力圖：

在自回歸推論（auto-regressive decoding）時，每次只輸入「新 token」，過去 token 的 K、V 向量已算好並緩存，不需重算。

MicroGPT 的 KV Cache 實作

# 初始化 KV Cache（每層各一個 list）
keys   = [[] for _ in range(n_layer)]   # keys[layer] = list of k vectors
values = [[] for _ in range(n_layer)]   # values[layer] = list of v vectors

# 每次前向傳播：append 新 K, V
keys[li].append(k)     # 新 K 加入 Cache
values[li].append(v)   # 新 V 加入 Cache

# Attention 使用全部歷史 K, V
k_h = [ki[hs:hs+head_dim] for ki in keys[li]]   # 讀 Cache
v_h = [vi[hs:hs+head_dim] for vi in values[li]]  # 讀 Cache

互動演示：逐步生成 token 觀察 Cache 增長

KV Cache 以 記憶體（存 K、V 向量）換取 時間（省略重複投影），是現代 LLM 推論加速的關鍵技術。

調整 logits 觀察 softmax 輸出如何變化：

def softmax(logits):
    max_val = max(val.data for val in logits)  # 減去 max 防止 exp 溢位
    exps    = [(val - max_val).exp() for val in logits]
    total   = sum(exps)
    return [e / total for e in exps]

為何在殘差後仍使用 RMSNorm？

程式碼中特別標注 # note: not redundant due to backward pass via the residual connection。殘差連接讓梯度可以直接流過，但前向路徑上的 norm 仍有意義——確保 attention 輸入的尺度一致。

losses = []
for pos_id in range(n):
    token_id, target_id = tokens[pos_id], tokens[pos_id + 1]
    logits = gpt(token_id, pos_id, keys, values)
    probs  = softmax(logits)
    loss_t = -probs[target_id].log()   # −log p(correct token)
    losses.append(loss_t)
loss = (1 / n) * sum(losses)            # 平均 loss

temperature = 0.5   # τ ∈ (0,1]，越低越確定性
probs = softmax([l / temperature for l in logits])
token_id = random.choices(range(vocab_size), weights=[p.data for p in probs])[0]

$\tau \to 0$：模型永遠選最高機率 token（貪婪）；$\tau = 1$：標準分布；$\tau > 1$：更隨機/創意。

Temperature 互動展示

MicroGPT 實作

learning_rate = 0.01
beta1, beta2, eps_adam = 0.85, 0.99, 1e-8
m = [0.0] * len(params)   # 一階矩緩衝 (momentum)
v = [0.0] * len(params)   # 二階矩緩衝 (variance)

# 每個 step 更新
lr_t = learning_rate * (1 - step / num_steps)  # 線性 LR decay
for i, p in enumerate(params):
    m[i] = beta1 * m[i] + (1 - beta1) * p.grad        # 更新 m
    v[i] = beta2 * v[i] + (1 - beta2) * p.grad ** 2  # 更新 v
    m_hat = m[i] / (1 - beta1 ** (step + 1))          # 偏差修正
    v_hat = v[i] / (1 - beta2 ** (step + 1))          # 偏差修正
    p.data -= lr_t * m_hat / (v_hat ** 0.5 + eps_adam)  # 參數更新
    p.grad = 0                                          # 梯度清零

互動：模擬 Adam 更新過程

MicroGPT 使用最簡單的線性衰減：訓練開始時 lr = 0.01，最後一步趨近於 0，防止後期震盪。

訓練完成後，模型以 BOS token 為起點，逐步採樣生成新 token，直到再次出現 BOS 為止。

temperature = 0.5
for sample_idx in range(20):
    keys, values = [[] for _ in range(n_layer)], [[] for _ in range(n_layer)]
    token_id = BOS          # 起始 token
    sample = []
    for pos_id in range(block_size):
        logits  = gpt(token_id, pos_id, keys, values)   # 前向 (KV Cache 累積)
        probs   = softmax([l / temperature for l in logits])
        token_id = random.choices(range(vocab_size),
                     weights=[p.data for p in probs])[0]
        if token_id == BOS: break  # EOS = BOS
        sample.append(uchars[token_id])

互動：模擬推論步驟動畫