《Attention Is All You Need》的 Encoder–Decoder Transformer 一次完整训练

下面按 2017 年原始论文《Attention Is All You Need》 的 Encoder--Decoder Transformer，整理一次完整训练过程。重点是：训练时使用 Teacher Forcing，并借助因果掩码并行计算所有目标位置，不像推理那样逐个生成 token。

1. 模型结构与主要符号

原论文基础模型的主要参数为：

参数	含义	基础模型取值
NNN	Encoder、Decoder 各自的层数	6
dmodeld_{\text{model}}dmodel	隐藏向量维度	512
hhh	注意力头数	8
dk=dvd_k=d_vdk=dv	每个头的键、值维度	64
dffd_{\text{ff}}dff	前馈网络中间维度	2048
PdropP_{\text{drop}}Pdrop	Dropout 概率	0.1

设一个训练样本为：

源序列：x=(x1,...,xn)x=(x_1,\ldots,x_n)x=(x1,...,xn)；
目标序列：y=(y1,...,ym)y=(y_1,\ldots,y_m)y=(y1,...,ym)，其中 ymy_mym 可以视为 <eos>；
词表大小：VVV；
所有向量最后一维通常为 dmodeld_{\text{model}}dmodel。

训练时构造 Decoder 输入：

y~=(⟨bos⟩,y1,...,ym−1) \widetilde y=(\langle bos\rangle,y_1,\ldots,y_{m-1}) y =(⟨bos⟩,y1,...,ym−1)

训练标签仍然是：

y=(y1,...,ym) y=(y_1,\ldots,y_m) y=(y1,...,ym)

也就是将目标序列右移一位后输入 Decoder。例如：

text 复制代码

目标标签：       我   喜欢   机器学习   <eos>
Decoder 输入：<bos>  我   喜欢   机器学习

Decoder 在位置 1 根据 <bos> 预测"我"，在位置 2 根据 <bos> 我 预测"喜欢"，以此类推。

2. 三个基础计算模块

2.1 缩放点积注意力

给定查询、键和值矩阵 Q,K,VQ,K,VQ,K,V，计算：

Attention⁡(Q,K,V)=softmax⁡(QK⊤dk+M)V \operatorname{Attention}(Q,K,V)= \operatorname{softmax} \left( \frac{QK^\top}{\sqrt{d_k}}+M \right)V Attention(Q,K,V)=softmax(dk QK⊤+M)V

其中：

QK⊤QK^\topQK⊤ 计算查询与键之间的相似度；
除以 dk\sqrt{d_k}dk ，避免点积过大导致 softmax 饱和；
MMM 是注意力掩码，允许的位置为 0，禁止的位置为 −∞-\infty−∞；
softmax 在键所在的维度上进行；
最后对 VVV 加权求和。

对于批量输入，若 QQQ 的形状为 $B,n,dk$ $B,n,d_k$ $B,n,dk$ ，K,VK,VK,V 的形状为 $B,m,dk$ $B,m,d_k$ $B,m,dk$ 和 $B,m,dv$ $B,m,d_v$ $B,m,dv$ ，则输出形状为 $B,n,dv$ $B,n,d_v$ $B,n,dv$ 。

2.2 多头注意力

第 iii 个注意力头为：

head⁡i=Attention⁡(QWiQ, KWiK, VWiV) \operatorname{head}_i= \operatorname{Attention} \left( QW_i^Q,\, KW_i^K,\, VW_i^V \right) headi=Attention(QWiQ,KWiK,VWiV)

将所有头拼接并做线性变换：

MHA⁡(Q,K,V)=Concat⁡(head⁡1,...,head⁡h)WO \operatorname{MHA}(Q,K,V)= \operatorname{Concat} (\operatorname{head}_1,\ldots,\operatorname{head}_h)W^O MHA(Q,K,V)=Concat(head1,...,headh)WO

其中：

WiQ∈Rdmodel×dk,WiK∈Rdmodel×dk,WiV∈Rdmodel×dv W_i^Q\in\mathbb R^{d_{\text{model}}\times d_k},\quad W_i^K\in\mathbb R^{d_{\text{model}}\times d_k},\quad W_i^V\in\mathbb R^{d_{\text{model}}\times d_v} WiQ∈Rdmodel×dk,WiK∈Rdmodel×dk,WiV∈Rdmodel×dv

WO∈Rhdv×dmodel W^O\in\mathbb R^{hd_v\times d_{\text{model}}} WO∈Rhdv×dmodel

基础模型中 hdv=8×64=512=dmodelhd_v=8\times64=512=d_{\text{model}}hdv=8×64=512=dmodel，因此输出仍可以接回残差连接。

2.3 位置前馈网络

对每个位置的向量独立使用同一个两层全连接网络：

FFN⁡(x)=max⁡(0,xW1+b1)W2+b2 \operatorname{FFN}(x)= \max(0,xW_1+b_1)W_2+b_2 FFN(x)=max(0,xW1+b1)W2+b2

其中：

W1∈R512×2048,W2∈R2048×512 W_1\in\mathbb R^{512\times2048},\qquad W_2\in\mathbb R^{2048\times512} W1∈R512×2048,W2∈R2048×512

不同序列位置之间不在这里交换信息，位置之间的信息交互由注意力完成。

3. 输入嵌入与位置编码

3.1 词嵌入

源序列和目标右移序列分别经过嵌入层：

Xemb=Ex $x$ ,Yemb=Ey $y\~$ X_{\text{emb}}=E_x $x$ ,\qquad Y_{\text{emb}}=E_y $\\widetilde y$ Xemb=Ex $x$ ,Yemb=Ey $y$

原论文将嵌入向量乘以 dmodel\sqrt{d_{\text{model}}}dmodel ：

Xemb←dmodel Xemb X_{\text{emb}}\leftarrow \sqrt{d_{\text{model}}}\,X_{\text{emb}} Xemb←dmodel Xemb

Yemb←dmodel Yemb Y_{\text{emb}}\leftarrow \sqrt{d_{\text{model}}}\,Y_{\text{emb}} Yemb←dmodel Yemb

原论文还采用了嵌入层与 softmax 前线性层之间的权重共享。

3.2 正弦位置编码

因为注意力本身不包含序列顺序信息，所以需要加入位置编码：

PE(pos,2i)=sin⁡(pos100002i/dmodel) PE_{(pos,2i)}= \sin\left( \frac{pos}{10000^{2i/d_{\text{model}}}} \right) PE(pos,2i)=sin(100002i/dmodelpos)

PE(pos,2i+1)=cos⁡(pos100002i/dmodel) PE_{(pos,2i+1)}= \cos\left( \frac{pos}{10000^{2i/d_{\text{model}}}} \right) PE(pos,2i+1)=cos(100002i/dmodelpos)

然后得到 Encoder、Decoder 的初始输入：

H(0)=Dropout⁡(Xemb+PEx) H^{(0)}= \operatorname{Dropout} \left( X_{\text{emb}}+PE_x \right) H(0)=Dropout(Xemb+PEx)

Z(0)=Dropout⁡(Yemb+PEy) Z^{(0)}= \operatorname{Dropout} \left( Y_{\text{emb}}+PE_y \right) Z(0)=Dropout(Yemb+PEy)

这里的位置编码是固定公式计算的，不需要训练。

4. 三种注意力掩码

训练时通常需要以下掩码。

4.1 Encoder 源序列填充掩码

源序列补齐产生的 <pad> 不应被关注：

Msrc $i,j$ ={0,xj≠⟨pad⟩−∞,xj=⟨pad⟩ M_{\text{src}} $i,j$ = \begin{cases} 0, & x_j\neq\langle pad\rangle\\ -\infty, & x_j=\langle pad\rangle \end{cases} Msrc $i,j$ ={0,−∞,xj=⟨pad⟩xj=⟨pad⟩

它用于 Encoder 自注意力和 Decoder 的交叉注意力。

4.2 Decoder 因果掩码

目标位置 iii 不能看到位置 j>ij>ij>i 的真实标签：

Mcausal $i,j$ ={0,j≤i−∞,j>i M_{\text{causal}} $i,j$ = \begin{cases} 0, & j\le i\\ -\infty, & j>i \end{cases} Mcausal $i,j$ ={0,−∞,j≤ij>i

它保证训练与自回归推理的条件一致：

p(yi∣y<i,x) p(y_i\mid y_{<i},x) p(yi∣y<i,x)

而不是让模型直接看到 yiy_iyi 或未来 token。

4.3 Decoder 目标填充掩码

目标右移输入中的 <pad> 也不应作为有效键参与注意力。实际实现中通常将目标填充掩码与因果掩码合并。

5. Encoder 的完整计算

Encoder 输入为 H(0)H^{(0)}H(0)，共执行 N=6N=6N=6 层。每一层包含两个子层。

5.1 多头自注意力子层

令当前输入为 HHH，则：

A=MHA⁡(H,H,H;Msrc) A=\operatorname{MHA}(H,H,H;M_{\text{src}}) A=MHA(H,H,H;Msrc)

这里 Q,K,VQ,K,VQ,K,V 都来自源序列，因此称为自注意力。

原论文采用 Post-LN，即先残差相加，再做层归一化：

H′=LayerNorm⁡(H+Dropout⁡(A)) H'= \operatorname{LayerNorm} \left( H+\operatorname{Dropout}(A) \right) H′=LayerNorm(H+Dropout(A))

5.2 前馈网络子层

F=FFN⁡(H′) F=\operatorname{FFN}(H') F=FFN(H′)

Hout=LayerNorm⁡(H′+Dropout⁡(F)) H_{\text{out}}= \operatorname{LayerNorm} \left( H'+\operatorname{Dropout}(F) \right) Hout=LayerNorm(H′+Dropout(F))

将 HoutH_{\text{out}}Hout 作为下一层输入。第六层输出记为：

Henc∈RB×n×dmodel H_{\text{enc}}\in\mathbb R^{B\times n\times d_{\text{model}}} Henc∈RB×n×dmodel

它就是 Decoder 交叉注意力使用的记忆表示。

6. Decoder 的完整计算

Decoder 输入为 Z(0)Z^{(0)}Z(0)，同样执行六层。每层包含三个子层。

6.1 掩码多头自注意力

A1=MHA⁡(Z,Z,Z;Mdec-self) A_1= \operatorname{MHA} \left( Z,Z,Z; M_{\text{dec-self}} \right) A1=MHA(Z,Z,Z;Mdec-self)

其中 Mdec-selfM_{\text{dec-self}}Mdec-self 同时处理：

目标序列的因果约束；
目标输入中的 <pad>。

然后进行残差连接和层归一化：

Z′=LayerNorm⁡(Z+Dropout⁡(A1)) Z'= \operatorname{LayerNorm} \left( Z+\operatorname{Dropout}(A_1) \right) Z′=LayerNorm(Z+Dropout(A1))

6.2 Encoder--Decoder 交叉注意力

查询来自 Decoder，键和值来自 Encoder：

Q=Z′,K=Henc,V=Henc Q=Z',\qquad K=H_{\text{enc}},\qquad V=H_{\text{enc}} Q=Z′,K=Henc,V=Henc

因此：

A2=MHA⁡(Z′,Henc,Henc;Msrc) A_2= \operatorname{MHA} \left( Z', H_{\text{enc}}, H_{\text{enc}}; M_{\text{src}} \right) A2=MHA(Z′,Henc,Henc;Msrc)

再执行：

Z′′=LayerNorm⁡(Z′+Dropout⁡(A2)) Z''= \operatorname{LayerNorm} \left( Z'+\operatorname{Dropout}(A_2) \right) Z′′=LayerNorm(Z′+Dropout(A2))

交叉注意力使每个目标位置都可以读取源序列的所有有效位置，这里没有因果掩码，因为完整的源序列在训练与推理时都已知。

6.3 前馈网络

A3=FFN⁡(Z′′) A_3=\operatorname{FFN}(Z'') A3=FFN(Z′′)

Zout=LayerNorm⁡(Z′′+Dropout⁡(A3)) Z_{\text{out}} = \operatorname{LayerNorm} \left( Z''+\operatorname{Dropout}(A_3) \right) Zout=LayerNorm(Z′′+Dropout(A3))

第六层输出记为：

Zdec∈RB×m×dmodel Z_{\text{dec}} \in \mathbb R^{B\times m\times d_{\text{model}}} Zdec∈RB×m×dmodel

7. 输出概率分布

将 Decoder 最后一层输出通过线性层：

O=ZdecWout+bout O=Z_{\text{dec}}W_{\text{out}}+b_{\text{out}} O=ZdecWout+bout

其中：

O∈RB×m×V O\in\mathbb R^{B\times m\times V} O∈RB×m×V

对每个目标位置分别做 softmax：

Pb,t,k=exp⁡(Ob,t,k)∑j=1Vexp⁡(Ob,t,j) P_{b,t,k} = \frac{\exp(O_{b,t,k})} {\sum_{j=1}^{V}\exp(O_{b,t,j})} Pb,t,k=∑j=1Vexp(Ob,t,j)exp(Ob,t,k)

Pb,tP_{b,t}Pb,t 就是第 bbb 个样本在目标位置 ttt 上，对下一个 token 的预测分布。

注意：**训练时不需要对每个位置执行 argmax 或采样，也不需要把预测 token 再嵌入后喂回模型。**所有位置已经在 Teacher Forcing 和因果掩码下并行算完。

8. 标签平滑与损失函数

原论文使用值为 0.10.10.1 的标签平滑，即：

εls=0.1 \varepsilon_{\text{ls}}=0.1 εls=0.1

对于真实类别 yyy，一种常见表示是构造平滑目标分布：

qk={1−εls+εlsV,k=yεlsV,k≠y q_k = \begin{cases} 1-\varepsilon_{\text{ls}}+\dfrac{\varepsilon_{\text{ls}}}{V}, & k=y\\ $6pt$ \dfrac{\varepsilon_{\text{ls}}}{V}, & k\neq y \end{cases} qk=⎩ ⎨ ⎧1−εls+Vεls,Vεls,k=yk=y

然后计算平滑后的交叉熵：

ℓt=−∑k=1Vqt,klog⁡Pt,k \ell_t=-\sum_{k=1}^{V}q_{t,k}\log P_{t,k} ℓt=−k=1∑Vqt,klogPt,k

只统计非 <pad> 的目标位置。设有效目标 token 集合为 T\mathcal TT， batch 损失可以写为：

L=1∣T∣∑t∈Tℓt \mathcal L = \frac{1}{|\mathcal T|} \sum_{t\in\mathcal T}\ell_t L=∣T∣1t∈T∑ℓt

具体实现对平滑概率分配到 VVV 个类还是 V−1V-1V−1 个非真实类可能略有差异，但核心都是：不给真实标签 100% 的概率质量。

9. 反向传播与参数更新

对损失关于所有可训练参数求梯度：

gt=∇θL(θt−1) g_t=\nabla_{\theta}\mathcal L(\theta_{t-1}) gt=∇θL(θt−1)

梯度会依次经过：

输出线性层；
六层 Decoder；
Decoder 交叉注意力连接到 Encoder 的分支；
六层 Encoder；
源、目标嵌入层；
所有注意力投影矩阵、FFN、LayerNorm 等参数。

原论文使用 Adam 优化器：

β1=0.9,β2=0.98,ϵ=10−9 \beta_1=0.9,\qquad \beta_2=0.98,\qquad \epsilon=10^{-9} β1=0.9,β2=0.98,ϵ=10−9

学习率随训练步数变化：

ηt=dmodel−0.5⋅min⁡(t−0.5, t⋅twarmup−1.5) \eta_t = d_{\text{model}}^{-0.5} \cdot \min \left( t^{-0.5},\, t\cdot t_{\text{warmup}}^{-1.5} \right) ηt=dmodel−0.5⋅min(t−0.5,t⋅twarmup−1.5)

基础模型中：

twarmup=4000 t_{\text{warmup}}=4000 twarmup=4000

因此：

前 4000 步，学习率近似线性增大；
之后，学习率按 t−1/2t^{-1/2}t−1/2 逐渐减小。

然后使用 Adam 更新参数：

θt=AdamUpdate⁡(θt−1,gt,ηt) \theta_t= \operatorname{AdamUpdate} (\theta_{t-1},g_t,\eta_t) θt=AdamUpdate(θt−1,gt,ηt)

10. 单步训练算法伪代码

下面省略部分张量广播细节，但保留完整的数据流：

text 复制代码

算法：原始 Transformer 的一次训练迭代

输入：
    源 token 批次 S，形状 [B, n]
    完整目标 token 批次 Y，形状 [B, m]
    当前参数 θ
    当前训练步数 step

1. 构造 Decoder 输入：
      dec_in = [BOS, Y[:, 0:m-1]]
      label  = Y

2. 构造掩码：
      src_mask       = 源序列 padding 掩码
      dec_self_mask  = 目标 padding 掩码 + 下三角因果掩码

3. 源序列输入表示：
      H = sqrt(d_model) * SourceEmbedding(S) + SourcePE
      H = Dropout(H)

4. Encoder，重复 6 层：
      A = MultiHeadAttention(
              Q=H, K=H, V=H,
              mask=src_mask)

      H = LayerNorm(H + Dropout(A))

      F = FFN(H)
      H = LayerNorm(H + Dropout(F))

   得到：
      memory = H

5. 目标右移序列输入表示：
      Z = sqrt(d_model) * TargetEmbedding(dec_in) + TargetPE
      Z = Dropout(Z)

6. Decoder，重复 6 层：
      A1 = MultiHeadAttention(
               Q=Z, K=Z, V=Z,
               mask=dec_self_mask)

      Z = LayerNorm(Z + Dropout(A1))

      A2 = MultiHeadAttention(
               Q=Z, K=memory, V=memory,
               mask=src_mask)

      Z = LayerNorm(Z + Dropout(A2))

      F = FFN(Z)
      Z = LayerNorm(Z + Dropout(F))

7. 输出词表分布：
      logits = OutputLinear(Z)
      probs  = Softmax(logits, dim=vocabulary)

8. 计算损失：
      loss = LabelSmoothedCrossEntropy(
                 probs,
                 label,
                 ignore=PAD,
                 smoothing=0.1)

9. 反向传播：
      gradient = Backward(loss)

10. 计算当前学习率：
      lr = d_model^(-0.5) *
           min(step^(-0.5),
               step * 4000^(-1.5))

11. 更新模型：
      θ = AdamUpdate(
              θ, gradient, lr,
              beta1=0.9,
              beta2=0.98,
              eps=1e-9)

输出：
    更新后的参数 θ

11. 整个计算过程的一行概括

源序列→嵌入+位置编码六层 Encoder→memory六层 Decoder←右移目标序列+因果掩码→线性层+Softmax各位置下一 token 分布→标签平滑交叉熵反向传播+Adam 更新 \boxed{ \text{源序列} \xrightarrow{\text{嵌入+位置编码}} \text{六层 Encoder} \xrightarrow{\text{memory}} \text{六层 Decoder} \leftarrow \text{右移目标序列+因果掩码} \xrightarrow{\text{线性层+Softmax}} \text{各位置下一 token 分布} \xrightarrow{\text{标签平滑交叉熵}} \text{反向传播+Adam 更新} } 源序列嵌入+位置编码六层 Encodermemory 六层 Decoder←右移目标序列+因果掩码线性层+Softmax 各位置下一 token 分布标签平滑交叉熵反向传播+Adam 更新

最关键的理解是：训练阶段计算的是"给定真实历史 y<ty_{<t}y<t，预测 yty_tyt"的所有位置损失；并不存在给未知的下一个 token 初始化一个向量，再通过迭代逐渐得到它的过程。