前言

Transformer 是目前深度学习中最重要的一类架构。从 GPT、BERT 这类大语言模型，到 ViT 这类视觉模型，近几年绝大多数突破背后都有它的身影。可以说，理解 Transformer 是读懂当代 AI 的必修课。

本文假设你已经读过本站的《正向传播与反向传播：神经网络到底如何学习》，理解"网络是一串函数复合、训练时误差沿链式法则反向流动"这两条主线。我们会沿用同样的思路拆解 Transformer：信息如何向前流动，梯度如何向后流动。

本文会从"为什么需要 Attention"讲起，逐步推导公式，最后用一个从零手写的迷你 Transformer（PyTorch）把代码落地。涉及 softmax 的数学细节——梯度、数值稳定性、为什么注意力要除以 $\sqrt{d_k}$——会链接到本站 ai-math 专栏的配套文章，避免在这里打断主线。

为什么要 Transformer

RNN 的局限

在 Transformer 之前，处理序列（文本、语音等）的主流架构是 RNN（循环神经网络）。RNN 像一条传送带，按顺序逐个处理 token，把之前看到的信息压缩在一个隐藏状态里。这个设计有三个天然的瓶颈：

  • 无法并行：第 $t$ 步的计算依赖第 $t-1$ 步的隐藏状态，只能从左到右依次计算，GPU 的大规模并行能力完全用不上。
  • 长距离依赖困难：信息每传一步都要经过一次变换，距离越远衰减越明显。处理长文本时，开头的词往往"传不到"结尾。
  • 信息瓶颈：所有历史都被压缩进一个固定大小的隐藏向量，容量有限。

Transformer 的思路

Transformer 的做法截然不同：一次看完整条序列。它不再按顺序传递状态，而是让序列中的任意两个位置直接交互。

先解释一个词：token（词元，我实在不喜欢它的官方中译） 是文本被切分后的基本单位，可以是一个词、一个字符或一个子词片段。比如"我吃了一个红苹果"可以切成 6 个 token：我 / 吃 / 了 / 一个 / 红 / 苹果。

flowchart LR subgraph RNN["RNN：信息一站一站传"] direction LR r1["第 1 个"] --> r2["第 2 个"] r2 --> r3["第 3 个"] r3 --> r4["第 4 个"] end subgraph ATTN["Transformer：任意两个直接对话"] direction LR a1["第 1 个"] --- a2["第 2 个"] a1 --- a3["第 3 个"] a1 --- a4["第 4 个"] a2 --- a3 a2 --- a4 a3 --- a4 end

上图为两种架构的token之间的交互示意： 从图论的角度来说，RNN 是一个 有向路径图，包含 $n$ 个节点（代表隐状态 $h_t$）和 $n-1$ 条有向边（$h_{t-1} \to h_t$）；而 Transformer 是一个 完全图，包含 $n$ 个节点和 $\frac{n(n-1)}{2}$ 条无向边（每两个 token 之间都有一条边）。

所以"任意两个 token 之间都只有一步"指的是：在 Transformer 里，任意两个位置之间的距离恒为一步，不随序列变长而增加。这种"所有人直接互相看"的机制就是 Self-Attention（自注意力） 的核心价值：GPU 可以并行处理整个序列（不用排队等前一个人），任何两个位置也能直接交互（不受距离衰减影响）。

复杂度的代价：为什么上下文窗口是瓶颈

并行和全局交互不是免费的午餐，Transformer 的代价是 $\mathcal{O}(N^2)$ 的时空复杂度。

先看 RNN：处理长度为 $N$ 的序列，时间复杂度是 $\mathcal{O}(N \cdot d^2)$——随序列长度线性增长，而且每个时刻只需保存一个隐藏状态，内存占用极小。

再看 Transformer：它要计算相似度矩阵 $QK^\top$，事情就完全不同了。

$N \times N$ 的相似度矩阵。这里的 $N$ 代表序列长度（也就是一句话里切分出的 token 数量）。因为 Transformer 不再像 RNN 那样从左到右挨个传纸条，而是要让序列里的每一个词，都和序列里的所有词比对一次关联度，这就形成了一个 $N \times N$ 的网格方阵。方阵中第 $(i, j)$ 个格子的数值，代表第 $i$ 个词对第 $j$ 个词的"注意力权重"。下图展示了这个 $N \times N$ 的注意力矩阵——每一行代表一个词，它在各个列上的颜色深浅，就是它对其他所有词的关注度：

注意力矩阵示意图

时间代价：为什么是 $\mathcal{O}(N^2 \cdot d)$。要填满上面这个 $N \times N$ 的方阵，需要做 $N \times N$ 次独立的计算；而每填一个格子，实际上是拿 Query 的一个向量（长度 $d$）和 Key 的一个向量（长度 $d$）做内积求和，这需要经历 $d$ 次基本的乘加运算。所以总的计算基本操作次数是：

$$ N \times N \times d $$

现实直觉：如果你把丢给模型的文本长度翻倍（比如从 1000 字变成 2000 字），模型计算相似度的时间不是变长 1 倍，而是直接变慢 4 倍（$2^2$）。$N$ 每翻一倍，$N^2$ 就放大到 4 倍，这就是为什么超长文本对 Transformer 来说是灾难性的。

回到苹果的例子，用"我吃了一个红苹果"这种短句（$N = 6$），$N^2 = 36$ 个点积毫无压力，GPU 上的矩阵乘法反而快得惊人；但输入一旦变成几十万 token 的长文档，$N^2$ 就会指数爆炸，显存（$\mathcal{O}(N^2)$，用来存储注意力矩阵）也会被直接撑爆。

这解释了为什么现代大模型都受困于 上下文窗口（Context Window） 的长度，也解释了为什么研究者要发明 FlashAttention（优化注意力计算的显存），以及 Mamba、RWKV 这类试图把复杂度降回线性的非 Transformer 架构。

Attention 的核心直觉

从静态词向量到上下文向量

先退一步问：Attention 到底在解决什么问题？

在 Word2Vec、GloVe 时代，每个词只有一个固定的向量，称为静态词向量（Static Embedding）。“苹果"这个词，无论在"我吃了一个红苹果"还是"苹果发布了新手机"里，向量完全一样。这显然有问题——同一个词在不同语境下意思截然不同：前者是可食用的水果，后者是一个科技品牌。

Self-Attention 的根本目的，就是把这种孤立、静态的词向量，融合成包含全局序列信息的动态上下文向量（Contextualized Embedding）。经过 Attention 层后，“苹果"的向量会吸收周围"吃"“红"或"手机"的信息，在特征空间中发生偏移，变成一个特化于当前语境的向量。

理解了这一点，你就能明白"为什么要让任意两个位置直接交互”——因为每个 token 都需要站在全局视角，重新定义自己。

在 Self-Attention 中，每个 token 会生成三个向量：Query、Key、Value。一个常用的比喻是图书馆检索：

  • Query（查询）：你想找什么，比如"关于深度学习的书”。
  • Key（标签）：每本书上的标签，比如"机器学习"“小说”。
  • Value（内容）：书的实际内容。

检索时，你会把 Query 与每本书的 Key 做匹配，匹配度越高越相关；最终你吸收的不是"标签"本身，而是标签对应的"内容”——即对 Value 做加权求和，权重就是匹配度。用一句话概括：Attention 让模型决定"该看哪里、看多重"。

为什么需要三个线性投影

图书馆比喻回答了"Q、K、V 各是什么"，但还有一个更根本的问题：为什么不直接用原始输入向量算相似度，而要先用三个权重矩阵 $W^Q$、$W^K$、$W^V$ 做线性投影？原因有二。

语义解耦。原始词向量里混着大量信息（词义、词性、单复数、语体……）。通过投影，模型可以专门剥离出"我要找什么"（Query）和"我是什么"（Key）这两类特征。一个语法注意力头可能学到：$Q$ 被投影成"寻找能修饰我的成分"，$K$ 被投影成"我本身是个名词"。回到苹果的例子，“红"的 $Q$ 会说"我要找被修饰的名词”，“苹果"的 $K$ 会说"我是名词”，于是二者匹配，Attention 把"红"的信息送到"苹果"那里。

打破对称性。如果直接用原始向量，$XX^\top$ 是一个对称矩阵——token A 对 B 的注意力必然等于 B 对 A 的注意力。但语言是有方向性的：“修饰语关注中心词"和"中心词关注修饰语"在语言学上并不对等。引入独立的 $W^Q$、$W^K$ 后，$XW^Q$ 与 $XW^K$ 来自不同的投影，“苹果"看"红"和"红"看"苹果"的权重可以不同，注意力矩阵才真正表达了一个有向图。

“Self"与"Cross”：同源与异源

我们已经反复提到 Self-Attention，但"Self（自）“到底指什么？关键在于 $Q$、$K$、$V$ 的来源。

所谓自注意力，是指 $Q$、$K$、$V$ 三个矩阵同源——它们都由同一个输入序列 $X$ 乘以不同的权重矩阵得到：

$$ Q = XW^Q, \qquad K = XW^K, \qquad V = XW^V $$

也就是说，“我吃了一个红苹果"里的每个 token，都在这个序列内部寻找与它相关的其他 token。这是序列在"自己"与自己建立关联。

相对的，如果是机器翻译的 Decoder，它的 $Q$ 来自当前正在生成的目标语言序列，而 $K$、$V$ 来自源语言序列——这种异源的交互才被称为 Cross-Attention（交叉注意力）。理解"Self 同源、Cross 异源"这个界定，读到后面 Encoder-Decoder 架构时就不会混淆了。

缩放点积注意力

把上面的直觉写成矩阵形式，就得到原论文的核心公式：

$$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V $$

其中 $Q$、$K$、$V$ 是同一批 token 分别经过三个线性变换得到的矩阵，$Q, K \in \mathbb{R}^{n \times d_k}$，$V \in \mathbb{R}^{n \times d_v}$（$n$ 为 token 数；$d_v$ 通常与 $d_k$ 相等，后文不区分）。这个公式做了四件事：

  1. $QK^\top$ 计算相似度：$Q$ 是 $n \times d_k$，$K^\top$ 是 $d_k \times n$，乘积是 $n \times n$ 的相似度矩阵。第 $(i, j)$ 项是第 $i$ 个 token 的 Query 与第 $j$ 个 token 的 Key 的点积，数值越大表示两者越相关。
  2. 除以 $\sqrt{d_k}$ 缩放：当 $d_k$ 较大时，点积的数值会偏大（方差随维数线性增长），直接送进 softmax 容易让输出"要么接近 0，要么接近 1”，导致梯度消失。除以 $\sqrt{d_k}$ 是为了把点积的方差压回 1。为什么偏偏是 $\sqrt{d_k}$，数学推导见本站 ai-math 文章《Softmax 与缩放点积的数学》。
  3. softmax 按行归一化：对相似度矩阵的每一行做 softmax，使该行权重和为 1，得到一个"注意力分布”。
  4. 乘 $V$ 加权求和：用注意力分布对 $V$ 的行做加权求和，得到每个 token 的输出。

一个具体的例子

为了不流于抽象，我们用极小的例子手算一遍。设 $d_k = 2$，两个 token 的 $Q$、$K$、$V$ 分别为：

$$ Q = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}, \quad K = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix}, \quad V = \begin{bmatrix} 10 & 20 \\ 30 & 40 \end{bmatrix} $$

这里故意让 $Q$、$K$ 都取单位阵，方便心算（事实上并不方便哈哈哈）。

第一步：计算 $QK^\top$。 第 $(i, j)$ 项是 $q_i \cdot k_j$：$q_1 \cdot k_1 = 1$，$q_1 \cdot k_2 = 0$，$q_2 \cdot k_1 = 0$，$q_2 \cdot k_2 = 1$，所以：

$$ QK^\top = \begin{bmatrix} 1 & 0 \\ 0 & 1 \end{bmatrix} $$

第二步：除以 $\sqrt{d_k} = \sqrt{2} \approx 1.414$， 得到 $\begin{bmatrix} 0.707 & 0 \\ 0 & 0.707 \end{bmatrix}$。

第三步：按行 softmax。 第一行 $[0.707, 0]$ 的 softmax：

$$ \frac{e^{0.707}}{e^{0.707} + e^{0}} \approx \frac{2.028}{3.028} \approx 0.67, \qquad \frac{e^{0}}{e^{0.707} + e^{0}} \approx 0.33 $$

第二行同理得到约 $[0.33, 0.67]$。

第四步：乘 $V$ 加权求和。 第一行输出：

$$ 0.67 \times \begin{bmatrix} 10 \\ 20 \end{bmatrix} + 0.33 \times \begin{bmatrix} 30 \\ 40 \end{bmatrix} \approx \begin{bmatrix} 16.6 \\ 26.6 \end{bmatrix} $$

第二行输出约为 $\begin{bmatrix} 23.4 \\ 33.4 \end{bmatrix}$。

结论：每个 token 的输出都是两个 token 的 $V$ 的加权混合，权重由它们之间的相似度决定。token 1 更关注自己（权重 $0.67$），但也没有完全忽略 token 2（权重 $0.33$）——这就是"注意力”。

Multi-Head Attention

单个注意力只能让模型关注一种"关系”。但一句话里同时存在句法关系、语义关系、指代关系……让多个注意力头并行、每个头在不同子空间里学习，能捕获更多样的模式。

$$ \text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \dots, \text{head}_h) W^O $$

其中 $\text{head}_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)$。

具体做法是：把 $d_{\text{model}}$ 维的向量通过不同的投影矩阵分成 $h$ 份，每份 $d_k = d_{\text{model}} / h$ 维，各自独立做一次缩放点积注意力；最后把 $h$ 个头的结果拼回 $d_{\text{model}}$ 维，再过一层输出投影 $W^O$。

flowchart LR x["输入 X"] --> proj["线性投影 ×h"] proj --> h1["head 1：缩放点积"] proj --> h2["head 2：缩放点积"] proj --> hh["…… head h：缩放点积"] h1 --> cat["Concat 拼接"] h2 --> cat hh --> cat cat --> wo["输出投影 W^O"] wo --> out["输出"]

位置编码

Self-Attention 本身对位置不敏感：把输入 token 的顺序打乱，注意力计算的结果不变（$QK^\top$ 只是两两之间在算相似度）。但"词序"在语言里至关重要——“我打你"和"你打我"完全不是一个意思。因此必须把位置信息注入进去。

原论文采用固定（不需要训练）的 sinusoidal 位置编码，对第 $pos$ 个位置的 $d_{\text{model}}$ 维向量，第 $2i$ 维和第 $2i+1$ 维分别填入：

$$ PE_{(pos, 2i)} = \sin\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right), \qquad PE_{(pos, 2i+1)} = \cos\left(\frac{pos}{10000^{2i / d_{\text{model}}}}\right) $$

频率随维度指数变化：低维对应高频（相邻位置差异大），高维对应低频（相邻位置差异小）。这样每个位置得到一个唯一的向量，把它加到 token 的词向量上，模型就能感知位置。

它的两个好性质：一是周期函数的线性组合可以表达相对位置（模型能学到"位置 $pos + k$ 与位置 $pos$ 之间的偏移”）；二是不需要额外参数、天然可以外推到比训练时更长的序列。

现代实现（如 GPT 系列）常用可学习的位置嵌入（一个普通的 nn.Embedding），效果相近。本文代码采用 sinusoidal 版本，与上面的公式一一对应。

完整架构：Encoder-Decoder

双重混合：Attention 负责空间，FFN 负责特征

在进入 Encoder-Decoder 之前，先看清每个 Transformer 块内部两大模块的分工。一个块由两个"混合器"组成，一个沿着序列方向工作，一个沿着特征方向工作，二者正交：

  • Self-Attention —— 空间混合（Token-mixing）：它沿着序列维度 $N$ 交换信息，把不同位置的 token 组合起来。抛开 softmax 和线性投影不谈，它本质上是一个加权求和过程，负责在词与词之间搬运信息。
  • FFN —— 特征混合（Channel-mixing）：它对每个 token 独立作用，沿着特征维度 $d_{\text{model}}$ 变换特征，没有任何跨位置的交互。

这个分工可以用苹果的例子理解：Attention 决定"看哪里"——把"吃"“红"的信息搬运到"苹果"那里；但搬运过来的原始信息是粗糙的，必须交给 FFN 做深度的加工、过滤和提炼，才能变成有用的特征。

FFN 先升维（通常是 $4d_{\text{model}}$）再降维的结构，在现代可解释性研究中被看作一种键值记忆网络（Key-Value Memory）：升维层的神经元隐式存储了大量具体的世界知识与模式（“苹果可食用"“苹果是科技品牌”），降维层则负责把这些模式提取并映射回模型的表征空间。

残差连接与 LayerNorm：让网络"做深"的物理基础

另一个要在概览阶段就讲清的，是残差连接与 LayerNorm 为什么是 Transformer 能堆到上百层的根本原因。

RNN 通常只能堆叠少数几层（2-4 层），因为梯度在时间和深度两个维度上都极易消失。Transformer 则采用类似"公路网"的设计：每个 token 都有一条贯穿整个模型的主干道（残差连接），Attention 和 FFN 只是主干道旁边的"加工厂”：

$$ x_{\text{out}} = x + \text{Sublayer}(x) $$

信息和误差可以沿这条主干道无损地穿透到任意一层。这正是 GPT-3 能堆 96 层而不崩溃的原因——无论中间经过多少加工厂，“苹果"的原始特征始终有一条直接通路从输入层通到输出层。LayerNorm 则像主干道上的稳压器，强制把每个 token 的特征约束在均值 0、方差 1 的空间内，防止深层网络前向传播时数值爆炸、方向震荡。

原论文把 Transformer 设计成 Encoder-Decoder 结构，用于机器翻译这类"序列到序列"任务：

flowchart LR subgraph ENC["Encoder 堆叠 N 层"] direction TB e1["输入嵌入 + 位置编码"] --> ea["Multi-Head Self-Attention"] ea --> eadd["+ 残差"] eadd --> eln["LayerNorm"] eln --> eff["Feed-Forward"] eff --> eadd2["+ 残差"] eadd2 --> eln2["LayerNorm"] end subgraph DEC["Decoder 堆叠 N 层"] direction TB d1["输出嵌入 + 位置编码"] --> da["Masked Self-Attention"] da --> dadd["+ 残差"] dadd --> dln["LayerNorm"] dln --> dc["Cross-Attention（Q 来自 Decoder，K、V 来自 Encoder）"] dc --> dadd2["+ 残差"] dadd2 --> dln2["LayerNorm"] dln2 --> dff["Feed-Forward"] dff --> dadd3["+ 残差"] dadd3 --> dln3["LayerNorm"] end ENC --> kv["Encoder 输出 K、V"] kv --> dc dln3 --> lin["Linear"] lin --> sm["Softmax"]
  • Encoder：每个块 = 一层 Self-Attention（这里的 $Q$、$K$、$V$ 都来自输入本身）+ 一层 Feed-Forward，每层外面都有残差连接和 LayerNorm。它的任务是把输入句子编码成语义向量。
  • Decoder：每个块在三层结构上又多了一层 Cross-Attention。三者分工明确：
    • Masked Self-Attention：带因果掩码的自注意力，位置 $i$ 只能"看见"位置 $\le i$ 的 token，保证生成时不能偷看未来——这是自回归生成的关键。
    • Cross-Attention：$Q$ 来自 Decoder 自己的上一层输出，而 $K$、$V$ 来自 Encoder 的输出。它让 Decoder 在生成每个词时"查询"源语言的语义——这是"翻译"发生的核心机制。
    • Feed-Forward：对每个 token 独立做两次全连接 + 激活：
$$ \text{FFN}(x) = W_2 \cdot \text{ReLU}(W_1 x + b_1) + b_2 $$

残差连接解决深层网络的优化困难（原理见本站《为什么 ResNet 能训练 152 层》）；LayerNorm 把每个 token 的特征归一化到均值 0、方差 1，稳定训练。原论文中这些子层的顺序是"先注意力、再加残差、再归一化”（Post-LN）；现代实现更常用"先归一化、再注意力、再加残差”（Pre-LN），训练更稳定，本文代码采用后者。

从 Encoder-Decoder 到 Decoder-only

原论文是 Encoder-Decoder 结构，适合翻译、摘要这类任务。但近几年的主流大语言模型（GPT 系列）走的是更简洁的 Decoder-only 路线：去掉 Encoder 和 Cross-Attention，只保留 Masked Self-Attention + Feed-Forward 的堆叠，用"预测下一个 token"的方式训练。这样模型只需要一个统一的接口：给一段文本，接着往下写。

理解完整架构有助于看懂原论文；而代码实现聚焦 Decoder-only，是因为它更贴近现代 LLM、代码也更短，且足以覆盖本文讲到的全部核心组件。

用动画走一遍 Transformer 前向流程

架构讲完了，让我们用一个动画把整个流程串起来。从"我吃了一个红苹果"这句话出发，观察它如何被切分成 token，再依次流经嵌入、位置编码、Q/K/V 投影、自注意力、前馈与归一化，最终输出对下一个 token 的预测。点"下一步"逐步观看，或"自动播放"连贯浏览。

1 / 8

代码实战：从零手写迷你 Transformer

下面的代码不借助 torch.nn.Transformer 之类的封装，从零实现一个 Decoder-only 的字符级语言模型。任务很直观：给它一段文本作为语料，训练它学会"预测下一个字符"。

数据准备

先做字符级分词：把语料拆成不重复的字符集合，每个字符对应一个整数索引。目标序列就是输入序列整体右移一位——输入"从前有座山"，目标"前有座山"（训练时模型要预测下一个字符）。

python 
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

# ---------- 超参数 ----------
BATCH_SIZE = 64      # 每次训练的序列条数
BLOCK_SIZE = 128     # 上下文长度：每个样本最多看多少个字符
N_EMBD = 128         # 嵌入维度（d_model）
N_HEAD = 4           # 注意力头数
N_LAYER = 4          # Transformer 块数量
DROPOUT = 0.2
LEARNING_RATE = 3e-4
MAX_ITERS = 5000
EVAL_INTERVAL = 500  # 每多少步打印一次 loss

# ---------- 数据：字符级分词 ----------
# input.txt 可以是任意一段中文或英文语料
with open('input.txt', 'r', encoding='utf-8') as f:
    text = f.read()

chars = sorted(list(set(text)))                  # 语料中所有不重复字符
vocab_size = len(chars)
stoi = {ch: i for i, ch in enumerate(chars)}     # 字符 -> 索引
itos = {i: ch for i, ch in enumerate(chars)}     # 索引 -> 字符
encode = lambda s: [stoi[c] for c in s]          # 字符串 -> 索引列表
decode = lambda l: ''.join(itos[i] for i in l)   # 索引列表 -> 字符串

data = torch.tensor(encode(text), dtype=torch.long)
n = int(0.9 * len(data))
train_data, val_data = data[:n], data[n:]

def get_batch(split):
    """随机取一批 (输入, 目标) 对：目标 = 输入右移一位"""
    d = train_data if split == 'train' else val_data
    ix = torch.randint(len(d) - BLOCK_SIZE, (BATCH_SIZE,))
    x = torch.stack([d[i:i + BLOCK_SIZE] for i in ix])
    y = torch.stack([d[i + 1:i + BLOCK_SIZE + 1] for i in ix])
    return x, y

位置编码

直接实现公式里的 sinusoidal 编码，预先算好一张 $(\text{max\_len} \times d_{\text{model}})$ 的表，在 forward 里加到 token 嵌入上。

python 
# ---------- 位置编码：sinusoidal，无需训练 ----------
class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_len: int = 512):
        super().__init__()
        pe = torch.zeros(max_len, d_model)                       # (max_len, d_model)
        pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)  # (max_len, 1)
        div = torch.exp(torch.arange(0, d_model, 2).float()
                        * (-math.log(10000.0) / d_model))        # 频率项 10000^(-2i/d)
        pe[:, 0::2] = torch.sin(pos * div)   # 偶数维：正弦
        pe[:, 1::2] = torch.cos(pos * div)   # 奇数维：余弦
        pe = pe.unsqueeze(0)                 # (1, max_len, d_model)，便于广播
        self.register_buffer('pe', pe)       # 不参与训练，随模型自动移动设备

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return x + self.pe[:, :x.size(1)]    # 直接加到 token 嵌入上

缩放点积注意力

一个类对应一个公式：$\text{Attention}(Q, K, V) = \text{softmax}\left(QK^\top / \sqrt{d_k}\right)V$。mask 用于把不允许"看见"的位置填成 -inf（softmax 后权重为 0）。

python 
# ---------- 缩放点积注意力 ----------
class ScaledDotProductAttention(nn.Module):
    def __init__(self, dropout: float = DROPOUT):
        super().__init__()
        self.dropout = nn.Dropout(dropout)

    def forward(self, q, k, v, mask=None):
        d_k = q.size(-1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(d_k)  # (B,H,T,T)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))       # 掩码位置设为 -inf
        weights = F.softmax(scores, dim=-1)                             # 按行归一化
        weights = self.dropout(weights)
        out = torch.matmul(weights, v)                                  # 加权求和
        return out, weights

多头注意力

$h$ 个投影矩阵把 $d_{\text{model}}$ 维输入切成 $h$ 份（每份 $d_k = d_{\text{model}} / h$），每个头独立做注意力，最后拼回并投影：

python 
# ---------- 多头注意力 ----------
class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_head: int, dropout: float = DROPOUT):
        super().__init__()
        assert d_model % n_head == 0
        self.n_head = n_head
        self.d_k = d_model // n_head
        self.wq = nn.Linear(d_model, d_model, bias=False)
        self.wk = nn.Linear(d_model, d_model, bias=False)
        self.wv = nn.Linear(d_model, d_model, bias=False)
        self.wo = nn.Linear(d_model, d_model)   # 输出投影
        self.attn = ScaledDotProductAttention(dropout)

    def forward(self, x, mask=None):
        B, T, C = x.size()
        # 投影后切分成 h 个头：(B, T, d_model) -> (B, h, T, d_k)
        q = self.wq(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)
        k = self.wk(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)
        v = self.wv(x).view(B, T, self.n_head, self.d_k).transpose(1, 2)
        ctx, _ = self.attn(q, k, v, mask)                  # (B, h, T, d_k)
        ctx = ctx.transpose(1, 2).contiguous().view(B, T, C)  # 拼回头维度
        return self.wo(ctx)

前馈网络与 Transformer 块

FFN 对每个 token 独立做"放大 → 激活 → 缩回"。TransformerBlock 采用 Pre-LN 顺序（先归一化再做注意力/前馈，再加残差）：

python 
# ---------- 前馈网络 ----------
class FeedForward(nn.Module):
    def __init__(self, d_model: int, dropout: float = DROPOUT):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(d_model, 4 * d_model),   # 先放大到 4 倍
            nn.GELU(),
            nn.Linear(4 * d_model, d_model),   # 再缩回原维度
            nn.Dropout(dropout),
        )

    def forward(self, x):
        return self.net(x)

# ---------- 单个 Transformer 块 ----------
class TransformerBlock(nn.Module):
    def __init__(self, d_model: int, n_head: int, dropout: float = DROPOUT):
        super().__init__()
        self.ln1 = nn.LayerNorm(d_model)                 # Pre-LN：先归一化
        self.attn = MultiHeadAttention(d_model, n_head, dropout)
        self.ln2 = nn.LayerNorm(d_model)
        self.ffn = FeedForward(d_model, dropout)

    def forward(self, x, mask):
        x = x + self.attn(self.ln1(x), mask)   # 子层 1 + 残差
        x = x + self.ffn(self.ln2(x))          # 子层 2 + 残差
        return x

# ---------- 因果掩码：位置 i 只能看到位置 <= i ----------
def build_causal_mask(B: int, n_head: int, T: int, device):
    mask = torch.tril(torch.ones(T, T, device=device))   # 下三角为 1，上三角为 0
    return mask.view(1, 1, T, T).expand(B, n_head, T, T)

完整模型与采样

把嵌入、位置编码、$N$ 层块、输出层串起来。generate 做自回归：每次只看最后一个位置的输出分布，采样一个 token 拼回输入，再继续。

python 
# ---------- 完整模型 ----------
class MiniTransformer(nn.Module):
    def __init__(self, vocab_size: int, d_model: int, n_head: int,
                 n_layer: int, block_size: int, dropout: float = DROPOUT):
        super().__init__()
        self.token_embedding = nn.Embedding(vocab_size, d_model)
        self.pos_encoding = PositionalEncoding(d_model, block_size)
        self.blocks = nn.ModuleList([
            TransformerBlock(d_model, n_head, dropout) for _ in range(n_layer)
        ])
        self.ln_f = nn.LayerNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size)
        self.block_size = block_size

    def forward(self, idx, targets=None):
        B, T = idx.size()
        x = self.pos_encoding(self.token_embedding(idx))   # (B, T, d_model)
        mask = build_causal_mask(B, self.blocks[0].attn.n_head, T, idx.device)
        for block in self.blocks:
            x = block(x, mask)
        logits = self.lm_head(self.ln_f(x))                # (B, T, vocab_size)
        loss = None
        if targets is not None:
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)),
                                   targets.view(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, max_new_tokens: int):
        """自回归采样：每次预测下一个 token，拼回输入继续"""
        for _ in range(max_new_tokens):
            idx_cond = idx[:, -self.block_size:]           # 只保留最后 block_size 个
            logits, _ = self.forward(idx_cond)
            probs = F.softmax(logits[:, -1, :], dim=-1)    # 只看最后一个位置的分布
            idx_next = torch.multinomial(probs, num_samples=1)  # 按概率采样
            idx = torch.cat((idx, idx_next), dim=1)
        return idx

训练与生成

训练循环只做三件事：前向算出 loss → backward() 反向传播 → 优化器更新参数。这正是本站《正向传播与反向传播》讲的那条主线。

python 
# ---------- 训练 ----------
model = MiniTransformer(vocab_size, N_EMBD, N_HEAD, N_LAYER, BLOCK_SIZE, DROPOUT)
optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE)

for step in range(MAX_ITERS):
    x, y = get_batch('train')
    _, loss = model(x, y)
    optimizer.zero_grad()
    loss.backward()          # 反向传播：误差沿链式法则流向每个参数
    optimizer.step()         # 参数更新
    if step % EVAL_INTERVAL == 0:
        print(f"step {step}, loss {loss.item():.4f}")

# ---------- 生成 ----------
context = torch.tensor([encode("从前有座山")], dtype=torch.long)  # 起始提示
print(decode(model.generate(context, max_new_tokens=200)[0].tolist()))

把上面的代码按顺序拼在一起，配一个 input.txt 语料，就可以在 CPU 上训练一个能"写句子"的迷你 Transformer。你可以把 N_EMBD、N_LAYER 调大、语料换得更长更有规律，生成的文本会明显更有模有样——模型的容量和语料决定了它的上限。

总结

回到开头那条主线：信息如何向前流动，梯度如何向后流动。 在 Transformer 里，输入 token 经过嵌入和位置编码，逐层穿过多头注意力、前馈网络和归一化层，最终得到预测；训练时，损失沿链式法则反向穿过每一层，告诉每个参数该往哪个方向调整。这与本站《正向传播与反向传播》讲的机制完全一致，只是每一层的具体结构换成了 Attention。

Transformer 的三个支柱：

  1. Self-Attention：让任意两个位置直接交互，并行处理序列、不受距离限制。
  2. Multi-Head：多组并行的注意力在各自子空间学习不同的关系模式。
  3. 位置编码：把顺序信息注入模型，弥补 Self-Attention 对位置不敏感的天性。

配合残差连接（原理见《为什么 ResNet 能训练 152 层》）和 LayerNorm，深层 Transformer 才能稳定训练。而 softmax 的数学细节——梯度、数值稳定性、为何除以 $\sqrt{d_k}$——请继续阅读本站 ai-math 文章《Softmax 与缩放点积的数学》。

参考