Transformer 架构

分词之后，一段文本变成了一串整数。语言模型的任务是：给定前面的 token，预测下一个 token 的概率分布。最朴素的模型只看前一个 token 查表（bigram），它能学会 q 后面常跟 u，但永远学不会”这个故事的主角叫 Lucy，所以后面的 She 指的是她”。要做到这一点，每个位置必须能从前面任意远的位置取信息，而且取哪些、取多少要由内容决定。这一篇从这个需求出发，一个零件一个零件地拼出一个现代 GPT（LLaMA 风格），最后在 CPU 上训练它写儿童故事。

前置：从线性回归到神经网络、序列模型与注意力的起源、线性代数

让 token 互相交流：从求平均到注意力

先用最笨的办法让位置 t 获得前文信息：把位置 0..t 的向量取平均。“只看前文”这个约束很关键，因为训练时我们让每个位置都预测它的下一个 token，如果位置 t 能看到 t+1，答案就泄露了。

用循环写很慢，但它可以写成一次矩阵乘：用一个下三角矩阵，每行归一化。

import torch
torch.manual_seed(0)
B, T, C = 1, 4, 2
x = torch.randn(B, T, C)
w = torch.tril(torch.ones(T, T))
w = w / w.sum(1, keepdim=True)
print(w)
print((w @ x)[0, 2], x[0, :3].mean(0))
# 等价写法：上三角填 -inf 再 softmax，-inf 经过 exp 变成 0
s = torch.zeros(T, T).masked_fill(torch.tril(torch.ones(T, T)) == 0, float("-inf"))
print(torch.softmax(s, -1)[2])

tensor([[1.0000, 0.0000, 0.0000, 0.0000],
        [0.5000, 0.5000, 0.0000, 0.0000],
        [0.3333, 0.3333, 0.3333, 0.0000],
        [0.2500, 0.2500, 0.2500, 0.2500]])
tensor([-0.5741, -0.3745]) tensor([-0.5741, -0.3745])
tensor([0.3333, 0.3333, 0.3333, 0.0000])

第三个写法是关键：权重矩阵来自 softmax(scores)，而 scores 现在全是 0，所以得到均匀平均。如果让 scores 依赖于数据，每个位置就能自己决定从哪些位置多取一点。这就是 self-attention。

每个位置的向量 x_i 用三个矩阵投影出三样东西：query q_i = x_i W_Q（我在找什么）、key k_j = x_j W_K（我包含什么）、value v_j = x_j W_V（如果你选中我，我给你什么）。位置 i 对位置 j 的打分是 q_i · k_j，经过因果 mask 和 softmax 变成权重，再对 value 加权求和：

$Attention (Q, K, V) = softmax (\frac{Q K ^{⊤}}{d _{k}} + M) V$

其中 $M$ 在上三角是 $- \infty$ ，其余是 0。mask 必须加在 softmax 之前：如果 softmax 之后再把未来位置的权重清零，剩下的权重和就不是 1 了。

为什么要除以 $d_{k}$ ？

import torch
torch.manual_seed(0)
for d in (16, 64, 256):
    q, k = torch.randn(10000, d), torch.randn(10000, d)
    s = (q * k).sum(-1)
    print(f"d={d:3d}  var(q.k)={s.var():7.1f}  var(q.k/sqrt(d))={(s / d**0.5).var():.2f}")
logits = torch.tensor([1.0, 2.0, 3.0])
print(torch.softmax(logits, -1), torch.softmax(logits * 16, -1))

d= 16  var(q.k)=   15.8  var(q.k/sqrt(d))=0.99
d= 64  var(q.k)=   63.5  var(q.k/sqrt(d))=0.99
d=256  var(q.k)=  261.3  var(q.k/sqrt(d))=1.02
tensor([0.0900, 0.2447, 0.6652]) tensor([1.2664e-14, 1.1254e-07, 1.0000e+00])

两个各分量独立、方差为 1 的向量，点积是 $d$ 项之和，方差是 $d$ 。不缩放的话，维度越大，scores 的绝对值越大，softmax 越接近 one-hot：一个位置拿走几乎全部权重，其他位置的梯度几乎为 0，训练初期就学不动。除以 $d_{k}$ 把方差拉回 1。

停一下，检查理解单选

训练语言模型时，self-attention 里的因果 mask 为什么必不可少？

因为不加 mask 时 softmax 会数值溢出因为训练时每个位置都在预测下一个 token，不加 mask 它能直接看到答案因为 mask 让注意力的计算量减半，是一种加速手段因为推理时只有一个 token，训练要和推理形状一致

多头，以及”注意力不知道顺序”

一组 Q/K/V 只能表达一种”关注方式”。但一个位置可能同时需要：找上一个名词、找句子的主语、找前一个引号。多头注意力把 d_model 维切成 h 份，每份 d_head = d_model / h 维，各自独立做注意力，结果拼回 d_model 维，再过一个输出矩阵 W_O 混合。参数量和单头一样（四个 d×d 矩阵），计算量也几乎一样，但表达力强得多。

注意力还有一个奇怪的性质：它完全不知道位置。把输入打乱顺序，输出只是跟着打乱，每个位置的结果不变。

import torch
torch.manual_seed(0)
T, d = 5, 8
x = torch.randn(T, d)
Wq, Wk, Wv = (torch.randn(d, d) for _ in range(3))
def attn(x):  # 不带 mask、不带位置信息
    q, k, v = x @ Wq, x @ Wk, x @ Wv
    return torch.softmax(q @ k.T / d**0.5, -1) @ v
perm = torch.tensor([4, 2, 0, 1, 3])
print(torch.allclose(attn(x)[perm], attn(x[perm]), atol=1e-5))

True

对注意力来说输入是一个集合，不是序列。“狗咬人”和”人咬狗”在它眼里一样（因果 mask 会带入一点顺序信息，但远远不够）。所以必须显式地注入位置。原始 Transformer 给每个位置加一个固定的正弦向量；GPT-2 给每个位置学一个 embedding，加到 token embedding 上。两者的问题一样：编码的是绝对位置，而语言里重要的通常是相对位置：“前一个词”在第 5 位和第 500 位应该是同一回事。

RoPE：用旋转编码相对位置

我们想要的是：位置 $m$ 的 query 和位置 $n$ 的 key 的内积只依赖 $m - n$ 。二维平面里有一个现成的工具：把两个向量各自旋转，内积只取决于它们的夹角之差。如果把位置 $m$ 的 query 旋转 $m θ$ ，位置 $n$ 的 key 旋转 $n θ$ ，那么

$⟨ R_{m θ} q, R_{n θ} k ⟩ = q^{⊤} R_{m θ}^{⊤} R_{n θ} k = q^{⊤} R_{(n - m) θ} k$

只依赖 $n - m$ 。这就是 RoPE（Rotary Position Embedding）。高维时把 d_head 维两两配对，第 $i$ 对用频率 $θ_{i} = 1000 0^{- 2 i / d}$ 旋转：前面的通道转得快，捕捉近距离；后面的转得慢，捕捉远距离，和正弦编码的多频率思想一样。RoPE 作用在每一层注意力内部的 q 和 k 上，不碰 value，也不加到残差流里。

import torch
def rope(x, pos, theta=10000.0):
    d = x.shape[-1]
    freqs = theta ** (-torch.arange(0, d, 2) / d)       # 每对通道一个频率
    ang = pos[:, None] * freqs[None, :]
    cos, sin = ang.cos(), ang.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], -1)
    return out.flatten(-2)
torch.manual_seed(0)
q, k = torch.randn(1, 64), torch.randn(1, 64)
for m, n in [(3, 1), (10, 8), (1002, 1000)]:
    s = (rope(q, torch.tensor([float(m)])) * rope(k, torch.tensor([float(n)]))).sum()
    print(f"m={m:4d} n={n:4d} m-n={m-n}  q.k={s:.4f}")

m=   3 n=   1 m-n=2  q.k=-10.1427
m=  10 n=   8 m-n=2  q.k=-10.1427
m=1002 n=1000 m-n=2  q.k=-10.1427

三组位置差都是 2，内积完全一样。RoPE 没有可学参数，cos/sin 表可以预先算好；它是 LLaMA 之后几乎所有开源模型的默认选择，长上下文扩展（调大 $θ$ 或对频率做插值）也都在它上面做。

停一下，检查理解多选

关于 RoPE，哪些说法正确？

它让 q 和 k 的内积只依赖两者的相对位置它没有可学习的参数它和 GPT-2 一样，把一个位置向量加到 token embedding 上它同时旋转 q、k 和 v

注意力之外：FFN、残差、Norm

注意力负责在位置之间搬运信息，但它本质是对 value 的加权平均，对每个位置自身的非线性加工很弱。所以每层还有一个逐位置的前馈网络（FFN）：先升维到 d_ff，过非线性，再降回 d_model。原始 Transformer 用 ReLU，d_ff = 4·d_model。

现代模型几乎都换成了门控版本 SwiGLU：

$FFN (x) = W_{2} (SiLU (W_{1} x) ⊙ W_{3} x), SiLU (z) = z σ (z)$

一支经过激活，另一支是纯线性的”门”，逐元素相乘。Shazeer 2020 的实验里 GLU 系列在同等计算下一致更好，作者自己的解释是”归功于神的仁慈”，至今没有公认的理论。门控多了一个矩阵（三个而不是两个），为了让参数量和计算量与原来的 4d 持平，把 d_ff 缩到 $\frac{2}{3} \times 4 d = \frac{8}{3} d$ ：原来是 $2 \cdot d \cdot 4 d = 8 d^{2}$ ，现在是 $3 \cdot d \cdot \frac{8}{3} d = 8 d^{2}$ 。实际实现会取整到 64 或 256 的倍数，对 GPU 的矩阵乘更友好。

把注意力和 FFN 叠几十层，靠的是两个东西。

残差连接：每个子层写成 x = x + f(x)。可以把那条贯穿所有层的 x 想成一条”残差流”，每个子层从流里读，再把自己的结果加回去。反向传播时，加法把梯度原样传下去，即使 f 的梯度很小，底层也能收到来自 loss 的信号。

归一化：每层的输出尺度如果不受控，几十层叠下来会漂移。原始 Transformer 是 post-norm：x = LN(x + f(x))，norm 放在残差相加之后，相当于残差主干本身每一层都被归一化一次，梯度要穿过几十个 LayerNorm 才能到底层。这种结构深了就难训练，需要小心的学习率 warmup。GPT-2 之后的模型几乎都是 pre-norm：x = x + f(LN(x))，norm 只放在支路的入口，残差主干从输入到输出是一条干净的加法链。代价是残差流的尺度会随层数增长，所以在最后输出前要再加一个 final norm。

Norm 本身也简化了。LayerNorm 减均值、除标准差、再乘 $γ$ 加 $β$ ；RMSNorm 去掉了减均值和 $β$ ，只除以均方根：

$RMSNorm (x) = \frac{x}{\frac{1}{d} \sum _{i} x _{i}^{2} + ϵ} ⊙ γ$

效果基本一样，但更简单。norm 在 FLOPs 里占比微不足道，但它要把整个激活读一遍写一遍，是受内存带宽限制的操作，少一次规约就少一次读写。同样的理由，现代模型的 Linear 普遍去掉了 bias。

停一下，检查理解单选

SwiGLU 的 FFN 有三个矩阵 W1、W2、W3。若希望它的参数量与 d_ff = 4·d_model 的普通两矩阵 FFN 相同，d_ff 应取多少？

4·d_model2·d_model8/3·d_model12·d_model

拼起来：一个能跑的 GPT

把上面的零件按 LLaMA 的方式组装：token embedding → N 个 pre-norm block（RMSNorm → 带 RoPE 的因果多头注意力 → 残差；RMSNorm → SwiGLU → 残差）→ final RMSNorm → 线性层输出 logits。下面的代码存成 gpt.py：

import math, torch
import torch.nn as nn
import torch.nn.functional as F
 
class RMSNorm(nn.Module):
    def __init__(self, d, eps=1e-5):
        super().__init__(); self.eps = eps; self.weight = nn.Parameter(torch.ones(d))
    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
 
def rope(x, theta=10000.0):              # x: (B, H, T, hd)
    T, hd = x.shape[-2], x.shape[-1]
    freqs = theta ** (-torch.arange(0, hd, 2, device=x.device) / hd)
    ang = torch.arange(T, device=x.device)[:, None] * freqs[None, :]
    cos, sin = ang.cos(), ang.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], -1).flatten(-2)
 
class Attention(nn.Module):
    def __init__(self, d, n_head):
        super().__init__()
        self.h = n_head
        self.qkv = nn.Linear(d, 3 * d, bias=False)   # 三个投影合成一次矩阵乘，GPU 上更高效
        self.proj = nn.Linear(d, d, bias=False)
    def forward(self, x):
        B, T, C = x.shape
        q, k, v = self.qkv(x).split(C, dim=-1)
        # (B, T, C) -> (B, H, T, hd)：每个头在自己的 hd 维子空间里独立做注意力
        q, k, v = (t.view(B, T, self.h, C // self.h).transpose(1, 2) for t in (q, k, v))
        q, k = rope(q), rope(k)
        att = q @ k.transpose(-2, -1) / math.sqrt(k.shape[-1])
        mask = torch.tril(torch.ones(T, T, dtype=torch.bool, device=x.device))
        att = att.masked_fill(~mask, float("-inf")).softmax(-1)
        y = (att @ v).transpose(1, 2).reshape(B, T, C)   # transpose 后不连续，只能 reshape
        return self.proj(y)
 
class SwiGLU(nn.Module):
    def __init__(self, d, d_ff):
        super().__init__()
        self.w1 = nn.Linear(d, d_ff, bias=False)
        self.w3 = nn.Linear(d, d_ff, bias=False)
        self.w2 = nn.Linear(d_ff, d, bias=False)
    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))
 
class Block(nn.Module):
    def __init__(self, d, n_head, d_ff):
        super().__init__()
        self.ln1, self.attn = RMSNorm(d), Attention(d, n_head)
        self.ln2, self.ffn = RMSNorm(d), SwiGLU(d, d_ff)
    def forward(self, x):
        x = x + self.attn(self.ln1(x))   # pre-norm：norm 只在支路上，残差主干保持干净
        x = x + self.ffn(self.ln2(x))
        return x
 
class GPT(nn.Module):
    def __init__(self, vocab, d=128, n_layer=4, n_head=4, d_ff=344):
        super().__init__()
        self.emb = nn.Embedding(vocab, d)
        self.blocks = nn.ModuleList(Block(d, n_head, d_ff) for _ in range(n_layer))
        self.ln_f = RMSNorm(d)
        self.head = nn.Linear(d, vocab, bias=False)
    def forward(self, idx, targets=None):
        x = self.emb(idx)
        for b in self.blocks:
            x = b(x)
        logits = self.head(self.ln_f(x))
        loss = None if targets is None else F.cross_entropy(logits.flatten(0, 1), targets.flatten())
        return logits, loss
    @torch.no_grad()
    def generate(self, idx, n, ctx):
        for _ in range(n):
            logits, _ = self(idx[:, -ctx:])     # 超过上下文长度就截断，RoPE 表只覆盖 ctx 以内
            nxt = torch.multinomial(logits[:, -1].softmax(-1), 1)
            idx = torch.cat([idx, nxt], 1)
        return idx

这里有一个形状细节：transpose(1, 2) 之后张量在内存里不再连续，.view 会报错，必须用 .reshape（原理见资源核算与 Scaling Laws 里的张量内存模型）。d_ff = 344 是 $\frac{8}{3} \times 128 \approx 341$ 向上取整到 8 的倍数。

训练数据用 TinyStories（GPT-4 生成的简单英文儿童故事）的前 100 万个字符，为了让例子短，直接按字符分词。训练目标是”输入右移一位”：

import time, torch
from gpt import GPT
torch.manual_seed(0)
text = open("stories.txt", encoding="utf-8").read()[:1_000_000]
chars = sorted(set(text)); stoi = {c: i for i, c in enumerate(chars)}
data = torch.tensor([stoi[c] for c in text])
n = int(0.9 * len(data)); train, val = data[:n], data[n:]
ctx, bs = 128, 32
def batch(d):
    ix = torch.randint(len(d) - ctx - 1, (bs,))
    x = torch.stack([d[i:i + ctx] for i in ix])
    return x, torch.stack([d[i + 1:i + ctx + 1] for i in ix])   # 目标 = 输入右移一位
model = GPT(len(chars))
print(f"vocab {len(chars)}, params {sum(p.numel() for p in model.parameters()):,}")
opt = torch.optim.AdamW(model.parameters(), lr=3e-3)
t = time.time()
for step in range(1501):
    x, y = batch(train)
    _, loss = model(x, y)
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 500 == 0:
        with torch.no_grad():
            vl = sum(model(*batch(val))[1].item() for _ in range(10)) / 10
        print(f"step {step:4d}  train {loss.item():.3f}  val {vl:.3f}")
print(f"{time.time()-t:.0f}s")
out = model.generate(torch.full((1, 1), stoi["O"]), 200, ctx)[0]
print("".join(chars[i] for i in out))

vocab 74, params 810,624
step    0  train 4.495  val 3.523
step  500  train 1.005  val 1.105
step 1000  train 0.953  val 1.038
step 1500  train 0.818  val 1.007
348s
One day, and the food. He wanted to make a magicmb. They tried to skip tretn on the kitchen.
One day, Lucy saw the bakery new friend. He smiled and smiled. Lucy was so happy to have fun with all the bo

先核对参数量：embedding 和输出层各 74×128 = 9472；每个 block 里注意力 4×128² = 65536、SwiGLU 3×128×344 = 132096、两个 norm 256，共 197888；4 层加 final norm，总数 810624，和打印的一致。初始 loss 4.495 接近均匀猜测的 ln 74 = 4.30（第 0 步的 val 是在一次更新之后测的，所以更低）。在一台笔记本的 CPU 上训练 6 分钟，80 万参数的模型已经会写句子结构、人名、引号和换行，偶尔拼错词。把字符级换成 BPE、把模型和数据放大几个数量级，就是 GPT-2 的配方。真正规模的训练还需要混合精度、多卡并行和 checkpoint，见并行训练。

停一下，检查理解单选

上面的 Attention.forward 里，如果把 (att @ v).transpose(1, 2).reshape(B, T, C) 改成 .view(B, T, C)，会发生什么？

结果一样，view 只是更快报错，因为 transpose 之后的张量在内存中不连续，view 无法不复制地改变形状不报错，但各个头的输出被错误地交织在一起不报错，但梯度无法回传

超参数怎么选：大家其实都差不多

一个 Transformer 的形状由几个数决定：d_model、n_layer、n_head、d_ff、词表大小、上下文长度。翻一遍公开模型的配置，会发现选择出奇地一致。

超参数	常见取值	为什么
`d_ff / d_model`	ReLU/GELU 用 4，GLU 用 8/3	前面推过，GLU 是为了与 4 倍的计算量持平
`n_head × d_head`	等于 `d_model`	多头只是切分，不增加参数；`d_head` 多为 64 或 128
`d_model / n_layer`	大约 100 到 200	太深则串行步骤多、难以并行；太宽则表达力利用率差；在这个区间内 loss 对形状不敏感
词表	单语 3 万到 5 万，多语言 10 万到 25 万	见分词里关于词表大小的权衡
dropout / weight decay	预训练通常不用 dropout，weight decay 多为 0.1	预训练数据只过一遍左右，几乎不会过拟合；weight decay 的作用更多是和学习率调度配合，让训练末期 loss 降得更好

例如 GPT-3 175B 是 d_model = 12288、96 层，比值 128；d_head = 128，96 个头。LLaMA 2 7B 是 4096 维、32 层，比值 128，d_ff = 11008（约为 2.69 倍，即 8/3 倍再向上取整）。

规模大了以后还要处理训练不稳定，几个常见补丁都瞄准 softmax：在输出层的 log-sum-exp 上加一个 z-loss，惩罚 logits 整体变大；在 q、k 进入点积之前各加一个 norm（QK norm），防止注意力 logits 爆炸；或者用 tanh 对 logits 做 soft-capping。它们的共同点是：softmax 对输入的尺度很敏感，而训练久了 logits 的尺度会漂移。

注意力本身的变体，比如为推理省内存的 GQA、MLA，长上下文用的滑动窗口和稀疏注意力，以及把 FFN 换成 MoE，见注意力变体与 MoE。注意力在长序列上被显存带宽卡住的问题见 FlashAttention。

停一下，检查理解多选

某团队设计一个 d_model = 4096 的模型，下列哪些配置属于主流选择？

32 个头，每个头 128 维使用 SwiGLU，d_ff 约为 1100864 个头，每个头 128 维，把注意力参数翻倍以提高表达力预训练时 dropout 设为 0.3 来防止过拟合

记住这三件事

self-attention 是”由数据决定权重的加权平均”：q·k 打分、除以 √d、加因果 mask、softmax、对 v 加权；它本身不知道顺序，位置要另外注入。
现代 GPT 的标准零件是 RoPE（旋转 q、k 以编码相对位置）、SwiGLU（门控 FFN，d_ff ≈ 8/3·d）、RMSNorm 与 pre-norm 残差流，Linear 不带 bias。
超参数有很强的共识：n_head × d_head = d_model、宽深比约 100 到 200、预训练不用 dropout；规模变大后的不稳定多来自 softmax 的尺度。

CS 与 LLM 学习知识库

探索

Transformer 架构

让 token 互相交流：从求平均到注意力

多头，以及”注意力不知道顺序”

RoPE：用旋转编码相对位置

注意力之外：FFN、残差、Norm

拼起来：一个能跑的 GPT

超参数怎么选：大家其实都差不多

记住这三件事

延伸阅读

目录

反向链接