Skip to content

Attention Is All 鼬 Need


Attention Is All You Need — Vaswani et al., 2017

The dominant sequence transduction models are based on complex recurrent or convolutional neural networks that include an encoder and a decoder. The best performing models also connect the encoder and decoder through an attention mechanism. We propose a new simple network architecture, the Transformer, based solely on attention mechanisms, dispensing with recurrence and convolutions entirely.

目前主流的序列转导模型基于包含编码器和解码器的复杂循环或卷积神经网络。表现最佳的模型还会通过注意力机制连接编码器与解码器。我们提出了一种全新的简单网络架构——Transformer,该架构完全基于注意力机制,彻底摒弃了循环结构和卷积运算。

疾旋鼬发现自己注意力完全不在学习上(!)天天想着桃旋鼬桃旋鼬……可是学习没有注意力怎么行……

伤心的疾旋鼬


从经典架构到 Transformer

Transformer 并非凭空出现。它站在巨人的肩膀上,解决了前人架构中一系列根深蒂固的瓶颈。让我们沿着时间线,逐一审视那些为 Transformer 铺路的关键工作。

前馈神经网络(FNN)

前馈神经网络是最基础的神经网络架构。信息从输入层出发,经过若干隐藏层,最终到达输出层——单向流动,没有环路

graph LR
    X["输入层<br/>x ∈ ℝᵈ"] --> H1["隐藏层 1<br/>σ(W₁x + b₁)"]
    H1 --> H2["隐藏层 2<br/>σ(W₂h₁ + b₂)"]
    H2 --> Y["输出层<br/>ŷ = softmax(W₃h₂ + b₃)"]

每一层所做的计算本质上是一个仿射变换 followed by a 非线性激活函数

\[ \mathbf{h} = \sigma(\mathbf{W}\mathbf{x} + \mathbf{b}) \]

其中 \(\mathbf{W}\) 是权重矩阵,\(\mathbf{b}\) 是偏置向量,\(\sigma\) 是激活函数(如 ReLU、Sigmoid)。

激活函数

ReLU $$ \text{ReLU}(x) = \max(0, x) $$ 最常用的隐藏层激活函数。它将所有负输入置零,正输入则原样输出。其计算高效且能有效缓解梯度消失问题,但可能导致“神经元死亡”。

Sigmoid $$ \sigma(x) = \frac{1}{1 + e^{-x}} $$ 传统激活函数,能将任意实数输入“挤压”到 (0, 1) 区间。因其输出可视为概率,常用于二分类问题的输出层。但其两端饱和区梯度接近于零,易导致梯度消失。

Softmax 对于一个包含 \(K\) 个类的输出向量 \(\mathbf{z} = [z_1, z_2, ..., z_K]\),第 \(i\) 类的输出为: $$ \text{Softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{K} e^{z_j}} $$ 专用于多分类网络的输出层。它将所有输出转换为一个概率分布(各值在0到1之间,且总和为1),从而表示输入样本属于各个类别的概率。

前馈网络的局限性:

  • 无记忆能力:每个输入独立处理,无法捕捉序列中的时序依赖关系。给它一个句子,它不知道 "I" 后面跟 "love" 还是 "hate"。
  • 固定输入长度:输入维度在设计时就已确定,处理变长序列需要额外的 padding 或截断策略。
  • 参数量爆炸:处理长序列时,全连接层的参数量随输入维度线性增长。

这些问题催生了一个关键的洞察:我们需要一种能处理序列、并且理解序列中元素顺序关系的架构。

循环神经网络(RNN)

为了解决前馈网络无法处理序列的问题,循环神经网络引入了一个革命性的概念:隐状态(hidden state)的递归传递

RNN 的核心思想是:在处理序列中的每个元素时,不仅考虑当前输入,还考虑之前所有元素累积的信息

graph LR
    X0["x₀"] --> H0["h₀"]
    H0 --> H1["h₁"]
    X1["x₁"] --> H1
    H1 --> H2["h₂"]
    X2["x₂"] --> H2
    H2 --> H3["h₃"]
    X3["x₃"] --> H3

在每个时间步 \(t\),RNN 执行:

\[ \mathbf{h}_t = \tanh(\mathbf{W}_{hh}\mathbf{h}_{t-1} + \mathbf{W}_{xh}\mathbf{x}_t + \mathbf{b}) \]

其中 \(\mathbf{h}_{t-1}\) 是上一时间步的隐状态,\(\mathbf{x}_t\) 是当前输入。这个递归公式使得信息可以在序列中流动。

RNN 解决了什么:

  • 变长序列处理:无需固定输入长度。
  • 时序建模:隐状态 \(\mathbf{h}_t\) 理论上编码了从 \(\mathbf{x}_0\)\(\mathbf{x}_t\) 的全部历史信息。

RNN 引入了什么新问题:

  • 梯度消失与梯度爆炸:在反向传播时,梯度需要沿时间步逐级传递。对于长度为 \(T\) 的序列,梯度需要连乘 \(T\)\(\mathbf{W}_{hh}\) 矩阵。当 \(\|\mathbf{W}_{hh}\|\) 的特征值小于 1 时,梯度指数级衰减;大于 1 时,梯度指数级增长。这使得 RNN 难以学习长距离依赖——句首的信息很难影响句尾的决策。
  • 串行计算\(\mathbf{h}_t\) 的计算依赖于 \(\mathbf{h}_{t-1}\),因此无法并行化。处理一个长度为 100 的序列,必须依次执行 100 步,无法利用 GPU 的大规模并行能力。

反向传播

反向传播是训练神经网络的核心算法,它通过链式法则计算损失函数对网络各参数的梯度,从而利用梯度下降法更新参数。在循环神经网络(RNN)中,反向传播需要沿时间序列展开,称为反向传播通过时间(Backpropagation Through Time, BPTT)

设神经网络的前向传播过程为:输入 \(\mathbf{x}\) 经过多层变换(如线性变换加激活函数)得到输出 \(\hat{\mathbf{y}}\),与真实标签 \(\mathbf{y}\) 比较后计算损失 \(L\)。反向传播的目标是计算损失对每个参数 \(\theta\) 的梯度 \(\frac{\partial L}{\partial \theta}\),用于参数更新:\(\theta \leftarrow \theta - \eta \frac{\partial L}{\partial \theta}\),其中 \(\eta\) 为学习率。

反向传播的核心是链式法则。例如,对于复合函数 \(L = f(g(h(x)))\),有:

\[ \frac{\partial L}{\partial x} = \frac{\partial L}{\partial f} \cdot \frac{\partial f}{\partial g} \cdot \frac{\partial g}{\partial h} \cdot \frac{\partial h}{\partial x} \]

在神经网络中,每一层对应一个函数,梯度从输出层逐层反向传递至输入层。

RNN 处理序列数据 \(\mathbf{x}_1, \mathbf{x}_2, \dots, \mathbf{x}_T\),其隐藏状态 \(\mathbf{h}_t\) 和输出 \(\mathbf{o}_t\) 的计算为:

\[ \mathbf{h}_t = \sigma(\mathbf{W}_{hh} \mathbf{h}_{t-1} + \mathbf{W}_{xh} \mathbf{x}_t + \mathbf{b}_h), \quad \mathbf{o}_t = \mathbf{W}_{hy} \mathbf{h}_t + \mathbf{b}_y \]

其中 \(\sigma\) 是激活函数(如 tanh),\(\mathbf{W}_{hh}, \mathbf{W}_{xh}, \mathbf{W}_{hy}\) 为权重矩阵,\(\mathbf{b}_h, \mathbf{b}_y\) 为偏置。

损失函数通常为各时间步损失之和:\(L = \sum_{t=1}^T L_t(\mathbf{o}_t, \mathbf{y}_t)\)

在 BPTT 中,梯度计算需沿时间反向传播。以损失 \(L\)\(\mathbf{W}_{hh}\) 的梯度为例:

\[ \frac{\partial L}{\partial \mathbf{W}_{hh}} = \sum_{t=1}^T \sum_{k=1}^t \frac{\partial L_t}{\partial \mathbf{h}_t} \cdot \frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_k} \cdot \frac{\partial \mathbf{h}_k}{\partial \mathbf{W}_{hh}} \]

其中 \(\frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_k}\) 表示第 \(t\) 步隐藏状态对第 \(k\) 步隐藏状态的依赖,可通过链式展开:

\[ \frac{\partial \mathbf{h}_t}{\partial \mathbf{h}_k} = \prod_{j=k+1}^t \frac{\partial \mathbf{h}_j}{\partial \mathbf{h}_{j-1}} = \prod_{j=k+1}^t \mathbf{W}_{hh}^\top \cdot \text{diag}(\sigma'(\mathbf{h}_{j-1})) \]

这里 \(\frac{\partial \mathbf{h}_j}{\partial \mathbf{h}_{j-1}}\) 是 Jacobian 矩阵,包含权重 \(\mathbf{W}_{hh}\) 和激活函数导数 \(\sigma'\)

LSTM(Hochreiter & Schmidhuber, 1997)和 GRU(Cho et al., 2014)通过引入门控机制部分缓解了梯度消失问题,但并未解决串行计算的根本瓶颈。

卷积神经网络(CNN)用于序列处理

CNN 最初为计算机视觉设计,但研究者很快发现,一维卷积同样可以用于序列建模。与 RNN 不同,CNN 通过局部感受野并行地从序列中提取特征。

graph LR
    subgraph "一维卷积:kernel size = 3"
        X["... x₁ x₂ x₃ x₄ x₅ ..."] --> C1["卷积核₁ → 特征 c₁"]
        X --> C2["卷积核₂ → 特征 c₂"]
    end

对于序列 \(\mathbf{x} = (x_1, x_2, \dots, x_n)\),一维卷积在每个位置 \(i\) 处计算:

\[ c_i = \sigma\left(\mathbf{w}^\top \mathbf{x}_{i:i+k-1} + b\right) \]

其中 \(k\) 是卷积核大小,\(\mathbf{x}_{i:i+k-1}\) 是以位置 \(i\) 为中心的局部窗口。

CNN 的优势:

  • 并行计算:所有位置的卷积可以同时执行,充分利用 GPU。
  • 局部特征提取:通过堆叠多层卷积,可以逐步扩大感受野,捕捉更大范围的上下文。

CNN 的局限:

  • 长距离依赖需要深层堆叠:单层卷积只能看到局部窗口(如 3 个词),要捕捉距离为 \(d\) 的依赖关系,至少需要 \(O(d/k)\) 层卷积。这意味着路径长度随距离线性增长
  • 感受野是固定的:无论两个词是否真正相关,卷积核的感受野大小不变。它缺乏动态选择关注位置的能力。

这引出了一个核心问题:能否有一种机制,让模型在每一步都能直接"看到"序列中的任意位置,并且根据内容动态决定关注哪里?

注意力机制

注意力机制(Attention Mechanism)最初由 Bahdanau et al. (2014) 在机器翻译任务中提出,用于解决 Seq2Seq 模型中信息瓶颈的问题。

在经典的 Encoder-Decoder 架构中,编码器将整个输入序列压缩为一个固定长度的上下文向量 \(\mathbf{c}\),解码器基于这个向量逐步生成输出。当输入序列很长时,这个单一向量无法承载全部信息。

graph TB
    subgraph "传统 Seq2Seq"
        E["编码器"] -->|"一个向量 c"| D["解码器"]
    end
    subgraph "带注意力的 Seq2Seq"
        E2["编码器 h₁...hₙ"] -->|"加权组合"| D2["解码器"]
        D2 -->|"查询"| A["注意力层"]
        A -->|"选择性地关注 h₁...hₙ"| E2
    end

Bahdanau 注意力的关键思想是:解码器在生成每个词时,动态地计算对编码器各隐状态的"关注程度",形成一个加权上下文向量。

具体地,在解码器时间步 \(t\)

  • 计算注意力分数(使用前馈网络):
\[ e_{t,i} = \mathbf{v}^\top \tanh(\mathbf{W}_s \mathbf{s}_{t-1} + \mathbf{W}_h \mathbf{h}_i) \]
  • 归一化得到注意力权重
\[ \alpha_{t,i} = \frac{\exp(e_{t,i})}{\sum_{j=1}^{n} \exp(e_{t,j})} \]
  • 计算上下文向量
\[ \mathbf{c}_t = \sum_{i=1}^{n} \alpha_{t,i} \mathbf{h}_i \]

注意力机制的本质: 它是一个软寻址(soft addressing)机制。与传统的哈希表硬寻址不同,注意力机制根据查询(query)与所有键(key)的相似度,对值(value)进行加权求和。相似度高的位置获得更大的权重。

Luong et al. (2015) 进一步简化了注意力计算,提出了点积注意力(dot-product attention):直接用向量点积代替前馈网络来计算相似度,计算效率更高。这正是 Transformer 所采用的形式。

残差连接

He et al. (2015) 在 ResNet 中提出了残差连接(Residual Connection),解决了深层网络的退化问题(degradation problem):当网络层数增加时,训练误差反而上升——这不是过拟合,而是优化困难。

残差连接的核心思想极其简洁:让网络学习残差映射而非直接映射。

\[ \mathbf{y} = \mathcal{F}(\mathbf{x}) + \mathbf{x} \]
graph LR
    X["输入 x"] --> F["ℱ(x)"]
    X -->|"恒等映射"| ADD["+"]
    F --> ADD
    ADD --> Y["输出 y = ℱ(x) + x"]

为什么残差连接有效?

  • 梯度高速公路:在反向传播时,梯度可以通过恒等映射直接流过,不经过权重层。这使得梯度信号可以无衰减地传播到网络的任意深度。
  • 学习更容易:如果某一层不需要变换,网络只需将 \(\mathcal{F}(\mathbf{x})\) 学习为零向量即可——这比学习一个恒等映射 \(\mathbf{y} = \mathbf{x}\) 容易得多。

残差连接使得训练数百甚至上千层的网络成为可能,成为后续所有深度学习架构(包括 Transformer)的标配组件。

回顾与动机:Transformer?

架构 并行性 长距离依赖 路径长度
前馈网络 ✅ 完全并行 ❌ 无序列建模
RNN ❌ 串行 ⚠️ 梯度消失 \(O(n)\)
CNN ✅ 并行 ⚠️ 需要深层堆叠 \(O(\log_k n)\)
Transformer ✅ 并行 ✅ 直接连接 \(O(1)\)

Transformer 的核心洞察是:如果我们用注意力机制完全替代循环和卷积,就能同时获得并行计算和任意距离的直接连接。

Transformer 架构详解

Transformer 是一个基于编码器-解码器(Encoder-Decoder)结构的模型,但其内部组件与之前的 Seq2Seq 模型有本质区别。

graph TB
    subgraph "编码器 Encoder(×N)"
        E_IN["输入嵌入 + 位置编码"] --> MHA_E["多头自注意力<br/>Multi-Head Self-Attention"]
        MHA_E --> ADD1["Add & Norm"]
        ADD1 --> FFN_E["前馈网络<br/>Feed-Forward Network"]
        FFN_E --> ADD2["Add & Norm"]
    end
    subgraph "解码器 Decoder(×N)"
        D_IN["输出嵌入 + 位置编码"] --> MASK_MHA["掩码多头自注意力<br/>Masked Multi-Head Self-Attention"]
        MASK_MHA --> ADD3["Add & Norm"]
        ADD3 --> CROSS_MHA["多头交叉注意力<br/>Multi-Head Cross-Attention"]
        CROSS_MHA --> ADD4["Add & Norm"]
        ADD4 --> FFN_D["前馈网络<br/>Feed-Forward Network"]
        FFN_D --> ADD5["Add & Norm"]
    end
    ADD2 -->|"K, V"| CROSS_MHA
    ADD5 --> LINEAR["线性层 + Softmax"]
    LINEAR --> OUT["输出概率"]

输入表示:嵌入与位置编码

Token Embedding

将离散的 token(词、子词)映射为连续的稠密向量。给定词汇表大小 \(V\) 和嵌入维度 \(d_{model}\)(所有 token 向量的统一长度),嵌入矩阵 \(\mathbf{W}_E \in \mathbb{R}^{V \times d_{model}}\)(可学习的参数表) 将 token ID \(i\) 映射为向量 \(\mathbf{e}_i \in \mathbb{R}^{d_{model}}\)

“疾旋鼬”和“工管”的向量距离,会比“疾旋鼬”和“桃旋鼬”的向量距离更近(疾旋鼬学习《无情道》中)

Positional Encoding

由于 Transformer 没有循环或卷积结构,它无法感知 token 的顺序

"疾旋鼬吃捣蛋猫"和"捣蛋猫吃疾旋鼬"在没有位置信息时是等价的(与FNN面临的问题一致)

原始论文使用正弦位置编码

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

其中 \(pos\) 是序列中的位置,\(i\) 是维度索引。

为什么选择正弦编码?

  • 每个位置的编码都是唯一的。
  • 任意两个位置的相对距离可以通过线性变换表示:\(PE_{pos+k}\) 可以表示为 \(PE_{pos}\) 的线性函数。
  • 可以外推到训练时未见过的序列长度。

相对位置的可线性表示

\(\omega_i = \frac{1}{10000^{2i/d_{\text{model}}}}\),则公式可写为: $$ PE_{(pos, 2i)} = \sin(\omega_i \cdot pos) $$ $$ PE_{(pos, 2i+1)} = \cos(\omega_i \cdot pos) $$

考虑同一个频率 \(\omega_i\) 对应的一个“维度对” \([PE_{(pos, 2i)}, PE_{(pos, 2i+1)}]^T = [\sin(\omega_i pos), \cos(\omega_i pos)]^T\)

根据三角函数的和角公式: $$ \sin(\omega_i (pos + k)) = \sin(\omega_i pos)\cos(\omega_i k) + \cos(\omega_i pos)\sin(\omega_i k) $$

\[ \cos(\omega_i (pos + k)) = \cos(\omega_i pos)\cos(\omega_i k) - \sin(\omega_i pos)\sin(\omega_i k) \]

这可以写成矩阵乘法形式: $$ \begin{bmatrix} PE_{(pos+k, 2i)}, \ PE_{(pos+k, 2i+1)} \end{bmatrix} = \begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k), \ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix} \cdot \begin{bmatrix} PE_{(pos, 2i)}, \ PE_{(pos, 2i+1)} \end{bmatrix} $$

关键洞察: - 变换矩阵 \(\mathbf{M}_i(k) = \begin{bmatrix} \cos(\omega_i k) & \sin(\omega_i k) \\ -\sin(\omega_i k) & \cos(\omega_i k) \end{bmatrix}\) 是一个旋转矩阵,它只依赖于相对距离 \(k\) 和频率 \(\omega_i\),与绝对位置 \(pos\) 无关。 - 这意味着,对于整个 \(d_{\text{model}}\) 维的位置编码向量 \(\mathbf{PE}_{pos}\),存在一个由多个 \(\mathbf{M}_i(k)\) 组成的块对角矩阵 \(\mathbf{M}(k)\),使得 \(\mathbf{PE}_{pos+k} = \mathbf{M}(k) \cdot \mathbf{PE}_{pos}\)

对模型的意义: 在自注意力机制中,模型要计算 token 之间的关系。这种线性性质意味着,模型无需记忆每一个绝对位置组合,只需学会根据相对距离 \(k\) 来调整注意力计算(即学习如何利用 \(\mathbf{M}(k)\) 或其蕴含的关系),就能自然地泛化到训练时未出现过的绝对位置关系上。这极大地提升了模型对序列结构的理解和泛化能力。

最终输入为:\(\mathbf{x}_{pos} = \mathbf{e}_{pos} + PE_{pos}\)

核心机制:缩放点积注意力

缩放点积注意力(Scaled Dot-Product Attention)是 Transformer 的心脏。给定查询(Query)键(Key)值(Value)三组向量,注意力机制通过以下步骤工作:

第一步:计算注意力分数 查询 \(\mathbf{Q}\) 与所有键 \(\mathbf{K}\) 做点积,衡量它们的相似度:

\[ \text{scores} = \mathbf{Q}\mathbf{K}^\top \]

关于Q 和 K

在缩放点积注意力机制中,Q 和 K 代表所有查询向量和键向量构成的矩阵。整个计算过程是并行化的矩阵运算。

假设输入序列包含 \( n \) 个 token,每个 token 的表示维度为 \( d_{model} \)。通过线性投影,我们得到三个矩阵:

  • 查询矩阵 (Query Matrix) \( Q \in \mathbb{R}^{n \times d_k} \)
  • 键矩阵 (Key Matrix) \( K \in \mathbb{R}^{n \times d_k} \)
  • 值矩阵 (Value Matrix) \( V \in \mathbb{R}^{n \times d_v} \)

其中,\( d_k \)\( d_v \) 是投影后的维度。矩阵的每一行对应序列中一个位置的向量

注意力分数通过矩阵乘法一次性计算得出: $$ S = Q K^\top \in \mathbb{R}^{n \times n} $$ 矩阵 \( S \) 中的每个元素 \( S_{ij} \) 是查询向量 \( \mathbf{q}_i \) 与键向量 \( \mathbf{k}_j \) 的点积:

\[ S_{ij} = \mathbf{q}_i^\top \mathbf{k}_j = \sum_{t=1}^{d_k} q_{i,t} \cdot k_{j,t} \]

因此,\( S \) 是一个 \( n \times n \) 的方阵,它完整地编码了序列中每个位置与所有位置(包括自身)的原始相关性强度。

第二步:缩放(Scaling) 除以 \(\sqrt{d_k}\)\(d_k\) 是键向量的维度):

\[ \text{scaledscores} = \frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}} \]

第三步:Softmax 归一化 将分数转换为概率分布:

\[ \alpha_{ij} = \frac{\exp(\text{score}_{ij})}{\sum_{k} \exp(\text{score}_{ik})} \]

为什么要缩放?

缩放因子 \( \frac{1}{\sqrt{d_k}} \) 被应用于分数矩阵 \( S \) 中的每一个元素: $$ S_{\text{scaled}} = \frac{S}{\sqrt{d_k}} = \frac{Q K^\top}{\sqrt{d_k}} $$ 即,对于所有 \( i, j \): $$ S_{\text{scaled}, ij} = \frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_k}} $$ 此操作的目的是控制点积值的方差。假设 \( \mathbf{q}_i \)\( \mathbf{k}_j \) 的各分量独立且方差为1,则 \( \text{Var}(\mathbf{q}_i^\top \mathbf{k}_j) = d_k \)

方差推导:由于 \(q_i\)\(k_i\) 独立,且 \(E[q_i k_i] = E[q_i]E[k_i] = 0\),则每个乘积项 \(q_i k_i\) 的方差为: $$ \text{Var}(q_i k_i) = E[q_i^2 k_i^2] - (E[q_i k_i])^2 = E[q_i^2] E[k_i^2] = \text{Var}(q_i) \text{Var}(k_i) = 1. $$ 点积是 \(d_k\) 个独立随机变量 \(q_i k_i\) 的和,因此其方差为: $$ \text{Var}(\mathbf{q}^\top \mathbf{k}) = \sum_{i=1}^{d_k} \text{Var}(q_i k_i) = d_k. $$

缩放后: $$ \text{Var}\left( S_{\text{scaled}, ij} \right) = \text{Var}\left( \frac{\mathbf{q}_i^\top \mathbf{k}_j}{\sqrt{d_k}} \right) = \frac{1}{d_k} \cdot \text{Var}(\mathbf{q}_i^\top \mathbf{k}_j) = \frac{1}{d_k} \cdot d_k = 1 $$ 这确保了所有 \( n \times n \) 个注意力分数具有稳定的数值范围,防止后续 Softmax 进入梯度饱和区。

注意力权重通过Softmax函数计算: $$ \text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^\top}{\sqrt{d_k}}\right)V $$ 其中,未缩放的分数矩阵 \(S = QK^\top\) 的每个元素 \(s_{ij} = \mathbf{q}_i^\top \mathbf{k}_j\)

解决的问题:Softmax函数对输入值的绝对大小非常敏感。考虑一个向量 \(\mathbf{z} = [z_1, z_2, ..., z_n]\),其Softmax输出为: $$ \text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{n} e^{z_j}} $$ 当某个 \(z_i\) 的绝对值极大时,\(e^{z_i}\) 会变得极其巨大(或极其接近0)。这会导致 \(\text{softmax}(z_i)\) 的输出无限逼近于 1(如果 \(z_i\) 是正极大值)或 0(如果 \(z_i\) 是负极大值)。

在注意力机制中,如果点积 \(s_{ij}\) 的值普遍很大(由于方差 \(d_k\) 大),那么经过Softmax后,大部分注意力权重会变得非常接近0或1,分布变得极其尖锐(Peaky)。

Softmax函数在输入值极大或极小时,会进入梯度饱和区。 - 对于 \(\text{softmax}(z_i)\),其关于输入 \(z_j\) 的梯度为: \(\frac{\partial \text{softmax}(z_i)}{\partial z_j} = \text{softmax}(z_i)(\delta_{ij} - \text{softmax}(z_j))\) 其中 \(\delta_{ij}\) 是克罗内克函数。 - 当 \(\text{softmax}(z_i)\) 接近0或1时,梯度 \(\text{softmax}(z_i)(1 - \text{softmax}(z_i))\)趋近于0

  • 梯度消失:极小的梯度意味着在反向传播时,用于更新模型参数(即生成 \(\mathbf{q}\)\(\mathbf{k}\) 的权重)的误差信号非常微弱,导致学习速度极其缓慢甚至停滞。
  • 训练不稳定:微小的参数更新可能会引起点积值的剧烈波动,进而导致注意力权重分布在“饱和”与“非饱和”状态间跳跃,使训练过程难以收敛。

第四步:加权求和 用注意力权重对值 \(\mathbf{V}\) 进行加权组合:

\[ \boxed{\text{Attention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}}\right)\mathbf{V}} \]

直觉理解: 注意力机制可以类比为一个图书馆检索系统

  • Query 是你的搜索问题(什么是疾旋鼬?)
  • Key 是每本书的索引标签("游戏"、"幻兽帕鲁"、"龙属性帕鲁"……)
  • Value 是每本书的实际内容(它会蜷缩身体,并以迅疾的速度旋转移动。过去人们会在被驯服的疾旋鼬身上装上放牛奶的袋子再放牧,以此利用它的旋转能力来制作黄油。)
  • 点积衡量 Query 与每个 Key 的匹配程度
  • Softmax 将匹配度转换为概率
  • 最终输出是所有书内容的加权混合——与你问题最相关的书贡献最大

"工管疾旋鼬"和"头目Boss疾旋鼬"的"疾旋鼬"Token与对上下文相同词的关心是不同的

多头注意力

单个注意力头只能学习一种"关注模式"。但在自然语言中,一个词可能同时需要关注多种不同类型的信息:

  • 句法关系:"The 疾旋鼬 is doing ___" → 关注 "疾旋鼬"(主动一致)
  • 语义关系:"It is cute." → 关注前面提到的 "疾旋鼬"
  • 位置关系:关注相邻的词,比如旁边出现的那只疾旋鼬

发现自己被火控雷达照的疾旋鼬

多头注意力(Multi-Head Attention)通过并行运行多个独立的注意力头来捕获不同类型的关系。

graph TB
    subgraph "Multi-Head Attention"
        Q["Q"] --> S1["Head 1<br/>Attention(QW₁ᴷ, KW₁ᴷ, VW₁ⱽ)"]
        Q --> S2["Head 2<br/>Attention(QW₂ᴷ, KW₂ᴷ, VW₂ⱽ)"]
        Q --> S3["Head h<br/>Attention(QWₕᴷ, KWₕᴷ, VWₕⱽ)"]
        K["K"] --> S1
        K --> S2
        K --> S3
        V["V"] --> S1
        V --> S2
        V --> S3
        S1 --> CONCAT["Concat"]
        S2 --> CONCAT
        S3 --> CONCAT
        CONCAT --> PROJ["Linear(Wᴼ)"]
        PROJ --> OUT["输出"]
    end

具体实现:

  • \(\mathbf{Q}, \mathbf{K}, \mathbf{V}\) 分别通过 \(h\) 组不同的线性投影,投影到低维空间(维度为 \(d_k = d_v = d_{model}/h\)):
\[ \text{head}_i = \text{Attention}(\mathbf{Q}\mathbf{W}_i^Q, \mathbf{K}\mathbf{W}_i^K, \mathbf{V}\mathbf{W}_i^V) \]

其中 \(\mathbf{W}_i^Q, \mathbf{W}_i^K \in \mathbb{R}^{d_{model} \times d_k}\)\(\mathbf{W}_i^V \in \mathbb{R}^{d_{model} \times d_v}\)

  • 将所有头的输出拼接(Concat)
\[ \text{MultiHead} = \text{Concat}(\text{head}_1, \text{head}_2, \dots, \text{head}_h) \]
  • 通过一个输出投影矩阵 \(\mathbf{W}^O \in \mathbb{R}^{hd_v \times d_{model}}\) 将维度映射回 \(d_{model}\)
\[ \text{MultiHead}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{Concat}(\text{head}_1, \dots, \text{head}_h)\mathbf{W}^O \]

多个头输出的简单拼接,可能不够;可学习的线性层\(\mathbf{W}^O \in \mathbb{R}^{hd_v \times d_{model}}\),目标是如何最优地加权、混合和转换这些来自不同子空间的特征

参数设置(原始论文): \(d_{model} = 512\)\(h = 8\),因此 \(d_k = d_v = 512/8 = 64\)

计算量与单头注意力相当: 虽然有 \(h\) 个头,但每个头的维度是 \(d_{model}/h\),因此总计算量约为 \(O(n^2 \cdot d_{model})\),与使用完整维度的单头注意力相同。

三种注意力的使用方式

Transformer 中注意力被用于三种不同的场景:

类型 Q 的来源 K/V 来源 作用
编码器自注意力 编码器当前层 编码器当前层 每个输入 token 关注所有其他输入 token
解码器掩码自注意力 解码器当前层 解码器当前层(掩码) 每个输出 token 只能关注它之前的 token
编码器-解码器交叉注意力 解码器当前层 编码器输出 解码器关注输入序列的相关部分

掩码自注意力(Masked Self-Attention):在解码器中,为了保持自回归生成的特性(生成第 \(t\) 个词时不能看到第 \(t+1\) 个词),需要在 Softmax 之前将未来位置的分数设为 \(-\infty\)

\[ \text{MaskedAttention}(\mathbf{Q}, \mathbf{K}, \mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^\top}{\sqrt{d_k}} + \mathbf{M}\right)\mathbf{V} \]

其中掩码矩阵 \(\mathbf{M}\) 的定义为:

\[ M_{ij} = \begin{cases} 0 & \text{if } i \geq j \\ -\infty & \text{if } i < j \end{cases} \]

为什么不少大模型,如GPT,只使用了解码器?

大模型的输出可以视为输入的“续写”,而翻译任务的输出是输入的“转换”,这恰恰是纯解码器架构与编码器-解码器架构的根本分野所在。

在续写任务中,训练数据就是一段段连续的文本。模型的学习目标极其纯粹:给定前N个token,预测第N+1个token。这不需要区分“源”和“目标”,整个文本序列既是输入也是(偏移一位的)输出。纯解码器的掩码注意力完美实现了这一点——它只允许看前面的词来预测下一个词。一个成功的续写模型,必须在内部隐式地学会理解输入。为了能写出合理的下文,它必须理解上文的语法、语义、逻辑和常识。因此,在纯解码器中,“理解”不是通过一个独立的编码器模块完成的,而是作为生成能力的一个必要组成部分,在自回归生成的过程中被学会的。

也许,你可以玩一玩操作系统Lab: GPT-2 并行推理

翻译需要捕捉输入句子的完整语义、句法结构和细微含义,这通常需要双向的、全局的上下文信息,编码器的双向注意力机制为此提供了强大支持。大模型选择纯解码器,是因为它们的首要目标是成为一个通用的、强大的续写引擎。而续写任务在数据形式和目标上的纯粹性,使得结构更简单、更易扩展的纯解码器架构成为了最优解。通过巧妙的提示工程,纯解码器模型用一个统一的“续写”机制,解决了千变万化的任务。这证明了:当模型规模足够大、训练数据足够丰富时,一个极度擅长“自回归续写”的单一模型(纯解码器),其内部涌现出的理解和推理能力,足以泛化到需要复杂转换的任务上,从而在许多场景下简化或替代了专门的编码器-解码器架构。

前馈网络

在每个注意力子层之后,Transformer 使用一个逐位置(position-wise)前馈网络(Position-wise Feed-Forward Network)——即对序列中的每个位置独立地、相同地应用一个两层的全连接网络:

\[ \text{FFN}(\mathbf{x}) = \max(0, \mathbf{x}\mathbf{W}_1 + \mathbf{b}_1)\mathbf{W}_2 + \mathbf{b}_2 \]

其中 \(\mathbf{W}_1 \in \mathbb{R}^{d_{model} \times d_{ff}}\)\(\mathbf{W}_2 \in \mathbb{R}^{d_{ff} \times d_{model}}\)\(d_{ff} = 2048\)

这个 FFN 可以理解为:注意力层负责信息的选择与汇聚("看哪里"),而 FFN 负责信息的变换与处理("怎么用")。

前馈网络,也很重要

这个两层的FFN结构(一个带有ReLU激活的升维线性层,再接一个降维线性层)解决了两个问题:

  • 引入关键的非线性:如果只有一层线性变换(\(\mathbf{x}\mathbf{W} + \mathbf{b}\)),那么无论堆叠多少层注意力,整个Transformer块在数学上仍然等价于一个巨大的线性变换,因为自注意力本身也是线性投影的加权和。这将严重限制模型的表达能力,无法拟合复杂的函数。中间的ReLU激活函数(\(\max(0, \cdot)\))提供了必需的非线性,使模型能够学习更复杂的模式。

  • 构建一个“通用函数逼近器”:这个特定的两层结构(升维 → 非线性激活 → 降维)本质上是一个小型但高容量的多层感知机。理论上,这样的网络结构可以逼近任何连续函数。其设计精髓在于:第一层(升维层):将每个位置的表示从 \(d_{\text{model}}\)(例如512)投影到一个高维隐藏空间 \(d_{ff}\)(例如2048)。这相当于将信息“展开”到一个更高维、更具表现力的空间中,便于进行复杂的特征加工。ReLU激活:在高维空间中进行非线性过滤,引入稀疏性(将部分神经元置零),这是学习复杂关系的关键。第二层(降维层):将处理后的信息从高维隐藏空间投影回原始的模型维度 \(d_{\text{model}}\),以便与残差连接相加,并输入到下一层。

一层无法实现非线性变换,而三层及以上则会大幅增加计算量且收益递减,这个两层设计是效果与效率的平衡。

Add & Norm

Transformer 的每个子层(注意力层或 FFN)都包裹在残差连接层归一化之中:

\[ \text{LayerNorm}(\mathbf{x} + \text{Sublayer}(\mathbf{x})) \]
  • 残差连接:保证梯度可以直接流过,使得 Transformer 可以堆叠到 6 层甚至更多。(缓解了梯度在深层传递中逐渐缩小消失或放大爆炸的问题)
  • 层归一化(Layer Normalization):对每个样本在特征维度上进行归一化,稳定训练过程。与 Batch Normalization 不同,Layer Normalization 不依赖于 batch 中的其他样本,更适合变长序列。

过去与现在,我们如何操作?

对于输入向量 \(\mathbf{x} \in \mathbb{R}^{d_{model}}\)(代表某个样本的某个时间步),层归一化计算: $$ \text{LayerNorm}(\mathbf{x}) = \gamma \odot \frac{\mathbf{x} - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta $$ 其中: * \(\mu, \sigma\)\(\mathbf{x}\) 在所有 d_model 个特征上计算的均值和标准差。 * \(\gamma\)\(\beta\) 是可学习的缩放和偏移参数,用于恢复模型可能需要的表达能力(例如,如果归一化后全为0-1分布,会丢失信息)。 * \(\epsilon\) 是一个极小值,防止除零错误。

为什么放在残差之后?将层归一化应用在“输入 + 子层输出”之后,其主要作用是:稳定前向传播:确保输入到下一层的信号具有稳定的均值和方差,防止数值在深层网络中出现漂移或饱和;平滑优化地形:归一化使得损失函数对参数的变化更加平滑,允许使用更大的学习率,并加速训练收敛。

LayerNorm(x + Sublayer(x)) 是原始论文的 Post-LN(层归一化在后)结构。在实践中,许多现代大模型(如GPT、LLaMA)采用了 Pre-LN 结构: 输出 = x + Sublayer(LayerNorm(x)) * Pre-LN的优势:将层归一化置于子层之前,能更早地稳定输入分布,使得训练更深层的模型时更加稳定,梯度流更顺畅。 * Post-LN vs Pre-LN:Post-LN在输出附近进行强归一化,有时能带来更好的最终性能,但训练更不稳定;Pre-LN则大大提升了训练的稳定性和收敛速度,成为当前训练超大规模模型的首选。

Transformer 的完整数据流

让我们追踪一个句子在 Transformer 中的完整旅程:

编码器侧:

  1. 输入序列 "疾旋鼬 love 工管"(love 工管,不是桃旋鼬) 经过词嵌入和位置编码,得到 \(\mathbf{X} \in \mathbb{R}^{3 \times 512}\)
  2. 进入第 1 层编码器:
  3. 多头自注意力:每个词关注所有词 → 输出 \(\mathbf{A}_1\)
  4. Add & Norm:\(\text{LayerNorm}(\mathbf{X} + \mathbf{A}_1)\)
  5. FFN:逐位置变换 → 输出 \(\mathbf{F}_1\)
  6. Add & Norm:\(\text{LayerNorm}((\mathbf{X} + \mathbf{A}_1) + \mathbf{F}_1)\)
  7. 重复 \(N\) 层(原始论文 \(N = 6\)
  8. 最终输出编码表示 \(\mathbf{C} \in \mathbb{R}^{3 \times 512}\)

解码器侧:

  1. 输入已生成的序列(训练时为 ground truth,推理时为已生成的部分)
  2. 进入解码器层:
  3. 掩码自注意力:每个词只关注前面的词
  4. Add & Norm
  5. 交叉注意力:\(\mathbf{Q}\) 来自解码器,\(\mathbf{K}, \mathbf{V}\) 来自编码器输出 \(\mathbf{C}\)
  6. Add & Norm
  7. FFN
  8. Add & Norm
  9. 最终经过线性层 + Softmax,输出下一个词的概率分布

为什么 Transformer 如此成功?

  1. 并行性:与 RNN 的串行计算不同,Transformer 中所有位置的注意力计算可以同时进行,训练速度大幅提升。

  2. 直接的长距离连接:任意两个位置之间只需一步注意力计算即可建立连接,路径长度为 \(O(1)\)。相比之下,RNN 需要 \(O(n)\) 步,CNN 需要 \(O(\log_k n)\) 层。

  3. 灵活性:多头注意力可以同时学习多种不同类型的关系(句法、语义、位置等),每种关系由不同的头负责。

  4. 可扩展性:Transformer 架构对算力的利用效率极高,可以轻松扩展到数十亿参数,催生了 GPT、BERT、T5 等大型语言模型。

回顾与习题:简易Transformer的实现

理论讲完了,现在让我们从零手写一个完整的 Transformer。

Attention 注意力机制

以 3 个 token、每个 4 维为例,展示完整的计算流程。

import torch
import math

# ===== 缩放点积注意力 =====
# 核心公式: Attention(Q, K, V) = softmax(Q @ K^T / sqrt(d_k)) @ V

# --- 设置参数 ---
seq_len, d_k, d_v = 3, 4, 4  # 3 个 token, Q/K 维度 4, V 维度 4

# --- 输入: 3 个 token, 每个 4 维 ---
X = torch.tensor([[1.0, 0.0, 1.0, 0.0],
                   [0.0, 2.0, 0.0, 2.0],
                   [1.0, 1.0, 1.0, 1.0]])  # (3, 4)

# --- 投影矩阵 ---
W_Q = torch.randn(d_k, d_k)  # (4, 4)
W_K = torch.randn(d_k, d_k)  # (4, 4)
W_V = torch.randn(d_k, d_v)  # (4, 4)

# --- Step 1: 计算 Q, K, V (线性投影) ---
Q = X @ W_Q  # (3, 4) @ (4, 4) = (3, 4)
K = X @ W_K  # (3, 4) @ (4, 4) = (3, 4)
V = X @ W_V  # (3, 4) @ (4, 4) = (3, 4)

# --- Step 2: 计算注意力分数 (Q 与 K 的点积) ---
scores = Q @ K.T  # (3, 4) @ (4, 3) = (3, 3)
# scores[i][j] = token_i 的 query 与 token_j 的 key 的点积 → 衡量相似度

# --- Step 3: 缩放 (防止点积过大导致 softmax 饱和) ---
scores = scores / math.sqrt(d_k)  # 每个元素除以 sqrt(4) = 2.0

# --- Step 4: Softmax 归一化 → 注意力权重 ---
weights = torch.softmax(scores, dim=-1)  # (3, 3), 每行之和为 1

# --- Step 5: 加权求和 ---
output = weights @ V  # (3, 3) @ (3, 4) = (3, 4)
# output[i] = 所有 token 的 value 的加权和, 权重由 token_i 的注意力决定

# --- 因果掩码 (解码器用): 下三角矩阵遮住未来位置 ---
causal_mask = torch.tril(torch.ones(seq_len, seq_len))  # (3, 3) 下三角
scores_masked = scores.masked_fill(causal_mask == 0, float('-inf'))
# 被遮住的位置填 -inf, softmax 后变为 0 → token_i 只能看到 token_0..token_i
weights_masked = torch.softmax(scores_masked, dim=-1)

完整代码见 code/01_attention.py

Multi-Head Attention

将 Q, K, V 投影到 h 个低维子空间,分别计算注意力后拼接。以 d_model=8, num_heads=2 为例。

import torch
import torch.nn as nn
import math

def scaled_dot_product_attention(Q, K, V, mask=None):
    """缩放点积注意力"""
    d_k = Q.size(-1)
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, float('-inf'))
    weights = torch.softmax(scores, dim=-1)
    return torch.matmul(weights, V), weights

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        assert d_model % num_heads == 0, "d_model 必须能被 num_heads 整除"
        self.d_model = d_model
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 每个头的维度
        # Q, K, V 各一个线性投影 (无偏置)
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        # 输出投影
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, Q, K, V, mask=None):
        B = Q.size(0)
        # Step 1: 线性投影  (B, seq, d_model) @ (d_model, d_model) = (B, seq, d_model)
        Q = self.W_q(Q)
        K = self.W_k(K)
        V = self.W_v(V)
        # Step 2: 拆分成多头  (B, seq, d_model) -> (B, seq, h, d_k) -> (B, h, seq, d_k)
        Q = Q.view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        K = K.view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        V = V.view(B, -1, self.num_heads, self.d_k).transpose(1, 2)
        # Step 3: 每个头独立计算注意力  (B, h, seq, d_k) -> (B, h, seq, d_k)
        attn_output, weights = scaled_dot_product_attention(Q, K, V, mask)
        # Step 4: 拼接所有头  (B, h, seq, d_k) -> (B, seq, h*d_k) = (B, seq, d_model)
        attn_output = attn_output.transpose(1, 2).contiguous().view(B, -1, self.d_model)
        # Step 5: 输出投影  (B, seq, d_model) @ (d_model, d_model) = (B, seq, d_model)
        output = self.W_o(attn_output)
        return output, weights

# --- 验证 ---
mha = MultiHeadAttention(d_model=8, num_heads=2)
x = torch.randn(1, 3, 8)  # batch=1, 3 个 token, 8 维
output, weights = mha(x, x, x)  # 自注意力: Q=K=V=x
# output shape: (1, 3, 8)
# weights shape: (1, 2, 3, 3) → 2 个头, 每个头 3×3 注意力矩阵

完整代码见 code/02_multi_head_attention.py

LayerNorm 与残差连接

import torch
import torch.nn as nn

class LayerNorm(nn.Module):
    """手写层归一化: 对每个样本在特征维度上归一化"""
    def __init__(self, d_model, eps=1e-5):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(d_model))   # 可学习缩放
        self.beta = nn.Parameter(torch.zeros(d_model))    # 可学习偏移
        self.eps = eps

    def forward(self, x):
        # Step 1: 计算均值 (沿最后一维, 即特征维度)
        mean = x.mean(dim=-1, keepdim=True)  # (B, seq, 1)
        # Step 2: 计算方差
        var = x.var(dim=-1, keepdim=True, unbiased=False)  # (B, seq, 1)
        # Step 3: 归一化
        x_norm = (x - mean) / torch.sqrt(var + self.eps)  # (B, seq, d_model)
        # Step 4: 仿射变换 (可学习参数)
        return self.gamma * x_norm + self.beta  # (B, seq, d_model)

# --- 手算验证: x = [1, 2, 3, 4] ---
# mean = (1+2+3+4)/4 = 2.5
# var  = [(1-2.5)^2 + (2-2.5)^2 + (3-2.5)^2 + (4-2.5)^2] / 4 = 1.25
# x_norm = (x - 2.5) / sqrt(1.25) = [-1.342, -0.447, 0.447, 1.342]
# 验证: 均值=0, 方差=1 ✓

# --- 与 PyTorch 内置 LayerNorm 对比 ---
ln = LayerNorm(4)
ln_torch = nn.LayerNorm(4)
ln_torch.weight.data = ln.gamma.data.clone()
ln_torch.bias.data = ln.beta.data.clone()
x = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
diff = (ln(x) - ln_torch(x)).abs().max().item()  # 差异 < 1e-6 ✓

# --- 残差连接 (Post-LN) ---
# 梯度可以通过 '+' 直接流回, 不经过子层权重 → 缓解梯度消失
# output = LayerNorm(x + sublayer(x))
sublayer_output = torch.randn(1, 4)
output = ln(x + sublayer_output)  # 先残差相加, 再归一化

完整代码见 code/03_layernorm_residual.py

FFN前馈网络

对每个位置独立应用的两层全连接网络——升维、ReLU、降维。

import torch
import torch.nn as nn

class FeedForward(nn.Module):
    """逐位置前馈网络: d_model -> d_ff (升维) -> ReLU -> d_model (降维)"""
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.linear1 = nn.Linear(d_model, d_ff)   # 升维: d_model -> d_ff
        self.linear2 = nn.Linear(d_ff, d_model)    # 降维: d_ff -> d_model
        self.relu = nn.ReLU()

    def forward(self, x):
        # Step 1: 第一层线性变换 (升维)
        h = self.linear1(x)    # (B, seq, d_model) -> (B, seq, d_ff)
        # Step 2: ReLU 激活 (引入非线性, 负值置零)
        h = self.relu(h)       # (B, seq, d_ff)
        # Step 3: 第二层线性变换 (降维)
        output = self.linear2(h)  # (B, seq, d_ff) -> (B, seq, d_model)
        return output

# --- 验证 ---
ffn = FeedForward(d_model=16, d_ff=64)
x = torch.randn(2, 5, 16)  # batch=2, 5 个 token, 16 维
output = ffn(x)  # (2, 5, 16), shape 不变
# 参数量: (16×64+64) + (64×16+16) = 1088 + 1040 = 2128

完整代码见 code/04_ffn.py

Transformer Block

import torch
import torch.nn as nn
import math

# (此处复用上面定义的 scaled_dot_product_attention, MultiHeadAttention, LayerNorm, FeedForward)

class EncoderBlock(nn.Module):
    """编码器块: 自注意力 -> Add&Norm -> FFN -> Add&Norm"""
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.self_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)

    def forward(self, x, mask=None):
        # Sublayer 1: 自注意力 (Q=K=V=x, 每个 token 关注所有 token)
        attn_out, _ = self.self_attn(x, x, x, mask)
        x = self.norm1(x + attn_out)  # 残差连接 + LayerNorm
        # Sublayer 2: 前馈网络 (逐位置独立变换)
        ffn_out = self.ffn(x)
        x = self.norm2(x + ffn_out)   # 残差连接 + LayerNorm
        return x

class DecoderBlock(nn.Module):
    """解码器块: 掩码自注意力 -> Add&Norm -> 交叉注意力 -> Add&Norm -> FFN -> Add&Norm"""
    def __init__(self, d_model, num_heads, d_ff):
        super().__init__()
        self.masked_self_attn = MultiHeadAttention(d_model, num_heads)
        self.cross_attn = MultiHeadAttention(d_model, num_heads)
        self.ffn = FeedForward(d_model, d_ff)
        self.norm1 = LayerNorm(d_model)
        self.norm2 = LayerNorm(d_model)
        self.norm3 = LayerNorm(d_model)

    def forward(self, x, enc_output, src_mask=None, tgt_mask=None):
        # Sublayer 1: 掩码自注意力 (因果掩码, token 只能看前面的 token)
        attn_out, _ = self.masked_self_attn(x, x, x, tgt_mask)
        x = self.norm1(x + attn_out)
        # Sublayer 2: 交叉注意力 (Q 来自解码器, K/V 来自编码器输出)
        cross_out, _ = self.cross_attn(x, enc_output, enc_output, src_mask)
        x = self.norm2(x + cross_out)
        # Sublayer 3: 前馈网络
        ffn_out = self.ffn(x)
        x = self.norm3(x + ffn_out)
        return x

# --- 验证 ---
enc = EncoderBlock(d_model=16, num_heads=4, d_ff=64)
dec = DecoderBlock(d_model=16, num_heads=4, d_ff=64)
src = torch.randn(2, 6, 16)  # 源序列: batch=2, 6 个 token
tgt = torch.randn(2, 4, 16)  # 目标序列: batch=2, 4 个 token
causal_mask = torch.tril(torch.ones(4, 4)).unsqueeze(0)  # 因果掩码

enc_out = enc(src)                   # (2, 6, 16)
dec_out = dec(tgt, enc_out, tgt_mask=causal_mask)  # (2, 4, 16)

完整代码见 code/05_transformer_block.py

Tiny Transformer

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

# (此处复用上面定义的所有组件)

class TinyTransformer(nn.Module):
    """完整的 Transformer: Embedding + 位置编码 + N层Encoder + N层Decoder + 输出投影"""
    def __init__(self, src_vocab_size, tgt_vocab_size, d_model=16,
                 num_heads=4, d_ff=64, num_layers=2, max_len=100):
        super().__init__()
        self.d_model = d_model
        # 嵌入层: token ID -> d_model 维向量
        self.src_embedding = nn.Embedding(src_vocab_size, d_model)
        self.tgt_embedding = nn.Embedding(tgt_vocab_size, d_model)
        # 正弦位置编码 (不可学习, 预计算)
        self.register_buffer('pos_encoding', self._get_positional_encoding(max_len, d_model))
        # N 层编码器 + N 层解码器
        self.encoder_layers = nn.ModuleList([
            EncoderBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
        ])
        self.decoder_layers = nn.ModuleList([
            DecoderBlock(d_model, num_heads, d_ff) for _ in range(num_layers)
        ])
        # 输出投影: d_model -> 词表大小
        self.output_proj = nn.Linear(d_model, tgt_vocab_size)

    def _get_positional_encoding(self, max_len, d_model):
        """生成正弦位置编码"""
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1).float()
        div_term = torch.exp(torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)  # 偶数维度用 sin
        pe[:, 1::2] = torch.cos(position * div_term)  # 奇数维度用 cos
        return pe.unsqueeze(0)  # (1, max_len, d_model)

    def encode(self, src):
        """编码器前向: token IDs -> 编码表示"""
        seq_len = src.size(1)
        # Embedding * sqrt(d_model) + 位置编码
        x = self.src_embedding(src) * math.sqrt(self.d_model) + self.pos_encoding[:, :seq_len, :]
        for layer in self.encoder_layers:
            x = layer(x)
        return x  # (B, src_len, d_model)

    def decode(self, tgt, enc_output):
        """解码器前向: token IDs + 编码输出 -> 解码表示"""
        seq_len = tgt.size(1)
        causal_mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0).to(tgt.device)
        x = self.tgt_embedding(tgt) * math.sqrt(self.d_model) + self.pos_encoding[:, :seq_len, :]
        for layer in self.decoder_layers:
            x = layer(x, enc_output, tgt_mask=causal_mask)
        return x  # (B, tgt_len, d_model)

    def forward(self, src, tgt):
        """完整前向: 源序列 + 目标序列 -> logits"""
        enc_output = self.encode(src)           # (B, src_len, d_model)
        dec_output = self.decode(tgt, enc_output)  # (B, tgt_len, d_model)
        logits = self.output_proj(dec_output)   # (B, tgt_len, vocab_size)
        return logits

# --- 完整验证 ---
model = TinyTransformer(src_vocab_size=50, tgt_vocab_size=50,
                        d_model=16, num_heads=4, d_ff=64, num_layers=2)

# 模拟输入 (token IDs)
src_ids = torch.tensor([[3, 7, 12, 5, 9]])   # 源序列: 5 个 token
tgt_ids = torch.tensor([[2, 15, 8, 4]])       # 目标序列: 4 个 token

# 前向传播
logits = model(src_ids, tgt_ids)  # (1, 4, 50) → 每个位置 50 维概率分布

# 训练一步
target = torch.tensor([[5, 12, 3, 8]])  # ground truth
loss = F.cross_entropy(logits.view(-1, 50), target.view(-1))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
optimizer.zero_grad()
loss.backward()
optimizer.step()
# loss ≈ 3.9, 反向传播成功 ✓

完整代码(含详细计算过程演示)见 code/06_tiny_transformer.py


参考文献

  • Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. NeurIPS. 【伟大,无需多言】Clink Here to Download
  • Bahdanau, D., Cho, K., & Bengio, Y. (2014). Neural Machine Translation by Jointly Learning to Align and Translate. ICLR.【首次将注意力机制成功应用于神经机器翻译的开创性工作】
  • Luong, M.T., Pham, H., & Manning, C.D. (2015). Effective Approaches to Attention-based Neural Machine Translation. EMNLP.【提出了更简洁的全局注意力(Global Attention)和计算效率更高的局部注意力(Local Attention)机制】
  • He, K., Zhang, X., Ren, S., & Sun, J. (2015). Deep Residual Learning for Image Recognition. CVPR.【提出了残差网络(ResNet),通过残差连接解决了深度网络训练中的退化问题】
  • Hochreiter, S., & Schmidhuber, J. (1997). Long Short-Term Memory. Neural Computation.【提出了长短期记忆网络(LSTM),通过门控机制有效缓解了循环神经网络中的梯度消失问题】
  • Cho, K., van Merriënboer, B., Gulcehre, C., et al. (2014). Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation. EMNLP.【提出了RNN编码器-解码器(Seq2Seq)框架以及门控循环单元(GRU)】

2026.05.09 深夜 不开心

昨天是我的生日……可是我没想到我那天凌晨是哭着睡着的……唉,鼬一直在哭

写到这里时正在听的音乐

疾旋鼬在闹