说起《This post is all you need——层层拨开Transformer》这份文档我想大家一定不会陌生，不止是你们常读常新，我自己读也是常读常新，好多读者甚至还把它说成是本人的成名作（不过好像也没啥名），毕竟当时花了半年多时间来写这篇文章。
但是成不成名那是后话，把这件事做好才是最重要的！
关键字：Transformer、注意力掩码、Attention Mask、Key Padding Mask
1. 背景#
前两天，有一个读者留言说 Transformer 中 Decoder 里的 tgt key padding mask 操作其实是不需要的。我看了以后回复到是需要的，你再好好想一想如果还是没想明白明天再来找我（主要是想让他自己思考一下，有利于自己理解），我告诉你。
第二天，这位读者如约而至，给我发了他画的图示等等。我又想了一下，觉得依旧是必要的。看在对方如此认真地询问这个问题，我就打算写篇文章顺便录个视频讲解一下这里的门道。
结果这下到好，当我开始写作这篇文章时，顺着计算逻辑一步一步发现，他说的好像是对的。但此时并不敢立马做出判断，因为这篇文章太多人看，得要为自己说出去的话负责！
于是，我又开始从头回顾整个 Transformer 的前向传播和损失计算过着，通过图示来进行梳理，同时也通过代码来进行了验证。
最终发现……
2. Transformer 计算过程回顾#
下面，就来带着大家一探究竟，到底整个 mask 操作是否真的不需要！
首先，我们先来通过一张图回顾一下 Transformer 中3个主要多头注意力部分的计算过程。
如图1所示便是整个 Transformer 在训练时的计算过程图，从左到右依次是 Encoder 中的多头注意力计算过程、Decoder 中的掩码多头注意力计算过程和 Decoder 中与 Encoder 交互的交叉多头注意力计算过程。
整个过程大家对着图示自行温习这里就不在赘述了。
我们都知道，整个 Transformer 结构中涉及到的 “掩码” 类型一共有两种：① 用于区分同一个 batch 中不同长度序列是否被填充的 key padding mask；② 在训练时，Decoder 中用于模仿推理过程中在编码当前时禁止看到未来信息的 attention mask（也叫做 casual mask 或 future mask）。
这里需要提醒一点的是，不同的代码实现或框架对这两种“掩码”的称呼不一样，有的地方把填充掩码也成为 attention mask，而把我们这里指代的attention mask 称为 casual mask，大家在阅读相关材料时需要根据上下文理解。
对于 key padding mask ，它在 Transformer 中有3个地方被用到，也就是图1中横着看的第二大块；而 attention mask 则只会被用在解码器中的 Masked-multi Head Attention 中。
那现在问题来，Decoder 中的 tgt key padding mask 是否真是可以不用呢？
3. padding mask 真可以不用吗？#
答：真的可以不用。
这个结论听起来好像真的有点难以相信，毕竟好像大家都用了。
首先，对于自回归模型来说，我们在训练阶段为了加速模型的训练过程，都是直接使用小批量样本来进行训练，而这就导致了不同样本长度不能进行计算的问题，因而要使用填充操作来使同一个 batch 中的序列长度保持一致。同时，在自注意力机制中，为了防止注意力被分配到 padding 位置上，所以我们需要通过 padding mask 来忽略掉这部分注意力值。
本来一切的合情合理，但这个时候 Decoder 中偏偏又出现了 Attention mask，如图2所示。
在图2中，由于 Attention mask 的作用，在得到 Query 中 “who am I P P P” 这个序列的编码结果时：对于 “who” 来说模型只会讲注意力分配到 V 中的第一行上；对于 “am” 来说模型只会将注意力分配到前面两行上；对于 “I” 来说模型只会将注意力分配到前面三行上。
下面关键点来了，敲黑板！！！！
此时我们可以发现，在对前面三个 token 依次进行编码时，由于有 Attention Mask 的存在，所以不需要tgt key padding mask 同样能够保证注意力部分配分到 padding 位置上。
进一步，对于 Query 中后三个位置的编码来说：
① 如果没有 tgt key padding mask 的存在，那么模型在对这三个位置的 token 进行编码时，会分别将注意力分配到 V 中的前四行、前五行和前六行上，换句话说注意力会被分配到 padding 位置上。
② 如果有tgt key padding mask 的存在，那么模型在对这三个位置的 token 进行编码时，都只会将注意力分配 V 中的前三行上，也就是注意力不会被分配到 padding 位置上。
到此时，我们依旧能够得出加与不加 tgt key padding mask 得到的编码结果 Query 的确是不一样的，但是，
你 Query 中的后三行本来就是被标记为填充位置的，你被编码成什么样子、是对的还是错的我们关心吗？
我们不 Care 呀！我们又不会用到你，我管你被编码成什么样了。
不 Care，真的不 Care !
那从哪里可以体现出它的的确确没有被使用到呢？
答：在最后计算损失的时候。
如图3所示，模型根据 Query 和 Encoder 的输出 Memory 交互后得到了输出 Output，再经过一个分类层得到了 y_pred。我们可以看到，如果根据 logits 我们将会得到 “no am who you are .” 这样一个结果（真实的标签为 “who am I P P P”，即 “4 2 5 0 0 0”）。
由于真实标签的最后三个 token 是填充的，所以我们在计算损失的使用会通过指定 ignore_index 来指定忽略的位置。因此，对于上述预测结果来说在计算损失时只会考虑 “no am who” 对应的损失，所以后三行被编码成什么结果（是否考虑到注意力分配到了 padding 位置上），都不会影响模型的训练过程。
$$ \text{loss} = -(\log0.12+\log0.22+\log0.31) = 4.806 $$到此，我们总算说清楚了在训练阶段原来真的可以不用 tgt key padding mask.
那推理阶段呢？当然更不需要啦，都没有 padding 这回事。
4. 编码验证#
上面我们从原理上分析了训练阶段可以不用 tgt key padding mask 这一结论，下面我们再来通过实际代码计算一下，看看两种情况下的损失是否相等。
1 import torch
2 import torch.nn as nn
3
4 torch.random.manual_seed(42)
5
6 d_model = 7
7 tgt_len = 6
8 src_len = 5
9
10 # =================== Encoder ===================
11 src_em_input = torch.randn(src_len, d_model)
12 q_proj0 = nn.Linear(d_model, d_model)
13 k_proj0 = nn.Linear(d_model, d_model)
14 v_proj0 = nn.Linear(d_model, d_model)
15 q, k, v = q_proj0(src_em_input), k_proj0(src_em_input), v_proj0(src_em_input)
16 src_key_padding_mask = torch.tensor([[False, False, False, True, True]])
17 weight = q @ k.transpose(0, 1) # [src_len,src_len]
18 print("Encoder attention weight before key padding:")
19 print(weight)
20 weight = weight.masked_fill(src_key_padding_mask, float('-inf'))
21 print("Encoder attention weight after key padding:")
22 print(weight)
23 weight = torch.softmax(weight, dim=-1) # [tgt_len,tgt_len]
24 memory = weight @ v # [src_len,d_model]
25
26 # =================== Decoder ===================
27 ## ------------------ Masked Attention
28 tgt_em_input = torch.randn(tgt_len, d_model)
29 tgt_key_padding_mask = torch.tensor([[False, False, False, True, True, True]])
30 attn_mask = (torch.triu(torch.ones((tgt_len, tgt_len))) == 1).transpose(0, 1)
31 attn_mask = attn_mask.float().masked_fill(attn_mask == 0, float('-inf')).masked_fill(attn_mask == 1, float(0.))
32 q_proj1 = nn.Linear(d_model, d_model)
33 k_proj1 = nn.Linear(d_model, d_model)
34 v_proj1 = nn.Linear(d_model, d_model)
35 q, k, v = q_proj1(tgt_em_input), k_proj1(tgt_em_input), v_proj1(tgt_em_input)
36 # [tgt_len, d_model]
37 weight = q @ k.transpose(0, 1) # [tgt_len,tgt_len]
38 weight += attn_mask
39 print("Decoder masked attention weight before key padding:")
40 print(weight)
41 # weight = weight.masked_fill(tgt_key_padding_mask, float('-inf')) #
42 # print("Decoder masked attention weight after key padding:")
43 print(weight)
44 weight = torch.softmax(weight, dim=-1) # [tgt_len,tgt_len]
45 output = weight @ v # [tgt_len,tgt_len] @ [tgt_len, d_model]
46
47 ## ------------------ Cross Attention
48 q_proj2 = nn.Linear(d_model, d_model)
49 k_proj2 = nn.Linear(d_model, d_model)
50 v_proj2 = nn.Linear(d_model, d_model)
51 q, k, v = q_proj2(output), k_proj2(memory), v_proj2(memory)
52 # q: [tgt_len, d_model], k,v: [src_len, d_model]
53 weight = q @ k.transpose(0, 1)
54 # [tgt_len, d_model] @ [d_model,src_len] = [tgt_len,src_len]
55 weight = weight.masked_fill(src_key_padding_mask, float('-inf'))
56 weight = torch.softmax(weight, dim=-1)
57 output = weight @ v # [tgt_len,src_len] @ [src_len,d_model] = [tgt_len,d_model]
58
59 # =================== Loss ===================
60
61 tgt = torch.tensor([1, 2, 3, 0, 0, 0], dtype=torch.long)
62 loss_fcn = torch.nn.CrossEntropyLoss(ignore_index=0)
63 loss = loss_fcn(output, tgt)
64 print(loss)打开是否需要tgt_key_padding_mask的开关在第41，42行，其对应的 weight 输出结果分别为：
tensor([[-1.2846, -inf, -inf, -inf, -inf, -inf],
[-0.0381, -0.1050, -inf, -inf, -inf, -inf],
[ 0.1754, -0.1007, 0.4034, -inf, -inf, -inf],
[ 0.1751, 0.0842, 0.0198, -0.4393, -inf, -inf],
[-1.5436, -0.4939, -1.5253, -0.9625, -0.4269, -inf],
[-0.3050, -0.0124, -0.0216, -0.5008, -1.5813, -0.0728]],和
tensor([[-1.2846, -inf, -inf, -inf, -inf, -inf],
[-0.0381, -0.1050, -inf, -inf, -inf, -inf],
[ 0.1754, -0.1007, 0.4034, -inf, -inf, -inf],
[ 0.1751, 0.0842, 0.0198, -inf, -inf, -inf],
[-1.5436, -0.4939, -1.5253, -inf, -inf, -inf],
[-0.3050, -0.0124, -0.0216, -inf, -inf, -inf]],
grad_fn=<MaskedFillBackward0>)两种情况下对应的损失值均为 tensor(1.8582, grad_fn=<NllLossBackward0>)，具体分析各位读者自行运行代码，或看文末视频。
关于这件事还想分享的一点是，我们每个人在表达当下观点时都是基于截至那一刻自己对于这件事情的认知，至少在这一刻我们认为它是对的。如果某一天，当有人质疑你先前给出的定论，那也没关系，顺着这个逻辑再去理解就是，如果真是那样纠正就是。