编者按：年初疫情在家期间开始大量阅读NLP领域的经典论文，在学习《Attention Is All You Need》时发现了一位现居日本的台湾数据科学家LeeMeng写的Transformer详解博客，理论讲解+代码实操+动画演示的写作风格，在众多文章中独树一帜，实为新手学习Transformer的上乘资料，在通读以及实操多遍之后，现在将其编辑整理成简体中文分享给大家。由于原文实在太长，为了便于阅读学习，这里将其分为四个部分：
- 透过机器翻译理解Transformer(一)：关于机器翻译
- 透过机器翻译理解Transformer(二)：师傅引进门，修行在个人—建立输入管道
- 透过机器翻译理解Transformer(三)：理解 Transformer 之旅：跟着多维向量去冒险
- 透过机器翻译理解Transformer(四)：打造 Transformer：叠叠乐时间
在涉及代码部分，强烈推荐大家在Google的Colab Notebooks中实际操作一遍，之所以推荐Colab Notebooks是因为1).这里有免费可以使用的GPU资源；2). 可以避免很多安装包出错的问题
本节目录
- 理解 Transformer 之旅：跟着多维向量去冒险
- 5.1 可视化原始句子
- 5.2 可视化 3 维词嵌入张量
- 5.3 Mask：Transformer 的秘密配方
- 5.4 Scaled dot product attention：一种注意函数
- 5.5 直观理解遮罩(mask)在注意力函数中的效果
- 5.6 Multi-head attention：你看你的，我看我的
- 理解 Transformer 之旅：跟着多维向量去冒险
5. 理解 Transformer 之旅：跟着多维向量去冒险
在实现 Transformer 及注意力机制这种高度并行运算的模型时，你将需要一点「空间想像力」，能够想像最高高达 4 维的向量是怎么在 Transformer 的各个元件中被处理与转换的。
如果你跟我一样脑袋并不是那么灵光的话，这可不是一件简单的事情。不过别担心，从这节开始我会把 Transfomer （主要针对注意力机制）里头的矩阵运算过程可视化（visualize）出来，让你在这个多维空间里头也能悠游自在。
就好像一般你在写代码时会追踪某些变量在函数里头的变化，一个直观理解 Transformer 的方法是将几个句子丢入其中，并观察 Transformer 对它们做了些什么转换。
首先让我们建立两个要拿来持续追踪的中英平行句子：
编者按：记住这里的两个例子，他们将贯穿本文中Transformer 的整个实现过程，并通过他们将矩阵运算过程可视化出来
demo_examples = [
("It is important.", "这很重要。"),
("The numbers speak for themselves.", "数字证明了一切。"),
]
pprint(demo_examples)
[('It is important.', '这很重要。'),
('The numbers speak for themselves.', '数字证明了一切。')]
接着利用之前建立数据集的方法将这 2 组中英句子做些预处理并以 Tensor 的方式读出：
batch_size = 2
demo_examples = tf.data.Dataset.from_tensor_slices((
[en for en, _ in demo_examples], [zh for _, zh in demo_examples]
))
# 将两个句子透过之前定义的字典转换成子词的序列（sequence of subwords）
# 并添加 padding token: <pad> 来确保 batch 里的句子有一样长度
demo_dataset = demo_examples.map(tf_encode)\
.padded_batch(batch_size, padded_shapes=([-1], [-1]))
# 取出这个 demo dataset 里唯一一个 batch
inp, tar = next(iter(demo_dataset))
print('inp:', inp)
print('' * 10)
print('tar:', tar)
inp: tf.Tensor(
[[8113 103 9 1066 7903 8114 0 0]
[8113 16 4111 6735 12 2750 7903 8114]], shape=(2, 8), dtype=int64)
tar: tf.Tensor(
[[4205 10 241 86 27 3 4206 0 0 0]
[4205 165 489 398 191 14 7 560 3 4206]], shape=(2, 10), dtype=int64)
上节建立的数据集尸骨未寒，你应该还记得inp shape 里头第一个维度的2 代表着这个batch 有2 个句子，而第二维度的8则代表着句子的长度（单位：子词）；tar则为中文子词序列（subword sequence），不过因为中文我们以汉字为单位作分词，长度一般会比对应的英文句子来的长（shape 中的10）。
2 维矩阵还很容易想像，但我担心等到你进入 3 维空间后就会想放弃人生了。所以还是先让我们用人类比较容易理解的方式来呈现这些数据。
5.1 可视化原始句子
如果我们把这 2 个 batch 用你比较熟悉的方式呈现的话会长这样：
这样清楚多了不是吗？
你可以清楚地看到每个原始句子前后都有 <start> 与 <end>。而为了让同个 batch 里头的序列长度相同，我们在较短的序列后面也补上足够的 0，代表着 <pad>。
这个可视化非常简单，但十分强大。我现在要你记住一些本文会使用的惯例：
- 不管张量（Tensor）变几维，其第一个维度
shape[0]永远代表batch_size，也就代表着句子的数目 - 不同句子我们用不同颜色表示，方便你之后对照这些句子在转换前后的差异
- x 轴（横轴）代表张量的最后一个维度
shape[-1]，以上例来说分别为8和10 - x, y 轴上的标签分别代表倒数两个维度
shape[-2]及shape[-1]其所代表的物理含义，如图中的句子与子词 - 图中张量的
name会对应到代码里头定义的变量名称，方便你对照并理解实现逻辑。我也会秀出张量的 shape 帮助你想像该向量在多维空间的长相。一个简单的例子是：(batch_size, tar_seq_len)
这些准则与信息现在看似多余，但我保证你很快就会需要它们。
5.2 可视化 3 维词嵌入张量
在将索引序列丢入神经网络之前，我们一般会做词嵌入（word embedding），将一个维度为字典大小的高维离散空间「嵌入」到低维的连续空间里头。
让我们为英文与中文分别建立一个词嵌入层并实际对inp及 tar做转换：
# + 2 是因为我们额外加了 <start> 以及 <end> tokens
vocab_size_en = subword_encoder_en.vocab_size + 2
vocab_size_zh = subword_encoder_zh.vocab_size + 2
# 为了方便 demo, 将词汇转换到一个 4 维的词嵌入空间
d_model = 4
embedding_layer_en = tf.keras.layers.Embedding(vocab_size_en, d_model)
embedding_layer_zh = tf.keras.layers.Embedding(vocab_size_zh, d_model)
emb_inp = embedding_layer_en(inp)
emb_tar = embedding_layer_zh(tar)
emb_inp, emb_tar
(<tf.Tensor: shape=(2, 8, 4), dtype=float32, numpy=
array([[[ 0.0041508 , 0.04106052, 0.00270988, -0.00628465],
[ 0.0261193 , 0.04892724, -0.03637441, 0.00032102],
[-0.0315491 , 0.03012072, -0.03764988, -0.00832593],
[-0.00863073, 0.01537497, 0.00647591, 0.01622475],
[ 0.01064278, 0.02867876, 0.0471475 , 0.02418466],
[-0.0357633 , -0.02500458, 0.00584758, 0.00984917],
[ 0.02766568, -0.02055204, 0.0366873 , -0.04519999],
[ 0.02766568, -0.02055204, 0.0366873 , -0.04519999]],
[[ 0.0041508 , 0.04106052, 0.00270988, -0.00628465],
[-0.03440493, 0.0245572 , -0.04154334, 0.01249687],
[-0.04102417, -0.04214551, -0.03087332, 0.03536062],
[ 0.00288613, -0.00550915, 0.02198391, -0.02721313],
[ 0.03594044, -0.02207484, 0.00774273, -0.01938369],
[-0.00556026, 0.04242435, 0.03270287, -0.00513189],
[ 0.01064278, 0.02867876, 0.0471475 , 0.02418466],
[-0.0357633 , -0.02500458, 0.00584758, 0.00984917]]],
dtype=float32)>, <tf.Tensor: shape=(2, 10, 4), dtype=float32, numpy=
array([[[-0.00084939, -0.02029408, -0.04978932, -0.02889797],
[-0.01320463, 0.00070287, 0.00797179, -0.00549082],
[-0.01859868, -0.04142375, 0.02479618, -0.00794141],
[ 0.04030085, -0.04564189, -0.03584541, -0.04098076],
[ 0.02629851, 0.01072141, -0.01055797, 0.04544314],
[-0.00223017, 0.02058548, 0.01649131, -0.01385387],
[ 0.00302396, -0.03152249, 0.0396189 , -0.03036447],
[ 0.00433234, 0.04481849, 0.04129448, 0.04720709],
[ 0.00433234, 0.04481849, 0.04129448, 0.04720709],
[ 0.00433234, 0.04481849, 0.04129448, 0.04720709]],
[[-0.00084939, -0.02029408, -0.04978932, -0.02889797],
[-0.04702241, 0.01816512, -0.02416607, -0.01993601],
[ 0.04391925, -0.03093947, -0.01225864, -0.03517971],
[ 0.03755457, 0.00626134, 0.04324439, 0.00490584],
[ 0.00495391, -0.03399891, 0.04144105, 0.02539945],
[ 0.0282723 , -0.0164601 , -0.00685417, -0.02280444],
[ 0.04738505, -0.01041915, -0.02054645, -0.00066562],
[-0.00438491, 0.02117647, -0.04890387, -0.01620366],
[-0.00223017, 0.02058548, 0.01649131, -0.01385387],
[ 0.00302396, -0.03152249, 0.0396189 , -0.03036447]]],
dtype=float32)>)
注意你的词嵌入层的随机初始值会跟我不同，结果可能会有一点差异。
但重点是你能在脑海中理解这两个 3 维张量吗？花了几秒钟？我相信在座不乏各路高手，而且事实上在这一行混久了，你也必须能直觉地理解这个表示方式。
但如果有更好的呈现方式帮助我们理解数据，何乐而不为呢？让我们再次可视化这两个 3 维词嵌入张量：
依照前面提过的准则，张量中第一个维度的 2 代表着句子数 batch_size。在 3 维空间里头，我会将不同句子画在 z 轴上，也就是你现在把脸贴近 / 远离屏幕这个维度。你同时也能用不同颜色来区分句子。
紧跟着句子的下一个维度则一样是本来的子词（subword）。只是现在每个子词都已从一个索引数字被转换成一个 4 维的词嵌入向量，因此每个子词都以 y 轴来表示。最后一维则代表着词嵌入空间的维度，一样以 x 轴来表示。
在学会怎么解读这个 3 维词嵌入张量以后，你就能明白为何 emb_tar 第一个中文句子里头的倒数 3 行（row) 都长得一样了：
print("tar[0]:", tar[0][-3:])
print("-" * 20)
print("emb_tar[0]:", emb_tar[0][-3:])
tar[0]: tf.Tensor([0 0 0], shape=(3,), dtype=int64)
--------------------
emb_tar[0]: tf.Tensor(
[[0.00433234 0.04481849 0.04129448 0.04720709]
[0.00433234 0.04481849 0.04129448 0.04720709]
[0.00433234 0.04481849 0.04129448 0.04720709]], shape=(3, 4), dtype=float32)
它们都是<pad>token（以 0 表示），理当有一样的词嵌入向量。
不同颜色也让我们可以很直观地理解一个句子是怎么从一个 1 维向量被转换到 2 维的。你后面就会发现，你将需要能够非常直觉地理解像是 emb_tar 这种 3 维张量里头每个维度所代表的意义。
5.3 Mask：Transformer 的秘密配方
我们在前面并没有仔细谈过遮罩（masking）的概念，但事实上它可以说是在实现 Transformer 时最重要却也最容易被搞砸的一环。它让 Transformer 在进行自注意力机制（Self-Attention Mechanism）时不至于偷看到不该看的。
在 Transformer 里头有两种 masks：
- padding mask
- look ahead mask
padding mask 是让 Transformer 用来识别序列实际的内容截止到哪里。此遮罩负责的就是将序列中被补 0 （也就是<pad>）的位置盖住，让 Transformer 可以避免「关注」到这些位置。
look ahead mask 人如其名，是用来确保Decoder 在进行自注意力机制时每个子词只会「往前看」：关注（包含）自己之前的字词，不会不小心关注「未来」Decoder 产生的子词。我们后面还会看到 look ahead mask 的详细介绍，但不管是哪一种遮罩向量，那些值为 1 的位置就是遮罩存在的地方。
因为 padding mask 的概念相对简单，让我们先看这种遮罩：
def create_padding_mask(seq):
# padding mask 的工作就是把索引序列中为 0 的位置设为 1
mask = tf.cast(tf.equal(seq, 0), tf.float32)
return mask[:, tf.newaxis, tf.newaxis, :] # broadcasting
inp_mask = create_padding_mask(inp)
inp_mask
<tf.Tensor: shape=(2, 1, 1, 8), dtype=float32, numpy=
array([[[[0., 0., 0., 0., 0., 0., 1., 1.]]],
[[[0., 0., 0., 0., 0., 0., 0., 0.]]]], dtype=float32)>
我们的第一个 4 维张量！不过别紧张，我们在中间加了 2 个新维度是为了之后可以做 broadcasting，现在可以忽视。喔！不过如果这是你第一次听到 broadcasting，我强烈建议你现在就阅读 numpy 官方的简短教学了解其概念。我们后面也会看到实际的 broadcasting 例子。
回到我们的 inp_mask 遮罩。现在我们可以先将额外的维度去掉以方便跟 inp 作比较：
print("inp:", inp)
print("-" * 20)
print("tf.squeeze(inp_mask):", tf.squeeze(inp_mask))
inp: tf.Tensor(
[[8113 103 9 1066 7903 8114 0 0]
[8113 16 4111 6735 12 2750 7903 8114]], shape=(2, 8), dtype=int64)
--------------------
tf.squeeze(inp_mask): tf.Tensor(
[[0. 0. 0. 0. 0. 0. 1. 1.]
[0. 0. 0. 0. 0. 0. 0. 0.]], shape=(2, 8), dtype=float32)
你可以看到 inp_mask将 inp 里头为0 的对应位置设为 1 凸显出来，这样之后其他函数就知道要把哪边「遮住」。让我们看看被降到 2 维的 inp_mask 是怎么被套用在 inp 身上的：
很好懂，不是吗？但这只是小暖身，等到之后要将遮罩 broadcast 到 3、4 维张量的时候你可能会黑人问号，所以最好做点心理准备
至于另外一种遮罩 look ahead mask，等我们说明完下节的注意力函数以后你会比较容易理解它的作用，所以先卖个关子。现在让我们进入 Tranformer 最核心的部分：注意力机制。
5.4 Scaled dot product attention：一种注意函数
我们在文中以及教授的影片已经多次看到，所谓的注意力机制（或称注意力函数，attention function）概念上就是拿一个查询（query）去跟一组key-values 做运算，最后产生一个输出。只是我们会利用矩阵运算同时让多个查询跟一组 key-values 做运算，最大化计算效率。
而不管是查询（query）、键值（keys）还是值（values）或是输出，全部都是向量（vectors）。该输出是 values 的加权平均，而每个 value 获得的权重则是由当初 value 对应的 key 跟 query 计算匹配程度所得来的。 （论文原文称此计算匹配程度的函数为 compatibility function）
将此运算以图表示的话则会长得像这样：
左右两边大致上讲的是一样的事情，不过右侧省略 Scale 以及 Mask 步骤，而左侧则假设我们已经拿到经过线性转换的 Q, K, V
我们是第一次秀出论文里头的图片（左），但右边你应该不陌生才对。
Scaled dot product attention 跟以往multiplicative attention 一样是先将维度相同的Q 跟K 做点积：将对应维度的值两两相乘后相加得到单一数值，接着把这些数值除以一个scaling factor sqrt(dk ) ，然后再丢入softmax 函数得到相加为1 的注意力权重（attention weights）。
最后以此权重对 V 作加权平均得到输出结果。
除以 scaling factor 的目的是为了让点积出来的值不会因为 Q 以及 K 的维度dk 太大而跟着太大（舌头打结）。因为太大的点积值丢入 softmax 函数有可能使得其梯度变得极小，导致训练结果不理想。
说完概念，让我们看看 Transformer 论文中的这个注意力函数怎么运作吧！首先我们得先准备这个函数的输入 Q, K, V 才行。我们在 Multi-head attention 一节会看到，在进行 scaled dot product attention 时会需要先将 Q、K 以及 V 分别做一次线性转换，但现在让我们先忽略这部分。
这边我们可以拿已经被转换成词嵌入空间的英文张量 emb_inp来充当左图中的 Q 以及 K，让它自己跟自己做匹配。 V 则让我随机产生一个 binary 张量（里头只有 1 或 0）来当作每个 K 所对应的值，方便我们直观解读 scaled dot product attention 的输出结果：
# 设定一个 seed 确保我们每次都拿到一样的随机结果
tf.random.set_seed(9527)
# 自注意力机制：查询 `q` 跟键值 `k` 都是 `emb_inp`
q = emb_inp
k = emb_inp
# 简单产生一个跟 `emb_inp` 同样 shape 的 binary vector
v = tf.cast(tf.math.greater(tf.random.uniform(shape=emb_inp.shape), 0.5), tf.float32)
v
<tf.Tensor: shape=(2, 8, 4), dtype=float32, numpy=
array([[[1., 0., 0., 0.],
[0., 1., 0., 1.],
[0., 0., 0., 1.],
[1., 0., 1., 0.],
[1., 0., 1., 0.],
[0., 1., 0., 1.],
[0., 0., 1., 0.],
[0., 1., 0., 1.]],
[[1., 0., 1., 1.],
[1., 0., 1., 0.],
[1., 0., 0., 0.],
[1., 0., 1., 0.],
[0., 1., 0., 1.],
[1., 1., 1., 1.],
[0., 0., 0., 0.],
[0., 0., 1., 0.]]], dtype=float32)>
现在应该能快速地解读 3 维张量了，但还是让我鸡婆点，将现在的 Q, K, V 都画出来让你参考：
注意不同颜色代表不同句子。虽然我们将拿 Q 跟 K 来做匹配，这个匹配只会发生在同个句子（同个颜色）底下（即 shape[1:]）。在深度学习世界，我们会为了最大化 GPU 的运算效率而一次将 64 个、128 个或是更多个 batch_size 的句子丢入模型。习惯 batch 维度的存在是非常实际的。
接着让我们看看 scaled dot product attention 在 TensorFlow 里是怎么被实现的：
def scaled_dot_product_attention(q, k, v, mask):
"""Calculate the attention weights.
q, k, v must have matching leading dimensions.
k, v must have matching penultimate dimension, i.e.: seq_len_k = seq_len_v.
The mask has different shapes depending on its type(padding or look ahead)
but it must be broadcastable for addition.
Args:
q: query shape == (..., seq_len_q, depth)
k: key shape == (..., seq_len_k, depth)
v: value shape == (..., seq_len_v, depth_v)
mask: Float tensor with shape broadcastable
to (..., seq_len_q, seq_len_k). Defaults to None.
Returns:
output, attention_weights
"""
# 将 `q`、 `k` 做点积再 scale
matmul_qk = tf.matmul(q, k, transpose_b=True) # (..., seq_len_q, seq_len_k)
dk = tf.cast(tf.shape(k)[-1], tf.float32) # 取得 seq_k 的序列長度
scaled_attention_logits = matmul_qk / tf.math.sqrt(dk) # scale by sqrt(dk)
# 将遮罩「加」到被丢入 softmax 前的 logits
if mask is not None:
scaled_attention_logits += (mask * -1e9)
# 取 softmax 是为了得到总和为 1 的比例之后对 `v` 做加权平均
attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1) # (..., seq_len_q, seq_len_k)
# 以注意权重对 v 做加权平均（weighted average）
output = tf.matmul(attention_weights, v) # (..., seq_len_q, depth_v)
return output, attention_weights
别被吓到了。除了Mask的运算部分我们还没解释，这 Python 函数事实上就是用 TensorFlow API 来实现刚刚才说的注意力机制逻辑罢了：
- 将
q和k做点积得到matmul_qk - 将
matmul_qk除以 scaling factorsqrt(dk) - 有遮罩的话在丢入 softmax 前套用
- 通过 softmax 取得加总为 1 的注意力权重
- 以该权重加权平均
v作为输出结果 - 回传输出结果以及注意力权重
扣掉注解事实上也就只有 8 行代码（当然有很多实作细节）。现在先让我们实际将 q, k, v输入此函数看看得到的结果。假设没有Mask的存在：
mask = None
output, attention_weights = scaled_dot_product_attention(q, k, v, mask)
print("output:", output)
print("-" * 20)
print("attention_weights:", attention_weights)
output: tf.Tensor(
[[[0.37510788 0.3749199 0.37497687 0.49994978]
[0.3750708 0.37495327 0.37484276 0.5000553 ]
[0.37498713 0.37496397 0.37476513 0.5001724 ]
[0.37511605 0.37494063 0.37501842 0.49995625]
[0.3752662 0.37487504 0.3752127 0.49974853]
[0.37497336 0.37501383 0.37500554 0.5000446 ]
[0.37485734 0.37506348 0.37516522 0.49988195]
[0.37485734 0.37506348 0.37516522 0.49988195]]
[[0.62507355 0.25005415 0.62513506 0.37515232]
[0.6252353 0.2498167 0.625135 0.37484628]
[0.6250702 0.24973223 0.62488395 0.37459227]
[0.6249318 0.25009662 0.6250206 0.37509808]
[0.6248611 0.25014114 0.62489796 0.37512854]
[0.62501526 0.25007918 0.625142 0.37516677]
[0.6249012 0.25008777 0.62500495 0.37513706]
[0.62500435 0.24987325 0.62496346 0.37478358]]], shape=(2, 8, 4), dtype=float32)
--------------------
attention_weights: tf.Tensor(
[[[0.1250733 0.12508997 0.1250299 0.12499584 0.12503874 0.12488762
0.12494231 0.12494231]
[0.12510203 0.12525114 0.125102 0.12499446 0.1249743 0.12482812
0.124874 0.124874 ]
[0.12506245 0.12512252 0.12520844 0.12501872 0.12490594 0.12500098
0.12484048 0.12484048]
[0.12502532 0.1250119 0.12501565 0.12503189 0.12505881 0.125001
0.12492773 0.12492773]
[0.12503879 0.12496231 0.12487347 0.12502936 0.12519807 0.12492746
0.12498529 0.12498529]
[0.12494987 0.1248783 0.12503074 0.1250338 0.12498969 0.1251535
0.12498205 0.12498205]
[0.12495281 0.12487246 0.12481848 0.12490878 0.12499575 0.12493028
0.12526071 0.12526071]
[0.12495281 0.12487246 0.12481848 0.12490878 0.12499575 0.12493028
0.12526071 0.12526071]]
[[0.12509817 0.1250309 0.12485092 0.12498978 0.12495036 0.12510379
0.12506361 0.12491246]
[0.12502958 0.12521692 0.12511879 0.12489447 0.12484112 0.12497558
0.12490506 0.12501846]
[0.12486006 0.12512927 0.12535231 0.12490248 0.12490615 0.12482607
0.12485762 0.12516604]
[0.12500146 0.12490747 0.12490503 0.12507936 0.12505813 0.12503849
0.12501612 0.12499388]
[0.12498739 0.12487943 0.12493402 0.1250835 0.12516432 0.12497681
0.12500364 0.12497085]
[0.12508757 0.12496072 0.12480079 0.12501062 0.12492362 0.12515557
0.1251336 0.12492751]
[0.1250493 0.12489209 0.12483419 0.12499013 0.1249523 0.12513547
0.12520859 0.12493796]
[0.12491032 0.12501764 0.12515475 0.12498006 0.1249317 0.12494154
0.12495013 0.12511387]]], shape=(2, 8, 8), dtype=float32)
scaled_dot_product_attention 函数输出两个张量：
-
output代表注意力机制的结果 -
attention_weights代表句子q里头每个子词对句子k里头的每个子词的注意力权重
而因为你知道目前的 q 跟 k 都代表同个张量 emb_inp，因此 attention_weights 事实上就代表了 emb_inp 里头每个英文序列中的子词对其他位置的子词的注意力权重。你可以再次参考之前 Transformer 是如何做 encoding 的动画。
output 则是句子里头每个位置的子词将 attention_weights 当作权重，从其他位置的子词对应的信息 v 里头抽取有用信息后汇总出来的结果。你可以想像 output 里头的每个子词都获得了一个包含自己以及周遭子词语义信息的新 representation。而因为现在每个字词的注意权重都相同，最后得到的每个 repr. 都长得一样。
下面则是我们实现的注意力函数的所有输入与输出张量。透过多次的矩阵运算，注意力机制能让查询 Q 跟键值 K 做匹配，再依据此匹配程度将值 V 做加权平均获得新的 representation。
动画里包含许多细节，但只要有矩阵运算的基本概念，你应该已经能够直观且正确地理解注意力函数是怎么运作的了。在真实世界里我们当然会用更长的序列、更大的 batch_size 来处理数据，但上面呈现的是代码的实际结果，而非示意图而已。这是注意力机制真正的「所见即所得」。
一般来说注意力函数的输出output张量维度会跟q张量相同（假设图上的 depth_v 等于 depth）。此张量也被称作「注意力张量」，你可以将其解读为 q在关注 k 并从 v 得到上下文信息后的所获得的新 representation。而注意力权重 attention_weights 则是 q里头每个句子的每个子词对其他位置的子词的关注程度。
P.S. 一般注意力函数只需输出注意力张量。而我们在这边将注意力权重attention_weights也输出是为了方便之后观察 Transformer 在训练的时候将「注意力」放在哪里。
5.5 直观理解遮罩在注意函数中的效果
刚刚为了让你消化注意力函数里头最重要的核心逻辑，我刻意忽略了遮罩（masking）的存在。现在让我们重新把scaled_dot_product_attention里头跟遮罩相关的代码拿出来瞧瞧：
...
# 将 `q`、 `k` 做点积再 scale
scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
# 将遮罩「加」到被丢入 softmax 前的 logits
if mask is not None:
scaled_attention_logits += (mask * -1e9)
# 取 softmax 是为了得到总和为 1 的比例做加权平均
attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
...
---------------------------------------------------------------------------
NameError Traceback (most recent call last)
<ipython-input-43-6891c377073d> in <module>()
2
3 # 将 `q`、 `k` 做点积再 scale
----> 4 scaled_attention_logits = matmul_qk / tf.math.sqrt(dk)
5
6 # 将遮罩「加」到被丢入 softmax 前的 logits
NameError: name 'matmul_qk' is not defined
如果你刚刚有仔细看上面的动画的话（17 秒之后），应该能想像 scaled_attention_logits 的 shape 为 （batch_size, seq_len_q, seq_len_k）。其最后一个维度代表某个序列 q 里的某个子词与序列k 的每个子词的匹配程度，但加总不为 1。而为了之后跟与k 对应的v 做加权平均，我们针对最后一个维度做 softmax 运算使其和为 1，也就是上图 axis=-1 的部分：
如果序列 k 里头的每个子词 sub_k 都是实际存在的中文字或英文词汇，这运算当然没问题。我们会希望序列 q 里头的每个子词 sub_q都能从每个 sub_k获得它所需要的语义资讯。
但李组长眉头一皱，发现案情并不单纯。
回想一下，我们的 q 跟 k 都是从emb_inp 来的。 emb_inp 代表着英文句子的词嵌入张量，而里头的第一个句子应该是有 <pad> token 的。啊哈！谁会想要放注意力在没有实际语义的家伙上呢？
...
if mask is not None:
scaled_attention_logits += (mask * -1e9) # 是 -1e9 不是 1e-9
attention_weights = tf.nn.softmax(scaled_attention_logits, axis=-1)
...
因此在注意力函数里头，我们将遮罩乘上一个接近负无穷大的-1e9，并把它加到进入 softmax 前的 logits 上面。这样可以让这些被加上极大负值的位置变得无关紧要，在经过 softmax 以后的值趋近于 0。这效果等同于序列q 中的某个子词 sub_q 完全没放注意力在这些被遮罩盖住的子词 sub_k 之上（此例中 sub_k 指是的<pad>）。
（动脑时间：为何遮罩要放在 softmax 之前而不能放之后？）
听我说那么多不如看实际的运算结果。让我们再次为英文句子 inp产生对应的 padding mask：
def create_padding_mask(seq):
# padding mask 的工作就是把索引序列中为 0 的位置设为 1
mask = tf.cast(tf.equal(seq, 0), tf.float32)
return mask[:, tf.newaxis, tf.newaxis, :] # broadcasting
print("inp:", inp)
inp_mask = create_padding_mask(inp)
print("-" * 20)
print("inp_mask:", inp_mask)
inp: tf.Tensor(
[[8113 103 9 1066 7903 8114 0 0]
[8113 16 4111 6735 12 2750 7903 8114]], shape=(2, 8), dtype=int64)
--------------------
inp_mask: tf.Tensor(
[[[[0. 0. 0. 0. 0. 0. 1. 1.]]]
[[[0. 0. 0. 0. 0. 0. 0. 0.]]]], shape=(2, 1, 1, 8), dtype=float32)
很明显地，第一个英文序列的最后 2 个位置是不具任何语义的 <pad>（图中为 0 的部分）。而这也是为何我们需要将遮罩 inp_mask 输入到注意力函数，避免序列中的子词关注到这 2 个家伙的原因。
我们这次把 inp_mask降到 3 维，并且将其跟刚刚的q、k 和 v 一起丢进注意函数里头，看看注意权重有什么变化：
# 这次让我们将 padding mask 放入注意函数并观察
# 注意权重的变化
mask = tf.squeeze(inp_mask, axis=1) # (batch_size, 1, seq_len_q)
_, attention_weights = scaled_dot_product_attention(q, k, v, mask)
print("attention_weights:", attention_weights)
attention_weights: tf.Tensor(
[[[0.16673873 0.16676098 0.16668089 0.16663547 0.16669267 0.1664912
0. 0. ]
[0.16674666 0.16694541 0.16674663 0.1666033 0.16657643 0.16638157
0. 0. ]
[0.16667904 0.16675909 0.1668736 0.16662073 0.16647044 0.16659711
0. 0. ]
[0.16666831 0.16665043 0.16665542 0.16667706 0.16671295 0.16663589
0. 0. ]
[0.16671184 0.16660987 0.16649143 0.16669928 0.1669242 0.1665634
0. 0. ]
[0.16659185 0.16649644 0.16669968 0.16670376 0.16664495 0.16686334
0. 0. ]
[0.16671966 0.16661246 0.16654043 0.16666092 0.16677696 0.16668959
0. 0. ]
[0.16671966 0.16661246 0.16654043 0.16666092 0.16677696 0.16668959
0. 0. ]]
[[0.12509817 0.1250309 0.12485092 0.12498978 0.12495036 0.12510379
0.12506361 0.12491246]
[0.12502958 0.12521692 0.12511879 0.12489447 0.12484112 0.12497558
0.12490506 0.12501846]
[0.12486006 0.12512927 0.12535231 0.12490248 0.12490615 0.12482607
0.12485762 0.12516604]
[0.12500146 0.12490747 0.12490503 0.12507936 0.12505813 0.12503849
0.12501612 0.12499388]
[0.12498739 0.12487943 0.12493402 0.1250835 0.12516432 0.12497681
0.12500364 0.12497085]
[0.12508757 0.12496072 0.12480079 0.12501062 0.12492362 0.12515557
0.1251336 0.12492751]
[0.1250493 0.12489209 0.12483419 0.12499013 0.1249523 0.12513547
0.12520859 0.12493796]
[0.12491032 0.12501764 0.12515475 0.12498006 0.1249317 0.12494154
0.12495013 0.12511387]]], shape=(2, 8, 8), dtype=float32)
加了 padding mask 后，第一个句子里头的每个子词针对倒数两个子词的「注意力权重」的值都变成 0 了。上句话非常饶舌，但我相信已经是非常精准的说法了。让我把这句话翻译成 numpy 的 slice 语法：
# 事实上也不完全是上句话的翻译，
# 因为我们在第一个维度还是把两个句子都拿出来方便你比较
attention_weights[:, :, -2:]
<tf.Tensor: shape=(2, 8, 2), dtype=float32, numpy=
array([[[0. , 0. ],
[0. , 0. ],
[0. , 0. ],
[0. , 0. ],
[0. , 0. ],
[0. , 0. ],
[0. , 0. ],
[0. , 0. ]],
[[0.12506361, 0.12491246],
[0.12490506, 0.12501846],
[0.12485762, 0.12516604],
[0.12501612, 0.12499388],
[0.12500364, 0.12497085],
[0.1251336 , 0.12492751],
[0.12520859, 0.12493796],
[0.12495013, 0.12511387]]], dtype=float32)>
第一个英文句子的最后2 个位置因为是<pad>所以被遮罩「盖住」而没有权重值（上方2 维阵列）；第二个句子的序列（下方2 维阵列）则因为最后2个位置仍是正常的英文子词，因此都有被其他子词关注。
如果听完我的碎碎念你还是无法理解以上结果，或是不确定有遮罩的注意力函数到底怎么运作，就实际看看其中的计算过程吧！
一张图胜过千言万语。在padding mask 的帮助下，注意力函数输出的新序列output里头的每个子词都只从序列k （也就是序列q 自己）的前6 个实际子词而非<pad> 来获得语义信息（最后一张图的黑框部分）。
再次提醒，因为我们输入注意力函数的q 跟k 都是同样的英文词嵌入张量emb_inp，事实上这边做的就是让英文句子里头的每个子词都去关注同句子中其他位置的子词的信息，并从中获得上下文语义，而这就是所谓的自注意力机制（self-attention）：序列关注自己。
当序列 q 换成 Decoder 的输出序列而序列k变成 Encoder 的输出序列时，我们就变成在计算一般 Seq2Seq 模型中的注意力机制。这点观察非常重要，且我们在前面就已经提过了。
打铁趁热，让我们看看前面提过的另一种遮罩 look ahead mask：
# 建立一个 2 维矩阵，维度为 (size, size)，
# 其遮罩为一个右上角的三角形
def create_look_ahead_mask(size):
mask = 1 - tf.linalg.band_part(tf.ones((size, size)), -1, 0)
return mask # (seq_len, seq_len)
seq_len = emb_tar.shape[1] # 注意這次我們用中文的詞嵌入張量 `emb_tar`
look_ahead_mask = create_look_ahead_mask(seq_len)
print("emb_tar:", emb_tar)
print("-" * 20)
print("look_ahead_mask", look_ahead_mask)
emb_tar: tf.Tensor(
[[[-0.00084939 -0.02029408 -0.04978932 -0.02889797]
[-0.01320463 0.00070287 0.00797179 -0.00549082]
[-0.01859868 -0.04142375 0.02479618 -0.00794141]
[ 0.04030085 -0.04564189 -0.03584541 -0.04098076]
[ 0.02629851 0.01072141 -0.01055797 0.04544314]
[-0.00223017 0.02058548 0.01649131 -0.01385387]
[ 0.00302396 -0.03152249 0.0396189 -0.03036447]
[ 0.00433234 0.04481849 0.04129448 0.04720709]
[ 0.00433234 0.04481849 0.04129448 0.04720709]
[ 0.00433234 0.04481849 0.04129448 0.04720709]]
[[-0.00084939 -0.02029408 -0.04978932 -0.02889797]
[-0.04702241 0.01816512 -0.02416607 -0.01993601]
[ 0.04391925 -0.03093947 -0.01225864 -0.03517971]
[ 0.03755457 0.00626134 0.04324439 0.00490584]
[ 0.00495391 -0.03399891 0.04144105 0.02539945]
[ 0.0282723 -0.0164601 -0.00685417 -0.02280444]
[ 0.04738505 -0.01041915 -0.02054645 -0.00066562]
[-0.00438491 0.02117647 -0.04890387 -0.01620366]
[-0.00223017 0.02058548 0.01649131 -0.01385387]
[ 0.00302396 -0.03152249 0.0396189 -0.03036447]]], shape=(2, 10, 4), dtype=float32)
--------------------
look_ahead_mask tf.Tensor(
[[0. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
[0. 0. 1. 1. 1. 1. 1. 1. 1. 1.]
[0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
[0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
[0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
[0. 0. 0. 0. 0. 0. 1. 1. 1. 1.]
[0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
[0. 0. 0. 0. 0. 0. 0. 0. 1. 1.]
[0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]
[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]], shape=(10, 10), dtype=float32)
我们已经知道demo 用的中文（目标语言）的序列长度为10，而look ahead 遮罩就是产生一个2 维矩阵，其两个维度都跟中文的词嵌入张量emb_tar 的倒数第2 个维度（序列长度）一样，且里头是一个倒三角形（1 的部分）。
我们前面曾经说过look_ahead_mask 是用来确保Decoder 在进行自注意力机制时输出序列里头的每个子词只会关注到自己之前（左边）的字词，不会不小心关注到未来（右边）理论上还没被Decoder 生成的子词。
运用从 padding mask 学到的概念，想像一下如果把这个倒三角的遮罩跟之前一样套用到进入 softmax 之前的 scaled_attention_logits，输出序列 output 里头的每个子词的 repr. 会有什么性质？
温馨小提醒：scaled_attention_logits 里头的每一 row 记录了某个特定子词对其他子词的注意力权重。
# 让我们用目标语言（中文）的 batch
# 来模拟 Decoder 处理的情况
temp_q = temp_k = emb_tar
temp_v = tf.cast(tf.math.greater(
tf.random.uniform(shape=emb_tar.shape), 0.5), tf.float32)
# 将 look_ahead_mask 放入注意函数
_, attention_weights = scaled_dot_product_attention(
temp_q, temp_k, temp_v, look_ahead_mask)
print("attention_weights:", attention_weights)
attention_weights: tf.Tensor(
[[[1. 0. 0. 0. 0. 0.
0. 0. 0. 0. ]
[0.49993625 0.5000637 0. 0. 0. 0.
0. 0. 0. 0. ]
[0.33313918 0.33324018 0.33362064 0. 0. 0.
0. 0. 0. 0. ]
[0.25015473 0.24959427 0.24974442 0.25050652 0. 0.
0. 0. 0. 0. ]
[0.19992094 0.19995634 0.19986811 0.19993235 0.20032226 0.
0. 0. 0. 0. ]
[0.16662028 0.16671094 0.16666561 0.16660225 0.16663651 0.1667644
0. 0. 0. 0. ]
[0.14276649 0.14282921 0.14297587 0.1428981 0.1426524 0.14282906
0.1430489 0. 0. 0. ]
[0.12475154 0.1250249 0.12494163 0.12469215 0.12516622 0.1250809
0.12494732 0.12539533 0. 0. ]
[0.1108513 0.11109421 0.11102021 0.11079852 0.11121978 0.11114396
0.11102527 0.11142336 0.11142336 0. ]
[0.09973814 0.09995669 0.09989012 0.09969065 0.10006968 0.10000146
0.09989467 0.10025284 0.10025284 0.10025284]]
[[1. 0. 0. 0. 0. 0.
0. 0. 0. 0. ]
[0.49974102 0.500259 0. 0. 0. 0.
0. 0. 0. 0. ]
[0.33343259 0.3327918 0.3337756 0. 0. 0.
0. 0. 0. 0. ]
[0.24972922 0.24968663 0.25012997 0.2504542 0. 0.
0. 0. 0. 0. ]
[0.19977222 0.19974756 0.19997004 0.2001723 0.20033783 0.
0. 0. 0. 0. ]
[0.16670689 0.16651376 0.16681753 0.16664357 0.16658409 0.16673416
0. 0. 0. 0. ]
[0.14287265 0.14264986 0.14297736 0.14284472 0.14276604 0.14290506
0.14298435 0. 0. 0. ]
[0.12512803 0.12510416 0.12499326 0.12483411 0.1247746 0.12498765
0.12500983 0.12516837 0. 0. ]
[0.11106904 0.11113531 0.11109053 0.11115386 0.11109442 0.11110445
0.11107941 0.11110792 0.11116511 0. ]
[0.09994661 0.09991619 0.10005405 0.10004354 0.10006738 0.10002076
0.09995341 0.09986328 0.09999042 0.10014433]]], shape=(2, 10, 10), dtype=float32)
答案呼之欲出，套用 look ahead mask 的结果就是让序列 q里的每个字词只关注包含自己左侧的子词，在自己之后的位置的字词都不看。比方说两个中文句子的第一个子词都只关注自己：
attention_weights[:, 0, :]
<tf.Tensor: shape=(2, 10), dtype=float32, numpy=
array([[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.],
[1., 0., 0., 0., 0., 0., 0., 0., 0., 0.]], dtype=float32)>
注意到了吗？两个句子的第一个子词因为自己前面已经没有其他子词，所以将全部的注意力 1 都放在自己身上。让我们看看第二个子词：
attention_weights[:, 1, :]
<tf.Tensor: shape=(2, 10), dtype=float32, numpy=
array([[0.49993625, 0.5000637 , 0. , 0. , 0. ,
0. , 0. , 0. , 0. , 0. ],
[0.49974102, 0.500259 , 0. , 0. , 0. ,
0. , 0. , 0. , 0. , 0. ]],
dtype=float32)>
两个句子的第 2 个子词因为只能看到序列中的第一个子词以及自己，因此前两个位置的注意权重加和即为 1，后面位置的权重皆为 0。而现在 2 个值都接近 0.5 是因为我们还没开始训练，Transformer 还不知道该把注意力放在哪里。
就跟一般的 Seq2Seq 模型相同，Transformer 里头的 Decoder 在生成输出序列时也是一次产生一个子词。因此跟左侧输入的英文句子不同，右侧输出的中文句子里头的每个子词都是在不同时间点产生的。所以理论上Decoder 在时间点t - 1（或者说是位置t - 1）已经生成的子词subword_t_minus_1 在生成的时候是不可能能够关注到下个时间点 t（位置t）所生成的子词subword_t的，尽管它们在Transformer 里头同时被做矩阵运算。
一个位置的子词不能去关注未来会在自己之后生成的子词，而这就像是祖父悖论一样有趣。
实际上look ahead mask 让Decoder 在生成第1 个子词时只看自己；在生成第2 个子词时关注前1 个子词与自己； 在生成第3 个子词时关注前两个已经生成的子词以及自己，以此类推。透过 look ahead mask，你可以想像 Transformer 既可以并行运算，又可以像是 RNN 一样，在生成子词时从前面已生成的子词获得必要的语义信息。
在实际做矩阵运算的时候我们当然还是会让注意力权重为 0 的位置跟对应的 v 相乘，但是上图的黑框才是实际会对最后的 output值造成影响的权重与 v。
我们在这节了解 Transformer 架构里头的两种遮罩以及它们的作用：
- padding mask：遮住
<pad>token 不让所有子词关注 - look ahead mask：遮住 Decoder 未来生成的子词不让之前的子词关注
你现在应该能够想像遮罩在注意力机制里头显得有多么重要了：它让注意力函数进行高效率的矩阵平行运算的时候不需担心会关注到不该关注的位置，一次获得序列中所有位置的子词各自应有的注意力权重以及新的reprsentation。
如果 Transformer 是变形金刚的话，注意力机制跟遮罩就是火种源了。
5.6 Multi-head attention：你看你的，我看我的
有好好听教授讲解 Transformer 的话，你应该还记得所谓的多头注意力（multi-head attention）概念。如果你到现在还没看课程影片或者想复习一下，我把 multi-head attention 的开始跟结束时间都设置好了，你只需观看大约 1 分半左右的影片：
mutli-head attention 的概念本身并不难，用比较正式的说法就是将Q、K 以及V 这三个张量先分别转换到d_model 维空间，再将其拆成多个比较低维的depth 维度N 次以后，将这些产生的小q、小k 以及小v分别丢入前面的注意力函数得到N 个结果。接着将这 N 个 heads 的结果串接起来，最后通过一个线性转换就能得到 multi-head attention 的输出
而为何要那么「搞刚」把本来d_model 维的空间投影到多个维度较小的子空间（subspace）以后才各自进行注意力机制呢？这是因为这给予了模型更大的弹性，让它可以同时关注不同位置的子词在不同子空间下的 representation，而不只是本来d_model 维度下的一个 representation。
我们在文章最开头看过的英翻中就是一个活生生的 mutli-head attention 例子：
在经过前面 2 节注意力函数的洗礼之后，你应该已经能够看出这里每张小图都是一个注意权重（为了方便渲染我做了 transpose）。而事实上每张小图都是 multi-head attention 里头某一个 head 的结果，总共是 8 个 heads。
你会发现每个 head 在 Transformer 生成同样的中文字时关注的英文子词有所差异：
- Head 4 在生成「们」与「再」时特别关注「renewed」
- Head 5 在生成「必」与「须」时特别关注「must」
- Head 6 & 8 在生成「希」与「望」时特别关注「hope」
透过这样同时关注多个不同子空间里头的子词的 representation，Transformer 最终可以生成更好的结果。
话是这么说，但代码该怎么写呢？
为了要实现 multi-head attention，得先能把一个 head 变成多个 heads。而这实际上就是把一个d_model 维度的向量「折」成 num_heads 个 depth维向量，使得：
num_heads * depth = d_model
让我们实现一个可以做到这件事情的函数，并将英文词嵌入张量 emb_inp 实际丢进去看看：
def split_heads(x, d_model, num_heads):
# x.shape: (batch_size, seq_len, d_model)
batch_size = tf.shape(x)[0]
# 我們要確保維度 `d_model` 可以被平分成 `num_heads` 個 `depth` 維度
assert d_model % num_heads == 0
depth = d_model // num_heads # 這是分成多頭以後每個向量的維度
# 將最後一個 d_model 維度分成 num_heads 個 depth 維度。
# 最後一個維度變成兩個維度，張量 x 從 3 維到 4 維
# (batch_size, seq_len, num_heads, depth)
reshaped_x = tf.reshape(x, shape=(batch_size, -1, num_heads, depth))
# 將 head 的維度拉前使得最後兩個維度為子詞以及其對應的 depth 向量
# (batch_size, num_heads, seq_len, depth)
output = tf.transpose(reshaped_x, perm=[0, 2, 1, 3])
return output
# 我們的 `emb_inp` 裡頭的子詞本來就是 4 維的詞嵌入向量
d_model = 4
# 將 4 維詞嵌入向量分為 2 個 head 的 2 維矩陣
num_heads = 2
x = emb_inp
output = split_heads(x, d_model, num_heads)
print("x:", x)
print("output:", output)
x: tf.Tensor(
[[[ 0.0041508 0.04106052 0.00270988 -0.00628465]
[ 0.0261193 0.04892724 -0.03637441 0.00032102]
[-0.0315491 0.03012072 -0.03764988 -0.00832593]
[-0.00863073 0.01537497 0.00647591 0.01622475]
[ 0.01064278 0.02867876 0.0471475 0.02418466]
[-0.0357633 -0.02500458 0.00584758 0.00984917]
[ 0.02766568 -0.02055204 0.0366873 -0.04519999]
[ 0.02766568 -0.02055204 0.0366873 -0.04519999]]
[[ 0.0041508 0.04106052 0.00270988 -0.00628465]
[-0.03440493 0.0245572 -0.04154334 0.01249687]
[-0.04102417 -0.04214551 -0.03087332 0.03536062]
[ 0.00288613 -0.00550915 0.02198391 -0.02721313]
[ 0.03594044 -0.02207484 0.00774273 -0.01938369]
[-0.00556026 0.04242435 0.03270287 -0.00513189]
[ 0.01064278 0.02867876 0.0471475 0.02418466]
[-0.0357633 -0.02500458 0.00584758 0.00984917]]], shape=(2, 8, 4), dtype=float32)
output: tf.Tensor(
[[[[ 0.0041508 0.04106052]
[ 0.0261193 0.04892724]
[-0.0315491 0.03012072]
[-0.00863073 0.01537497]
[ 0.01064278 0.02867876]
[-0.0357633 -0.02500458]
[ 0.02766568 -0.02055204]
[ 0.02766568 -0.02055204]]
[[ 0.00270988 -0.00628465]
[-0.03637441 0.00032102]
[-0.03764988 -0.00832593]
[ 0.00647591 0.01622475]
[ 0.0471475 0.02418466]
[ 0.00584758 0.00984917]
[ 0.0366873 -0.04519999]
[ 0.0366873 -0.04519999]]]
[[[ 0.0041508 0.04106052]
[-0.03440493 0.0245572 ]
[-0.04102417 -0.04214551]
[ 0.00288613 -0.00550915]
[ 0.03594044 -0.02207484]
[-0.00556026 0.04242435]
[ 0.01064278 0.02867876]
[-0.0357633 -0.02500458]]
[[ 0.00270988 -0.00628465]
[-0.04154334 0.01249687]
[-0.03087332 0.03536062]
[ 0.02198391 -0.02721313]
[ 0.00774273 -0.01938369]
[ 0.03270287 -0.00513189]
[ 0.0471475 0.02418466]
[ 0.00584758 0.00984917]]]], shape=(2, 2, 8, 2), dtype=float32)
观察 output与 emb_inp之间的关系，你会发现 3 维词嵌入张量emb_inp 已经被转换成一个 4 维张量了，且最后一个维度 shape[-1] = 4被拆成两半。
不过如果你不熟 TensorFlow API 或是矩阵运算，或许无法马上理解 head 的维度在哪里、还有不同 heads 之间有什么差异。为了帮助你直观理解 split_heads 函数，我将运算过程中产生的张量都可视化出来给你瞧瞧：
观察split_heads的输入输出，你会发现序列里每个子词原来为 d_model 维的 reprsentation 被拆成多个相同但较短的 depth 维度。而每个 head 的 2 维矩阵事实上仍然代表原来的序列，只是里头子词的 repr. 维度降低了。
透过动画，你现在应该已经能够了解要产生 multi-head 就是将输入张量中本来是 d_model 的最后一个维度平均地「折」成想要的 head 数，进而产生一个新的 head 维度。一个句子里头的子词现在不只会有一个d_model的 repr.，而是会有 num_heads 个 depth 维度的 representation。
接下来只要把3 维的Q、K 以及V 用 split_heads 拆成多个heads 的4 维张量，利用broadcasting 就能以之前定义的Scaled dot product attention 来为每个句子里头的每个head 并行计算注意结果了，超有效率！
在明白如何产生 multi-head 的 4 维张量以后，multi-head attention 的实现就比较容易理解了：
# 实作一个执行多头注意力机制的 keras layer
# 在初始的时候指定输出维度 `d_model` & `num_heads，
# 在呼叫的时候输入 `v`, `k`, `q` 以及 `mask`
# 输出跟 scaled_dot_product_attention 函数一样有两个：
# output.shape == (batch_size, seq_len_q, d_model)
# attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
class MultiHeadAttention(tf.keras.layers.Layer):
# 在初始的時候建立一些必要參數
def __init__(self, d_model, num_heads):
super(MultiHeadAttention, self).__init__()
self.num_heads = num_heads # 指定要將 `d_model` 拆成幾個 heads
self.d_model = d_model # 在 split_heads 之前的基底維度
assert d_model % self.num_heads == 0 # 前面看過，要確保可以平分
self.depth = d_model // self.num_heads # 每個 head 裡子詞的新的 repr. 維度
self.wq = tf.keras.layers.Dense(d_model) # 分別給 q, k, v 的 3 個線性轉換
self.wk = tf.keras.layers.Dense(d_model) # 注意我們並沒有指定 activation func
self.wv = tf.keras.layers.Dense(d_model)
self.dense = tf.keras.layers.Dense(d_model) # 多 heads 串接後通過的線性轉換
# 這跟我們前面看過的函数有 87% 相似
def split_heads(self, x, batch_size):
"""Split the last dimension into (num_heads, depth).
Transpose the result such that the shape is (batch_size, num_heads, seq_len, depth)
"""
x = tf.reshape(x, (batch_size, -1, self.num_heads, self.depth))
return tf.transpose(x, perm=[0, 2, 1, 3])
# multi-head attention 的實際執行流程，注意參數順序（這邊跟論文以及 TensorFlow 官方教學一致）
def call(self, v, k, q, mask):
batch_size = tf.shape(q)[0]
# 將輸入的 q, k, v 都各自做一次線性轉換到 `d_model` 維空間
q = self.wq(q) # (batch_size, seq_len, d_model)
k = self.wk(k) # (batch_size, seq_len, d_model)
v = self.wv(v) # (batch_size, seq_len, d_model)
# 前面看過的，將最後一個 `d_model` 維度分成 `num_heads` 個 `depth` 維度
q = self.split_heads(q, batch_size) # (batch_size, num_heads, seq_len_q, depth)
k = self.split_heads(k, batch_size) # (batch_size, num_heads, seq_len_k, depth)
v = self.split_heads(v, batch_size) # (batch_size, num_heads, seq_len_v, depth)
# 利用 broadcasting 讓每個句子的每個 head 的 qi, ki, vi 都各自進行注意力機制
# 輸出會多一個 head 維度
scaled_attention, attention_weights = scaled_dot_product_attention(
q, k, v, mask)
# scaled_attention.shape == (batch_size, num_heads, seq_len_q, depth)
# attention_weights.shape == (batch_size, num_heads, seq_len_q, seq_len_k)
# 跟我們在 `split_heads` 函数做的事情剛好相反，先做 transpose 再做 reshape
# 將 `num_heads` 個 `depth` 維度串接回原來的 `d_model` 維度
scaled_attention = tf.transpose(scaled_attention, perm=[0, 2, 1, 3])
# (batch_size, seq_len_q, num_heads, depth)
concat_attention = tf.reshape(scaled_attention,
(batch_size, -1, self.d_model))
# (batch_size, seq_len_q, d_model)
# 通過最後一個線性轉換
output = self.dense(concat_attention) # (batch_size, seq_len_q, d_model)
return output, attention_weights
是的，就算你有自己实现过 keras layer，multi-head attention layer 也不算短的实现。如果这是你第一次碰到定制化的 keras layer，别打退堂鼓，你可以多看几次我写给你的注解，或是参考等等下方的动画来加深对 multi-head attention 的理解。
split_heads 函数我们在前面就已经看过了，你应该还有印象。 call 函数则定义了这个 multi-head attention layer 实际的计算流程，而这流程跟我在本节开头讲的可以说是有 87% 相似：
将Q、K 以及V 这三个张量先个别转换到d_model 维空间，再将其拆成多个比较低维的depth 维度N 次以后，将这些产生的小q、小k 以及小v 分别丢入前面的注意函数得到N 个结果。接着将这 N 个 heads 的结果串接起来，最后通过一个线性转换就能得到 multi-head attention 的输出
差别只在于实际上我们是利用矩阵运算以及 broadcasting 让 GPU 一次计算整个 batch 里所有句子的所有 head 的注意力结果。
定义了一个新 layer 当然要实际试试。现在让我们初始一个 multi-head attention layer 并将英文词嵌入向量 emb_inp 输入进去看看：
# emb_inp.shape == (batch_size, seq_len, d_model)
# == (2, 8, 4)
assert d_model == emb_inp.shape[-1] == 4
num_heads = 2
print(f"d_model: {d_model}")
print(f"num_heads: {num_heads}\n")
# 初始化一个 multi-head attention layer
mha = MultiHeadAttention(d_model, num_heads)
# 简单将 v, k, q 都设置为 `emb_inp`
# 顺便看看 padding mask 的作用。
# 别忘记，第一个英文序列的最后两个 tokens 是 <pad>
v = k = q = emb_inp
padding_mask = create_padding_mask(inp)
print("q.shape:", q.shape)
print("k.shape:", k.shape)
print("v.shape:", v.shape)
print("padding_mask.shape:", padding_mask.shape)
output, attention_weights = mha(v, k, q, mask)
print("output.shape:", output.shape)
print("attention_weights.shape:", attention_weights.shape)
print("\noutput:", output)
d_model: 4
num_heads: 2
q.shape: (2, 8, 4)
k.shape: (2, 8, 4)
v.shape: (2, 8, 4)
padding_mask.shape: (2, 1, 1, 8)
output.shape: (2, 8, 4)
attention_weights.shape: (2, 2, 8, 8)
output: tf.Tensor(
[[[ 0.0041762 -0.00345554 0.00731694 0.0152422 ]
[ 0.0041797 -0.00345984 0.00730736 0.0152359 ]
[ 0.00415735 -0.00343554 0.00734071 0.01524936]
[ 0.00418272 -0.00346123 0.00731704 0.01524691]
[ 0.00419747 -0.00347528 0.00730165 0.01524222]
[ 0.00417466 -0.00345488 0.00733047 0.01525742]
[ 0.00418267 -0.00347137 0.00730291 0.01524842]
[ 0.00418267 -0.00347137 0.00730291 0.01524842]]
[[ 0.00347738 -0.00410287 -0.00210699 0.00332957]
[ 0.00347175 -0.00409482 -0.00208811 0.00333377]
[ 0.0034817 -0.0041021 -0.00207374 0.00335214]
[ 0.00347859 -0.00410586 -0.00210912 0.0033298 ]
[ 0.00347468 -0.00410707 -0.00211646 0.00332593]
[ 0.00348684 -0.00410802 -0.00210311 0.0033386 ]
[ 0.00349733 -0.0041161 -0.00209817 0.00335264]
[ 0.0034855 -0.00410557 -0.00208397 0.00334841]]], shape=(2, 8, 4), dtype=float32)
你现在应该明白为何我们当初要在padding mask 加入两个新维度了：一个是用来遮住同个句子但是不同head 的注意权重，一个则是用来broadcast 到2 维注意权重的（详见直观理解遮罩一节）。
没意外的话你也已经能够解读 mutli-head attention 的输出了：
-
output：序列中每个子词的新 repr. 都包含同序列其他位置的信息 -
attention_weights：包含每个 head 的每个序列q中的字词对序列k的注意权重
如果你还无法想像每个计算步骤，让我们看看 multi-head attention 是怎么将输入的 q、k 以及v 转换成最后的 output 的：
这应该是你这辈子第一次也可能是唯一一次有机会看到 multi-head 注意力机制是怎么处理 4 维张量的。
细节不少，我建议将动画跟代码比较一下，确保你能想像每一个步骤产生的张量以及其物理意义。到此为止，我们已经把 Transformer 里最困难的 multi-head attention 的概念以及运算都看过一遍了。
如果你脑袋还是一团乱，只要记得最后一个画面：在q、k 以及v 的最后一维已经是d_model 的情况下，multi-head attention 跟scaled dot product attention 一样，就是吐出一个完全一样维度的张量output。
multi-head attention 的输出张量output 里头每个句子的每个字词的repr. 维度d_model 虽然跟函数的输入张量相同，但实际上已经是从同个序列中不同位置且不同空间中的repr. 取得语义信息的结果。
要确保自己真的掌握了 multi-head attention 的精神，你可以试着向旁边的朋友（如果他 / 她愿意听的话）解释一下整个流程。
不用担心我们做 multi-head 以后计算量会大增。因为 head 的数目虽然变多了，每个子空间的维度也下降了。跟 single-head attention 使用的计算量是差不多的。
