回复
The Annotated Transformer注释加量版，读懂代码就真的懂了Transformer 原创
本文是在The Annotated Transformer这篇文章基础上的二次加工。
1.给代码加了更详细的注释。
2.输出详细日志跟踪数据。
原文地址：https://nlp.seas.harvard.edu/annotated-transformer/
或者后台回复taf获取pdf下载链接。
The Andnotated Transformer
Attention is All You Need
- v2022: Austin Huang, Suraj Subramanian, Jonathan Sum, Khalid Almubarak, and Stella Biderman.
- Original: Sasha Rush
阅读方法
由于原文内容过长，我没有把原文拷贝过来，阅读本文时，请打开原文链接或者我添加注释的notebook。
1、给代码加了更详细的注释。
原文基于pytorch从0开始复现了transformer模型，我在原文代码基础上追加了更详细的注释，代码可以在下面链接找到。https://github.com/AIDajiangtang/annotated-transformer/blob/master/AnnotatedTransformer_comment.ipynb
另外，我还在模型结构上加了注释，我将代码中重要的类名或者函数名标注在Transforner结构的图片上，阅读代码时请结合图片上的名称，这样有助于快速理解代码。
2、输出日志跟踪数据。
原文提供了一个训练德译英模型的代码，我在此基础上加了一些日志，打印数据的维度来辅助对Transformer的理解。
我将按照图片上标注数字顺序来跟踪数据。
原始论文中，Transformer是一种Encoder-Decoder架构，左边是Encoder，用于提取源语言的表征，右边是Decoder，根据表征结合目标语言语法生成目标语言。
先从Encoder这边开始。
0、Inputs：
假设batch size为2，所以每个batch包含两个样本，每个样本由（德语，英语）文本对组成。
[
('Eine große Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.', 'A large group of young adults are crammed into an area for entertainment.'),
('Zwei Arbeiter stellen Laternen auf.', 'Two workers working on putting up lanterns.')
]（batch size的意义：模型每次都是基于batch size个样本的损失来更新参数，batch size需要根据内存，显存大小确定）
对于Encoder而言，它只需要源语言，也就是德语。
'Eine große Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'
'Zwei Arbeiter stellen Laternen auf'
1、Embedding：
1.1.先将文本转换成tokens，并添加起始和结束符token。
(load_tokenizers函数,
tokenize函数,
build_vocabulary函数}
["<s>", "</s>", "<blank>", "<unk>"]起始符token id：0，结束符token id：1，padding token id：2
'Eine große Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'的tokens如下
torch.Size([11])
tensor([ 0, 14, 176, 38, 683, 7, 6, 116, 7147, 4, 1],
device='cuda:0')（通过结果看是基于词的tokenization方法）
1.2.因为文本长度不一致，通过padding的方式将序列长度统一为72。
{collate_batch函数}
（padding不是必须的，只是出于方便和效率考虑，72是个经验值，通过对训练数据的统计得出）
'Eine große Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'padding后的tokens如下
torch.Size([72])
[tensor([ 0, 14, 176, 38, 683, 7, 6, 116, 7147, 4, 1, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]一个batch下有两个样本，对另一个样本的德语进行同样的转换最终得到编码器输入：X，维度[2, 72]。
在训练过程中，无论是计算注意力还是交叉注意力，每个样本是相互独立的，所以可以将一个batch下所有数据组织成矩阵的形式输入到模型进行并行计算。
1.3.最后将上一步的tokens通过一个Embedding线性层转换成词嵌入，设置d_model=512，所以词嵌入维度为512。
{Embeddings类}
Embedding层输入就是前面的X；维度是torch.Size([2, 72])。
Embedding层的输出维度是torch.Size([2, 72，512])，也就是每个token id都被转换成512维的向量。
tensor([[[-0.6267, -0.0099, 0.3444, ..., 0.5949, -0.4107, -0.6037],
[ 0.4183, -0.1788, -0.3128, ..., 0.5363, -0.5519, 0.4621],
[ 0.4645, -0.2748, -0.4109, ..., -0.6270, 0.4595, -0.4259],
...,
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066]],
[[-0.6267, -0.0099, 0.3444, ..., 0.5949, -0.4107, -0.6037],
[-0.2121, 0.4323, -0.0869, ..., 0.1337, -0.2679, -0.4689],
[ 0.0751, -0.1048, -0.1263, ..., -0.5541, -0.4463, 0.5209],
...,
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066],
[-0.1489, 0.6431, -0.0301, ..., -0.0163, 0.4261, 0.3066]]],
device='cuda:0', grad_fn=<MulBackward0>)（Embedding过程相当于用512个属性值表示单词的语义信息，经过每个EncoderLayner时属性值会被修改，使其充分吸收上下文信息，属性越多，能表示的语音信息越丰富，但计算量和参数也会增加）
2、PositionalEncoding
{PositionalEncoding类}
在计算注意力分数时，如果调整单词的位置，注意力的输出结果不变，也就是自注意力这种计算方式没有考虑单词的位置信息。
所以需要通过一个额外的位置编码，位置编码与词嵌入维度相同，也是512维向量，最后与词嵌入相加。
前面Embedding层输出维度torch.Size([2, 72, 512])，将其与位置编码相加，输出也是torch.Size([2, 72, 512])。
（位置编码可以通过训练方法得到，也可以采用固定计算方式，本例采用固定计算方式）
所有样本共用同一个位置编码，本例序列长度为72，可以提前计算好位置编码备用。
pos表示位置，第一个词位置是0，第二个词位置是1....本例中就是0-71。
对于512维向量，偶数位置和奇数位置的值分别用上面两个公式计算。
tensor([[[ 0.0000e+00, 1.0000e+00, 0.0000e+00, ..., 1.0000e+00,
0.0000e+00, 1.0000e+00],
[ 8.4147e-01, 5.4030e-01, 8.2186e-01, ..., 1.0000e+00,
1.0366e-04, 1.0000e+00],
[ 9.0930e-01, -4.1615e-01, 9.3641e-01, ..., 1.0000e+00,
2.0733e-04, 1.0000e+00],
...,
[-8.9793e-01, 4.4014e-01, 3.6763e-01, ..., 9.9997e-01,
7.0490e-03, 9.9998e-01],
[-1.1478e-01, 9.9339e-01, -5.5487e-01, ..., 9.9997e-01,
7.1527e-03, 9.9997e-01],
[ 7.7389e-01, 6.3332e-01, -9.9984e-01, ..., 9.9997e-01,
7.2564e-03, 9.9997e-01]]], device='cuda:0')可视化出来就是下面效果。
(上图每一行都是一个位置编码向量，一共生成50个位置编码，每个位置编码是128维向量，而本例需要生成72个，每个512维）
3.MultiHeadedAttention
{MultiHeadedAttention类,
attention函数}
MultiHeadedAttention类的输入是query, key, value，维度都是torch.Size([2, 72, 512])，其实他们的内容也是一样的，就是上一步输出的Embedding+位置编码。
然后query, key, value分别经过一个独立的线性层，线性层的维度[512, 512]，两个样本的[72, 512]分别与[512, 512]矩阵乘法，所以线性层的输出维度仍是[2, 72, 512]，最后经过reshape和转置将[2, 72, 512]转换成torch.Size([2, 8, 72, 64])，8代表有8个头，其实就是将512转换成了8*64来实现多头注意力机制。
(虽然是8个头，但与一个头的情况相比，参数并没有增加)
接下来计算单个头的注意力，Attention函数的输入query, key, value的维度都是torch.Size([2, 8, 72, 64])，注意力分数矩阵维度torch.Size([2, 8, 72, 72])，输出torch.Size([2, 8, 72, 64])。
最后将多个头的输出拼接在一起，也就是通过reshape和转置将torch.Size([2, 8, 72, 72])转换成[2, 72, 512]，最后经过一个[512, 512]的线性层输出[2, 72, 512]。
4、SublayerConnection
{SublayerConnection类}
将多头注意力的输出经过层归一化和输入进行残差链接，不改变维度，输入输出都是[2, 72, 512]。
5、PositionwiseFeedForward
{PositionwiseFeedForward类}
这其实是一个MLP层，输入维度512，隐藏层维度2048，输出层维度512，也就是2*72个tokens并行与[512, 2048]矩阵乘升维至[2, 72, 2048]，然后再与矩阵[2048，512]乘恢复到原来维度[2, 72, 512]。最后再经过层归一化和残差链接。
6、EncoderLayer
{EncoderLayer类}
将3，4，5重复6次，这里需要注意下，这6个EncoderLayer只是结构一致，但参数是独立的，原始的Embedding经过6个EncoderLayer后维度是不变的，仍然是[2, 72, 512]，只不过内容被改变了。
7、LaynerNorm
{LayerNorm类}
为了计算稳定，整个Encoder的输出会再次经过层归一化处理，然后输入到Decoder层作为key和value，维度仍然是[2, 72, 512]。
Encoder把key和value传递给Decoder，它的使命就算完成了。剩下的就是根据那边的损失等着更新参数了。
让我们来到Decoder这边。
0、Inputs：
[
('Eine große Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.', 'A large group of young adults are crammed into an area for entertainment.'),
('Zwei Arbeiter stellen Laternen auf.', 'Two workers working on putting up lanterns.')
]对于Decoder，除了Encoder的key和value，还要有query，这个query就是目标语言，也就是英语。
'A large group of young adults are crammed into an area for entertainment.'
'Two workers working on putting up lanterns.'
1、Embedding
Decoder和Encoder的Embedding几乎一致，也是先转换成tokens。
'A large group of young adults are crammed into an area for entertainment.'->tokens
torch.Size([16])
tensor([ 0, 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1], device='cuda:0')然后进行padding。
'A large group of young adults are crammed into an area for entertainment.'->padding tokens
torch.Size([72])
[tensor([ 0, 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]对另一个样本进行同样的操作得到编码器的输入Y，维度[2, 72]。
最后将其转换成Embedding，维度是torch.Size([2, 72, 512])。
但有一点需要注意。
Decoder在训练时输入的是整个batch的英语文本，也就是torch.Size([2, 72, 512])。
但在训练过程中预测当前token的输出时，为了让其只能看到当前以及之前位置的输入，避免看到后面的内容，需要采用遮罩的方式，也就是要构造一个mask。
torch.Size([2, 72, 72])
tensor([[[ True, False, False, ..., False, False, False],
[ True, True, False, ..., False, False, False],
[ True, True, True, ..., False, False, False],
...,
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False]],
[[ True, False, False, ..., False, False, False],
[ True, True, False, ..., False, False, False],
[ True, True, True, ..., False, False, False],
...,
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False],
[ True, True, True, ..., False, False, False]]], device='cuda:0')
2、PositionalEncoding
与Encoder一样，输入输出都是[2, 72, 512]
3、MultiHeadedAttention
Decoder中的DecoderLayner有两个MultiHeadedAttention，第一个是Mask MultiHeadedAttention，与Encoder中的计算一致，只不过使用了上一步计算的Mask。
另一个MultiHeadedAttention中的key和value来自Encoder，我们称之为交叉注意力，与自注意力要区分开，query来自前一层的输出，维度都是[2, 72, 512]。
4，5，9，7和Encoder都是一样的。
同样输入Embedding经过6个DecoderLayner后维度不变[2, 72, 512]。
4、Generator
{Generator类}
这其实是一个没有隐藏层的MLP，输入维度512，输出维度vocab，2*72个token的Embedding与矩阵[512,vocab]相乘,输出[2, 72, vocab]，vocab为词表的单词个数，本例中英语单词个数为6291。经过softmax后输出一个概率分布，最大概率对应的位置的词就是模型预测的下一个词。
这样就得到了Decoder的最终输出，输出可以是[2, 72]，里面是英语词表下的id。也可以是[2, 72, vocab]直接输出概率分布，输出形式不同，损失函数也是不同的。
对于其中一个样本，训练过程中Decoder的输入是：
'A large group of young adults are crammed into an area for entertainment.'
torch.Size([72])
[tensor([ 0, 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]如果想更新参数就必须计算损失，计算损失就必须有标签，那标签是什么？
对于Decoder，输入也是输出，标签就是将输入向左移动了一位：
[tensor([ 6, 62, 39, 13, 25, 348, 17, 5318, 71, 28, 179,
55, 4285, 5, 1, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2，2],
device='cuda:0')]也就是起始符0对应的标签是A：6，输入A对应的标签是large：62，Decoder输出维度[2, 72]，标签维度也是[2, 72]，最后通过均方误差计算损失，或者输出概率分布，通过KL损失函数计算损失来更新Decoder和Encoder的参数。
再强调一下，整个batch下所有数据是一起输入到模型的，也就是通过将数据组织成矩阵实现了整个batch的数据并行计算。
训练完成后，就可以用它进行德译英翻译了。
假设输入这么一句德语。
'Eine große Gruppe Jugendlicher in einem kleinen Unterhaltungsbereich.'
德语先经过Encoder进行并行编码，输出[1, 72, 512]作为Decoder的value和key。
在推理过程中就Deocder就不能并行计算了，只能自回归的方式每次前向计算只产生一个token。
刚开始只有一个起始符token 0输入到Deocder。
[tensor([ 0, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]decoder输出6，将6加到0后面再次输入到decoder。
[tensor([ 0, 6, 1, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2,
2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
device='cuda:0')]decoder输出62，以此类推，直到输出终止符token 1。
人的大脑在学习复杂事物时，往往习惯使用一种整体到细节，抽象到具体的渐进的方式。
虽然我在作者的源代码添加了更多的注释和维度信息，但它仍然是细节，为了更好地理解大模型的工作原理，我建议先阅读我之前的图解和动画Transformer系列，以次获得对Transformer有一个高层次的认知。
另外，如果你如果弄明白了Encoder-Decoder架构，那么就能轻松搞懂GPT和BERT了，因为它们一个只用了Encoder，另一个只用了Decoder。
本文转载自公众号人工智能大讲堂
©著作权归作者所有，如需转载，请注明出处，否则将追究法律责任
标签
赞
收藏
回复
相关推荐
