‘Annotated’《The Annotated Transformer》(Ep.2)
104–156 分钟
“A vast similitude interlocks all.”
Your Sincerely
Self-Attention

承接Ep.1文,本文我们继续进一步讲解上一文中构建的模型的实例化(代码需要承接之前的!)

全文阅读时间较长,主要借助网络经典材料加以个人思考解读,如有误欢迎斧正

书接上回,我们在构建好Transformer整体模型之后要将其投入使用,然而一个最基本的深度学习任务,一般需要以下这些环节:

Data
Model
Criterion
Loop
Optimizer (Update ↺)

现在我们只是拥有了model模型这个大脑(由Transformer完整架构支持),其他步骤仍需要我们慢慢实现

Notice:Ep.1末尾我们提到,《The Annotated Transformer》原文后续实现的各种工具对于本人初学阶段有些浪费时间/必要性不大,所以以下我们着重于实现原文中所述的“一个简单任务”(Copy Task),即复制任务:让机器学会复制我们提供的数字序列,下文即围绕这个任务进行相应的全步骤分析处理准备

5%

我们先来看Data这个大块,这一步处理十分重要,而且还可以划分成几个小版块:

Data Block
1. Generation
data_generator

产生随机数字序列
确定起始符(BOS)

2. Optimization
batch_size_fn

动态决策 Batch 大小
防止显存溢出(OOM)

3. Processing
Batch Class

封装数据生成 Mask
确定训练目标(target_y)

data_generator版块:负责产生随机数字序列训练模型是“生产者”。一般的深度学习任务都会提供相应的数据集用于训练、测试,但由于我们目标任务是Copy Task,故不妨直接内存生成是最快、最方便的。

所以在data_generator板块我们的Goal是:随机生成两个完全一样的Tensor张量,分别作为source和target,而且要求第一个符号都是起始符<BOS>对应的编码(通常取1)

batch_size_fn版块:主要是负责动态优化调整每个Batch所含数据量,是“调度者”。既然每一批(Batch)的句子长度可能不同,它要根据总 Token 数而不是句子条数来决定这一波给 GPU 喂多少,从而榨干显存又不至于让显卡炸掉(OOM)

所以在batch_size_fn版块我们的Goal是:该函数的输出一个整数,该整数代表:为了不让显存爆炸,当前这一批(Batch)最多可以包含多少条数据(样本数)

当然,在我们的Copy Task任务中由于data_generator总是输出大小一样的数据,所以这个Dynamic Batching方法在我们这个任务中还没有用武之地

Batch版块:主要就是负责封装各个需要的数据是“加工者”。在模型训练中,我们需要全部数据是:Encoder所用的source、source_mask;Decoder所用的target、target_y、target_mask。其中target是Decoder的输入,而target_y是正确答案会在Decoder之后计算损失值时使用,target_mask相当于source_mask与subsequent_mask所取的一个并集,作用在target上面用于Decoder预测输出阶段

所以在Batch版块我们的Goal是:封装零散的数据,形成一个大的Batch抽象类,然后之后可以直接用这个Batch类进行数据分发

知道了Data版块中各个小板块的功能作用,我们现在对其依次代码实现:

  • data_generator模块:

三个传入的参数分别表示:词汇表大小为V,一个Batch里面要包含batch_size条数据,一个Epoch总共需要传入nbatches个Batch的数据量

于是,用循环遍历的时相当于生成数据是一个Batch一个Batch地生成,所以一共循环nbatches次;由于一个Batch由若干条句子序列组成,所以一个Batch的数据就可以看作是一个数字矩阵,而且是“类独热编码”之后得到的数字矩阵,所以里面的数字取值范围应当为1~(V-1)(注意:Copy Task中每条生成的数据都等长所以不会有PAD = 0编码);第一个编码设置成1(因为是BOS起始符)

使用yield而非return可节省内存,流水线吐出数据,最后data_generator吐出的数据是source和target打包形成的Batch对象

  • 结束了data_generator,现在我们来快速过一下第2个Dynamic Batching版块(在接下来的简单Copy Task中使用不到),即一个batch_size_fn类

回顾一下,之前我们说过,由于输入数据长度不一,长短句混杂会导致大量的 PAD 填充。虽然 Mask 能屏蔽这些 PAD 对模型结果的影响,但它们依然会消耗 GPU 算力和显存空间。Dynamic Batching 的目标是: 动态调整每个 Batch 里的句子数量(batch_size),确保每个批次产生的矩阵总面积(batch_size × 当前批次最大序列长度)维持在一个恒定且接近显存上限的水平。这样可以实现:短句多塞点,长句少塞点,从而最大限度减少 Padding 带来的计算浪费,并提升训练效率。

上述代码先定义了2个全局变量max_source_in_batch, max_target_in_batch分别表示 当前batch中源语言序列(原文)的最大长度 和 当前batch中目标语言序列(译文)的最大长度(含起止符)。需要这两个变量的原因出于PAD填充是以填充至矩形为标准,所以这个矩形的长一定由当前batch里面的“最大长度”所决定

主要的3个参数分别为:new:准备加入batch的下一个样本对象;count:如果这一条加入,batch里的总句数(正在更新中的batch_size);sofar:截至目前已经累计的batch尺寸(类似于内存)。

我们是把样本数据一个一个扔进去考虑,所以每一波batch的开头会有清零操作:如果是当前batch里面的第1条样本,就把两个全局变量进行重置为0(清空上一个batch留下的‘最长纪录’),之后不断更新最大值,最后返回Encoder与Decoder所用数据占用内存的最大值如果我们预先设置一个内存上限,那么这样能够保证不超越上限又能充分压榨显存。这样操作之后,每一个批次的batch_size都会不同,但是总内存量却保持在一个接近的水平

  • Batch版块

前两个版块我们用data_generator产出了数据,用batch_size_fn知道了打包数据的最大条数,现在就只差Batch类进行封装即可,并且再把Batch类接入到data_generator里面使用就可以使用数据产出+封装的流程

由于data_generator中产出的是source和target封装起来的Batch对象,现在的Batch类就是告诉我们这个Batch类收到source和target之后另外构造出一个批次中模型训练所要用到的source_mask、target_y、target_mask,并把这5个数据全部封装起来便于之后把数据分发给Encoder和Decoder

回到data_generator,产出的source和target是两个完整编码序列,从<BOS>代表的1开始一直有10位。我们知道:模型永远预测模型输入的后一位,即Decoder的输入target与正确答案之间具有一步的延后性,故有:

这也是为什么原论文中模型架构图中:

特别标注了对于Decoder的输入进行了 Shifted Right的处理

接下来我们来看make_std_mask这个函数的构建:

@staticmethod 表示把下面这个函数设置为一个静态函数,这样这个函数虽然处在Batch这个大类里面,但是传入的参数没有self;而如果没有使用这个方法,那么该函数会被”私有化“,如果要使用这个函数就必须先将Batch类实例化为b,然后才能调用:b.make_std_mask( )

注意:在函数静态化之后,这个函数不会与Batch大类进行绑定,也就是给它传入参数target和pad时它只会被动接收这两个参数,并不知道这两个参数是不是Batch里面的成员变量。所以我们在 self.target_mask 处调用这个函数并传参时用的是处理之后(切掉最后一行)的target,并最终返回一个 滤PAD(类似于source_mask)和防偷看的subsequent_mask的并集

除了一个批次训练过程中所需的5大数据的封装之外,Batch类新增了一个ntokens,即算出整个批次里面所有真实标签的总个数,之后会用于计算每个标签的平均损失值

30%

回顾一下我们最开始提出的整个流程图:

Data ✓
Model
Criterion
Loop
Optimizer (Update ↺)

我们上述已经完成了Data这一版块,又因为Ep.1中我们已然完成了Model板块的建构:

Data ✓
Model ✓
Criterion
Loop
Optimizer (Update ↺)

下面进入损失函数(criterion/loss function)阶段

我们知道,Decoder最终的输出应该是一个概率预测分布,而一般而言,这个输出是“非黑即白”的,即:Decoder仅会把其交叉注意力观察到的最大可能性的词汇标识为概率 = 1,而其余词汇在最终的概率分布输出中概率值均 = 0,这会导致一个问题:Decoder处理模糊问题能力较低,模型易过拟合

解决措施,使用标签平滑化(Label Smoothing),即让模型对于其预测的信心不为100%,而将剩余的信心值均分给其余词汇。这是避免模型过拟合的正则化方法中的一种。

在上述代码中的具体实施:实际思路是我们传入期望设置的超参数smoothing,首先在内存中开辟一块与输出x形状相同的矩阵(batch_size, vocab_size),随后先在所有位置上填充一个微小值(由于除去正确答案为confidence和PAD处的0,所以还剩vocab_size – 2个空,微小值 = smoothing/(vocab_size – 2)),把正确答案处覆盖成confidence(= 1 – smoothing),把PAD处覆盖为0,这样就生成了一个“标准答案”的概率分布矩阵(即Target)。

接着我们把这个目标概率分布(P(x))与经过LogSoftmax之后的输出x(Log Q(x))概率分布代入KL散度公式:

DKL(PQ)=P(x)logP(x)Q(x)=P(x)(logP(x)logQ(x))D_{KL}(P \parallel Q) = \sum P(x) \log \frac{P(x)}{Q(x)} = \sum P(x) (\log P(x) – \log Q(x))

最终散度公式得到的值就表示用预测分布P(x)去表示目标正确分布Q(x)时所得的损失值(信息损失)

数学详情可见:https://www.countbayesie.com/blog/2017/5/9/kullback-leibler-divergence-explained

40%

上述给出了在我们任务中的损失值计算标准,我们知道,Loop循环阶段肯定是要把当前训练轮次内要做的所有事情全部打包,那么肯定会涉及数据先打包,然后传送给model,model向前传播进行此刻的损失计算和参数更新,最后会有一个类似于日志打印的结果

  • 既然是几个离散的过程拼接在一起,那么我们不妨先来实现里面的一个小步骤:构建一个优化器,其主要作用就是计算梯度、更新权重参数

在Transformer模型中我们使用经典的Adam优化器内核:https://arxiv.org/abs/1412.6980

其中Adam的核心公式为:

Update=+ϵ×Learning_RateUpdate = \frac{一阶矩(方向建议)}{\sqrt{二阶矩(波动阻力)} + \epsilon} \times Learning\_Rate

分子:往哪走?(顺着惯性走)
分母:走多快?(波动大的走慢点,波动小的走快点)

虽然Adam优化器可以自适应地调整权重更新的快慢和学习率大小,但是在Transformer训练初期,由于层数过多,导致开始时反向传播得到的梯度过大,而此时的Adam优化器还没有建立好稳定的统计量(二阶矩趋于0)(波动阻力),将会导致参数优化过快

Noam的核心公式是:

lr=dmodel0.5min(step0.5,stepwarmup1.5)lr = d_{\text{model}}^{-0.5} \cdot \min(\text{step}^{-0.5}, \text{step} \cdot \text{warmup}^{-1.5})

其中warmup是我们人为设置的热身步数,step是指我们优化到目前进行到了多少步,这样min( )式里面左边随step衰减,右边随step线性递增,取min之后会对学习率learningrate进行限制大小,使学习率像爬山一样:先慢慢爬升(预热),然后再缓缓下降(衰减),并在 step = warmup 的时候取得最大值

整体优化器代码实现如下:

再额外阐释一下Noam调度策略与Adam优化器二者是怎样形成结合体的:

最终每个权重参数更新的步长公式可以这样表示:

最终位移=RateNoam全局限速×mtvt+ϵAdam 的个体微调\text{最终位移} = \underbrace{\text{Rate}_{\text{Noam}}}_{\text{全局限速}} \times \underbrace{\frac{m_t}{\sqrt{v_t} + \epsilon}}_{\text{Adam 的个体微调}}

所以Noam修复了Adam起步时候的风险,而Adam在中期则可以根据自身优势自适应地更新模型的各个参数

最终我们建立这样的完整版优化器:

举个简单例子实例化一下:

我们把这三个优化器每一个step处得到的learning rate画出来对比一下:

最终可视化呈现为:

50%

在我们编写Loop之前,我们先思考一下,我们一般执行的循环Loop一般来讲跟Optimizer不应该是两个呈先后顺序的步骤,相反,我们通常希望在Loop里面把模型训练一个Epoch中的所有操作全部囊括,所以,我们之前虽然独立地解决了Criterion(LabelSmoothing + KL散度) 和 Optimizer(Noam + Adam),我们希望有一个步骤能够直接调用这两个类,并实现从 “计算Loss ——> 参数更新” 的一个连贯操作

  • SimpleLossCompute抽象类进行上述这个连贯的操作

在上述代码中我们看到SimpleLossCompute的实现思路:首先是计算出预测对数概率分布,将其输入给LabelSmoothing,然后调用criterion计算损失loss并计算出相应梯度,再调用Optimizer进行参数更新优化,最终返回的是这一个Batch的总损失。

等等,为什么首先一步是计算处预测对数概率分布? 回顾一下,在make_model的最后一行我们看到:

而之前我们在定义Generator类的时候有:

这表明Generator类的实例化的forward函数具备输出对数概率分布的功能!

看到这里你更疑惑了,那既然model里面有Generator这一层,那么我们在run_epoch里面如果调用model的时候应该就会得到对数概率分布啊?那为啥调用model之后还要在SimpleLossCompute里面再算一次???

事实上,run_epoch里面我们将只调用model.forward( )函数,而回到EncoderDecoder的forward( ):

可以看到,forward函数里面根本不会用到Generator这个大类,也就是说在run_epoch中我们第一步使用model只是调用了它的forward函数,其最终返回的是 Decoder 最后一层输出的特征向量(Hidden States),形状通常为(batch_size, seq_Len, d_model),并没有得到最终的对数概率分布!!!

所以,我们在SimpleLossCompute中第一步是收到Decoder输出的特征向量之后先使用generator得到对数概率分布输出,再进行后续处理(送入KL散度公式时PyTorch要求预测分布必须取对数后再传入,目标正确分布则原封不动地传入)…

  • 补充: 既然model调用forward函数的时候没有用到Generator这个大类,那为什么在make_model的时候要将Generator这个大类囊括进去?!

这里暂时地要用到后续实例化优化器时的做法:

实例化优化器的时候我们会传入model.parameters( ),即模型全部的参数,在SimpleLossCompute中,当我们使用:

会从loss最后一步开始,反向传播,依次经过EncoderDecoder中的所有层,算出其梯度,并进一步使用优化器进行参数优化更新。

现在清晰了:ANS:在make_model中我们囊括Generator的大类是为了将其参数包含于model里面,这样model全部参数可以传入优化器,Generator的参数也就顺带传入了优化器,并在后续优化更新的时候参数可以得到更新

最后思考一个问题:为什么Generator里面的参数那么重要?为啥要更新它?它的作用是什么?

可以回Ep.1文章复习一下,因为它的作用是 将Decoder的输出转换为词汇表上的对数概率分布(Linear + LogSoftMax) ,在run_epoch里面,我们实际上是又把它拎出来再作为参数传到run_epoch中进行单独的一步处理操作。

60%

解决了计算损失、更新参数离散步骤的打包为SimpleLossCompute,回顾下我们的进度:

Data ✓
Model ✓
Criterion ✓
Loop …
Optimizer ✓ (Update ↺)

前面说到,Loop阶段的主要工作是首先将封装好的Batch类型数据分发给Encoder与Decoder,然后在创建好的model中向前传播得到预测分布,将预测分布输入criterion算出损失值,最后在optimizer中进行反向传播与参数更新

综上所述,Loop的主要作用是串联各个打包好步骤,并组装成一套完整的操作流程

我们看到在核心for循环中我们先根据单独拎出来传入的generator实例得到模型的预测概率分布输出out,随即计算损失loss,然后根据每个batch里面的ntokens数计算并打印日志:一轮Batch中的平均损失值。

这就是run_epoch要做的所有事情。

65%

讲了这么久,进度拉满!

Data ✓
Model ✓
Criterion ✓
Loop ✓
Optimizer ✓ (Update ↺)

相信如果看到这里的诸位已经磨刀霍霍了,话不多说,模型先来跑一跑:

如果输出cuda说明GPU配置好了,用它跑这个训练快一些,没有的话只用CPU也不要紧

执行后会输出日志,输出样例:

大概30s跑完模型训练之后,我们期望用我们训练的模型做一个交互式程序,看看能否正确完成Copy Task。现在的情形是:我们手上有完整的Transformer模型架构,也有一套它训练了30个Epoch后得到的一套权重参数,而我们的期望是我们自己设定一个seq_Len = 10 (因为我们之前训练时设置的V = 11,里面还包含<PAD> = 0)的数字序列为输入,让模型能够完成这个Copy Task

我们知道,在SimpleLossCompute中我们第一步的操作是根据Decoder输出的特征向量得到了模型的对数概率预测分布,随后我们再将这个概率分布与真实分布对比得到损失。而在现在的任务中,模型需要预测得到这个真实分布。

自然地,我们可以认定模型输出的就是真实分布,这就是模型最终的预测结果。

由于generator处理(batch_size, seq_Len, d_model)得到的是一个概率分布(batch_size, seq_Len, vocab_size),准确来讲这算是一个二维矩阵,因为我们预测任务中给模型的输入只有一个,所以batch_size = 1。假设我们规定生成的序列最大长度为10,这个矩阵大概长成(这个地方我乱画的):

Sequence Generation Probability Matrix
Idx Position 0 (EOS)12 345 678 910
1-8.42-5.88-7.10-6.34-6.15-0.02-6.78-7.21-6.95-7.05-6.34
2-7.51-4.32-0.35-4.12-4.55-4.28-4.66-4.38-4.47-4.19-4.41
3-9.88-6.56-7.12-6.98-7.21-6.77-7.01-6.95-6.92-0.01-6.82
4-8.21-0.09-5.42-5.05-5.18-5.32-5.25-5.38-5.10-5.24-5.14
5-7.34-5.12-5.67-5.28-5.45-5.20-5.51-5.62-0.07-5.48-5.38
6-6.21-3.05-3.45-3.32-0.65-3.11-3.28-3.52-3.15-3.25-3.19
7-8.05-5.67-5.92-5.75-5.82-5.70-5.95-0.15-5.96-5.90-5.85
8-7.77-4.55-4.82-0.12-4.58-4.62-4.85-4.92-4.71-4.80-4.74
9-9.12-4.98-5.25-5.05-5.18-5.01-0.04-5.32-5.11-5.21-5.15
10-0.03-6.12-6.54-6.32-6.48-6.25-6.61-6.70-6.42-6.52-6.38

其中绿色荧光部分就是模型在当前位置认为最有可能的输出(对数概率值最大)

那么我们可以使用一种“贪心解码”的方式,让模型在每一步都取概率最大的数字即可:

78%

现在模型和解码器逻辑都已经准备就绪,现在手动初始化一个source序列,然后让model跑一遍

得到输出:

80%

我们现在其实可以来看一下模型这一步是怎么预测出这个序列的,我们可以把generator处理后的概率分布调出来看一下:

虽然在上述过程中我们实际画图并不会用到yielded_sequence,但由于循环只有在上一个确定的答案后继续进行,所以里面有在不断延长yielded_sequence

最后对应的图像为:

图中每一行深色的地方就是模型在每一步的预测输出

85%

最后我们来实现注意力可视化:

  • 交叉注意力可视化

上面中这行代码是关键:

表示深入到decoder的第一个字曾中的cross_attention小块中取出其attention权重矩阵的第1个样本,也就是4个头对于输入这条句子每个位置预测值的观测位置

由于:

Attention=Softmax(QKTdk)Attention = \text{Softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)

所以 Attention矩阵的形状为:

(Batch_Size,Heads,Query_Len,Key_Len)(Batch\_Size, Heads, Query\_Len, Key\_Len)

在Copy Task中 Query_Len = source_Len = “seq_Len” = target_Len = Key_Len (此题中是10)

在预测时batch_size = 1,但还是需要一个索引[0]把四个Head的所有观测数据全部取出来,之后再在这一步再用一个索引分别取4个Head的数据:

92%
  • Encoder中的自注意力可视化

自注意力可视化与交叉注意力可视化大同小异:
①不需要运行greedy_decode因为我们此处只是可视化Encoder的自注意力,不涉及Decoder
②自注意力访问Encoder里面的self_attention;交叉注意力访问的是Decoder里面的cross_attention

进一步地,可以简单解释一下上述两幅可视化热力图的具体含义:

  • 交叉注意力图: 纵轴(Output Sequence):代表“正在生成的词”。每一行代表模型在生成第i个数字时的状态。横轴(Input Sequence):代表“原文中的词”。每一列代表输入序列中的第 j个数字。
    热力图的一行,描述了模型为了写出当前这个输出词,分别从原文的各个位置“搬运”了多少信息。深色表示“强关联”
  • 自注意力图:纵轴(Input Sequence):代表序列中的每一个词(作为查询者 Query),横轴(Input Sequence)代表序列中的每一个词(作为被查询者 Key),由于是自注意力,纵轴和横轴的标签完全一样。纵轴(Query)在观察横轴(Key),所以可以先锁定一个y坐标,表示“正站在y的视角进行观察”
99.9%

到此为止,基础的Transformer模型架构及一个简单的Copy Task实现已经完全结束!

更多资料可参考如下(不全)(在本文中虽并未参考,但读者可另行翻阅):

https://medium.com/@mayanksultania/transformers-101-tokens-attention-and-beyond-b080a900ca6c
https://jalammar.github.io/illustrated-transformer/
https://www.ruder.io/optimizing-gradient-descent/

2026.3.4晚

暂无评论

发送评论 编辑评论


				
|´・ω・)ノ
ヾ(≧∇≦*)ゝ
(☆ω☆)
(╯‵□′)╯︵┴─┴
 ̄﹃ ̄
(/ω\)
∠( ᐛ 」∠)_
(๑•̀ㅁ•́ฅ)
→_→
୧(๑•̀⌄•́๑)૭
٩(ˊᗜˋ*)و
(ノ°ο°)ノ
(´இ皿இ`)
⌇●﹏●⌇
(ฅ´ω`ฅ)
(╯°A°)╯︵○○○
φ( ̄∇ ̄o)
ヾ(´・ ・`。)ノ"
( ง ᵒ̌皿ᵒ̌)ง⁼³₌₃
(ó﹏ò。)
Σ(っ °Д °;)っ
( ,,´・ω・)ノ"(´っω・`。)
╮(╯▽╰)╭
o(*////▽////*)q
>﹏<
( ๑´•ω•) "(ㆆᴗㆆ)
😂
😀
😅
😊
🙂
🙃
😌
😍
😘
😜
😝
😏
😒
🙄
😳
😡
😔
😫
😱
😭
💩
👻
🙌
🖕
👍
👫
👬
👭
🌚
🌝
🙈
💊
😶
🙏
🍦
🍉
😣
Source: github.com/k4yt3x/flowerhd
颜文字
Emoji
小恐龙
花!
上一篇