Transformer 的完整结构

现在，我们已经会把文字查成向量，也会让一个位置从其他位置混合信息。离一个能训练的语言模型还差什么？

至少还差三个安排：告诉网络这些词元在哪里，让每个位置对信息做进一步加工，以及让多层计算能够稳定地传递数值与梯度。Transformer 就是把这些部件组织在一起的一类网络。

1. 先分清三种常见组织方式

最初的 Transformer 面向序列到序列任务，例如翻译，包含编码器和解码器。后来，这些部件也被组合成不同用途的结构。

这里的“解码器”是网络结构名称，不是第十七章把词元 ID 还原成文字的解码函数。

在编码器—解码器结构中，解码器内部通常有两种注意力：一层查看已经给出的目标前缀，另一层用当前查询访问编码器的键和值。后者就是交叉注意力。源输入本来就是已知材料，可以完整读取；目标未来词元仍然不能偷看。

仅解码器模型通常省去这条独立的编码器和交叉注意力通路，可以把提示和回答组织在同一条序列中。采用这种结构时，还需要明确训练目标、位置处理和归一化安排，不能只根据“解码器”这个名称推断全部细节。

本章实现一个小型的仅解码器网络，为下一章的文本训练做准备。它使用因果注意力和可学习的绝对位置向量，采用下面要解释的 Pre-LN 结构。

2. 一个块里面，信息经过哪些步骤

设当前表示为 $X$，形状是 $B\times T\times D$。我们使用两个子层：多头注意力，以及逐位置前馈网络。两者都配上残差连接：

$$\begin{aligned} U&=X+\operatorname{Attention}(\operatorname{LN}_1(X))\\ Y&=U+\operatorname{FFN}(\operatorname{LN}_2(U)) \end{aligned}$$

$\operatorname{LN}$ 表示 LayerNorm。归一化放在子层计算之前，因此这种安排叫 Pre-LN。两个 LN 各自有参数。

图 1：两个“加号”把原表示与新计算的结果逐项相加。图中各个主干位置都保持 $B\times T\times D$ 的形状。

原始 Transformer 采用的是另一种顺序，常叫 Post-LN：

$$U=\operatorname{LN}(X+\operatorname{Attention}(X))$$

两种顺序并不完全等价，对梯度和训练设置也会产生影响。关于 Transformer 中 LayerNorm 的研究分析了这些差异。本书选 Pre-LN 来搭建小模型，不把这一选择解释成所有任务上唯一正确的结构。

残差让已有表示可以直接向后传递，同时容纳新计算的修正。它不保证任意深度、初始化和学习率都能稳定工作；训练时仍要观察损失与梯度。

3. LayerNorm 与前馈网络各做什么

对当前一个位置的 $D$ 个特征 $x_1,\ldots,x_D$，LayerNorm 先计算：

$$\mu=\frac1D\sum_{j=1}^D x_j,\qquad s^2=\frac1D\sum_{j=1}^D(x_j-\mu)^2$$

再变换每个特征：

$$y_j=\gamma_j\frac{x_j-\mu}{\sqrt{s^2+\epsilon}}+\beta_j$$

$\epsilon$ 是防止分母过小的正数；$\gamma_j,\beta_j$ 可以学习。本章的 nn.LayerNorm(D) 对最后一个特征维计算，不会跨不同词元位置求均值，因此不会把未来位置混入当前表示。

接着，前馈网络对每个位置分别执行相同的非线性变换。按行向量写：

$$\operatorname{FFN}(\boldsymbol x) =\phi(\boldsymbol xW_1+\boldsymbol b_1)W_2+\boldsymbol b_2$$

这里 $W_1$ 把 $D$ 维扩展到 $4D$，$W_2$ 再映回 $D$ 维。扩展后的维度并不是增加了词元数量；序列长度 $T$ 始终不变。

本例选择 GELU 作为 $\phi$。可以把它理解成一种平滑的非线性门控：较大的正输入大多保留，负输入被抑制，却不像 ReLU 那样把所有负数都变成零。具体计算交给 nn.GELU()。

注意力在不同位置之间交换信息，前馈网络加工每个位置已经拿到的特征。两者配合，才不仅仅是在不断做加权平均。

4. 顺序从哪里进入网络

如果没有位置相关信息，完整的双向自注意力会随输入排列一起重排输出：它知道各个向量之间的关系，却没有一个独立的坐标告诉它“谁在前，谁在后”。

因果掩码本身会提供部分顺序约束，但我们通常还会显式加入位置信息。本例为第 $t$ 个位置准备一个可学习向量 $\boldsymbol p_t$：

$$X_t=E[x_t]+\boldsymbol p_t$$

两个向量同为 $D$ 维，可以逐项相加。词元嵌入说明“这个编号是什么”，位置向量说明“它出现在当前窗口的哪里”；这些含义最终仍由训练塑造，不是手工给每个维度命名。

可学习位置表简单直接，但只有预设数量的行。本章实现会拒绝超出 max_context 的输入；不能把访问越界或随意复制最后一行，当作已经具备更长上下文能力。

另一种方式不学习位置表，而是用正弦和余弦计算位置向量。对偶数宽度 $D$，第 $i$ 对分量为：

$$\begin{aligned} \operatorname{PE}(t,2i)&=\sin\left(t\,10000^{-2i/D}\right)\\ \operatorname{PE}(t,2i+1)&=\cos\left(t\,10000^{-2i/D}\right) \end{aligned}$$

$i$ 从 0 到 $D/2-1$。不同分量以不同速度变化，像几组快慢不同的刻度。位置 0 的编码会交替出现 0 和 1。

图 2：展示公式中的部分分量，宽度取 8。能计算一个更远位置的向量，不等于训练好的模型就能可靠处理那个长度。

5. RoPE：把位置变成旋转

旋转位置编码常写作 RoPE。它可以先对查询和键做位置相关的旋转，再计算它们的点积，而不必把位置向量直接加到每个输入上。

先看二维。把向量旋转角度 $\theta$，可以乘一个矩阵：

$$R(\theta)= \begin{bmatrix} \cos\theta&-\sin\theta\\ \sin\theta&\cos\theta \end{bmatrix}$$

位置 $m$ 的查询旋转 $m\theta$，位置 $n$ 的键旋转 $n\theta$，将它们的点积记作 $s_{mn}$：

$$\begin{aligned} s_{mn} &=\left(R(m\theta)\boldsymbol q\right)^\mathsf T \left(R(n\theta)\boldsymbol k\right)\\ &=\boldsymbol q^\mathsf T R(m\theta)^\mathsf T R(n\theta)\boldsymbol k\\ &=\boldsymbol q^\mathsf T R((n-m)\theta)\boldsymbol k \end{aligned}$$

这里用到了旋转的两个性质：转置相当于反向旋转，连续旋转的角度可以相加。于是，在固定查询、键向量的这一步计算里，位置通过差值 $n-m$ 进入点积。

图 3：为了看清几何关系，二维示意每步旋转 0.35 弧度。代码中的多维版本则给不同分量对使用不同频率。

更高维的 RoPE 把特征两两分组，每一对使用自己的旋转频率。它保持向量长度，又把位置信息带入匹配分数。常见用法旋转查询和键，值向量继续提供要混合的内容。

注意这条恒等式没有保证“距离越远，每一对向量的点积都单调变小”，也没有保证模型可以无限延长上下文。旋转频率、训练长度、模型参数和推理方式仍会影响实际表现。

下方脚本提供正弦位置编码和 RoPE 的数值演示，但小型网络本身使用可学习位置表。这样可以分别观察三种方法，不把它们不加说明地叠在一起。RoFormer 论文给出了 RoPE 的系统介绍。

6. 从词元编号到下一词元分数

把前面的部件连接起来，完整路径是：

词元 ID：B×T
    ↓ 词元嵌入 + 位置嵌入
输入表示：B×T×D
    ↓ 若干个 Pre-LN Transformer 块
隐藏表示：B×T×D
    ↓ 最终 LayerNorm
    ↓ 映射到词表的线性层
预测分数：B×T×V

这里的 $V$ 是词表大小。每个位置都有一行 $V$ 个分数，训练时与错开一位的目标计算交叉熵。网络返回的是 logits；交叉熵接口会处理相应的对数概率，不需要先手工做 Softmax。

代码里的多头计算先将 $B\times T\times D$ 拆成 $B\times h\times T\times(D/h)$，做完注意力再合回去。transpose 改变的是维度顺序，不是把数值乘上某个变换；contiguous 则让之后的形状调整拥有合适的内存排列。

为方便逐行阅读，这个模型没有加入 Dropout、输入输出权重共享或其他结构变体，也没有调用完整的 Transformer 封装。PyTorch 仍负责张量、基本层和自动微分；“从零实现”指亲手组织模型与训练过程，而不是重写底层矩阵库。

可以下载模型文件并运行结构演示：

python src/examples/tiny_transformer.py

"""第 20 章：可训练的微型解码器，以及位置编码的数值演示。"""

from math import sqrt
import torch
from torch import nn


class CausalAttention(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        if width % heads:
            raise ValueError("width 必须能被 heads 整除")
        self.heads = heads
        self.qkv = nn.Linear(width, 3*width, bias=False)
        self.output = nn.Linear(width, width, bias=False)

    def forward(self, x):
        batch, length, width = x.shape
        head_width = width//self.heads
        projected = self.qkv(x).chunk(3, dim=-1)
        query, key, value = [
            part.reshape(batch, length, self.heads, head_width).transpose(1, 2)
            for part in projected
        ]
        scores = query @ key.transpose(-2, -1) / sqrt(head_width)
        allowed = torch.ones(length, length, dtype=torch.bool, device=x.device).tril()
        weights = scores.masked_fill(~allowed, float("-inf")).softmax(-1)
        mixed = (weights @ value).transpose(1, 2).contiguous().reshape(batch, length, width)
        return self.output(mixed)


class Block(nn.Module):
    def __init__(self, width, heads):
        super().__init__()
        self.norm1, self.norm2 = nn.LayerNorm(width), nn.LayerNorm(width)
        self.attention = CausalAttention(width, heads)
        self.feedforward = nn.Sequential(nn.Linear(width, 4*width), nn.GELU(),
                                         nn.Linear(4*width, width))

    def forward(self, x):
        x = x+self.attention(self.norm1(x))
        return x+self.feedforward(self.norm2(x))


class TinyTransformer(nn.Module):
    def __init__(self, vocab_size, max_context=32, width=32, heads=4, layers=2):
        super().__init__()
        self.config = dict(vocab_size=vocab_size, max_context=max_context,
                           width=width, heads=heads, layers=layers)
        self.tokens = nn.Embedding(vocab_size, width)
        self.positions = nn.Embedding(max_context, width)
        self.blocks = nn.Sequential(*[Block(width, heads) for _ in range(layers)])
        self.final_norm = nn.LayerNorm(width)
        self.output = nn.Linear(width, vocab_size, bias=False)
        self.apply(self.initialize)

    @staticmethod
    def initialize(module):
        if isinstance(module, (nn.Linear, nn.Embedding)):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if isinstance(module, nn.Linear) and module.bias is not None:
                nn.init.zeros_(module.bias)

    def forward(self, ids):
        length = ids.shape[1]
        if not 1 <= length <= self.config["max_context"]:
            raise ValueError("输入长度必须在 1 与 max_context 之间")
        positions = torch.arange(length, device=ids.device)
        x = self.tokens(ids)+self.positions(positions)
        return self.output(self.final_norm(self.blocks(x)))


def sinusoidal(length, width):
    if width % 2:
        raise ValueError("本演示要求偶数维度")
    positions = torch.arange(length, dtype=torch.float64)[:, None]
    frequencies = 10000.0**(-torch.arange(0, width, 2, dtype=torch.float64)/width)
    angles = positions*frequencies
    result = torch.empty(length, width, dtype=torch.float64)
    result[:, 0::2], result[:, 1::2] = angles.sin(), angles.cos()
    return result


def rotate_pairs(vector, position, base=10000.0):
    if vector.shape[-1] % 2:
        raise ValueError("旋转维度必须是偶数")
    width = vector.shape[-1]
    frequencies = base**(-torch.arange(0, width, 2, dtype=vector.dtype,
                                      device=vector.device)/width)
    angles = position*frequencies
    even, odd = vector[..., 0::2], vector[..., 1::2]
    rotated = torch.stack([even*angles.cos()-odd*angles.sin(),
                           even*angles.sin()+odd*angles.cos()], dim=-1)
    return rotated.flatten(-2)


def main():
    torch.set_num_threads(1)
    torch.manual_seed(2026)
    model = TinyTransformer(vocab_size=20, max_context=8, width=16, heads=4, layers=2).double()
    model.eval()
    ids = torch.tensor([[1, 3, 5, 7], [1, 4, 6, 8]])
    with torch.no_grad():
        original = model(ids)
        changed_ids = ids.clone()
        changed_ids[:, 3] = 9
        changed = model(changed_ids)
    print(f"输入 ID 形状：{tuple(ids.shape)}；输出分数形状：{tuple(original.shape)}")
    print(f"可学习参数数：{sum(p.numel() for p in model.parameters())}")
    print(f"改变末项后，前三个位置的输出不变："
          f"{torch.equal(original[:, :3], changed[:, :3])}")
    print(f"正弦位置编码第 0 行：{sinusoidal(3, 4)[0].tolist()}")
    query = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float64)
    key = torch.tensor([2.0, 0.0, 1.0, -1.0], dtype=torch.float64)
    left = torch.dot(rotate_pairs(query, 2), rotate_pairs(key, 5))
    right = torch.dot(query, rotate_pairs(key, 3))
    print(f"旋转前后长度相同：{torch.allclose(query.norm(), rotate_pairs(query, 2).norm())}")
    print(f"位置 2 与 5 的点积={left:.6f}；相对位置 3 的点积={right:.6f}")
    print(f"相对位置恒等式成立：{torch.allclose(left, right, atol=1e-12)}")


if __name__ == "__main__":
    main()

输出为：

输入 ID 形状：(2, 4)；输出分数形状：(2, 4, 20)
可学习参数数：7232
改变末项后，前三个位置的输出不变：True
正弦位置编码第 0 行：[0.0, 1.0, 0.0, 1.0]
旋转前后长度相同：True
位置 2 与 5 的点积=-2.205086；相对位置 3 的点积=-2.205086
相对位置恒等式成立：True

这次检查的不只是一层注意力，而是两层完整的块、位置嵌入和输出层。改变末项之后，前面三个位置的输出仍不变，说明当前实现保留了预期的因果约束。

模型此时仍是随机初始化。参数量、形状正确以及因果性检查通过，证明的是它具有预期的计算结构，还不能证明它会写句子。

7. 自测与小实验

1. 编码器—解码器中的交叉注意力，查询、键和值通常分别来自哪里？
2. 残差相加为什么要求两个分支形状匹配？
3. 本例 LayerNorm 会用到同一句话里未来词元的均值吗？
4. $D=32$、前馈扩展倍数为 4 时，扩展后的特征数和序列长度怎样变化？
5. 有 100 行的可学习位置表，为什么不能直接读取第 101 个位置？
6. 正弦公式能计算任意整数位置，是否就能保证模型在任意长文本上工作正常？
7. 同时将 RoPE 中的位置 $m,n$ 加 10，在固定查询、键向量时，点积是否变化？
8. $B=4,T=12,V=100$ 时，本模型返回多少个分数？为什么每个位置都需要一行词表分数？
9. 将模型改成三层，再做“改变最后词元”的检查。若失败，优先查哪些操作？

下一章，我们会用这些代码训练一个真正能生成短文本的小模型，并比较它什么时候记住了有用关系，什么时候只是把局部句子接得像样。

继续阅读：Transformer 原始论文可用来对照编码器—解码器结构和正弦位置编码；实现接口可查阅 PyTorch LayerNorm 文档。阅读时留意结构差异：本章采用的是 Pre-LN、GELU 与可学习位置表。