Transformer

本 blog 的线索是先讲清楚 GPT-3 的内部工作原理。然后再去对着 transformer 论文里的图讲一讲。

关于对于 GPT-3 内部机制的讲解，强烈推荐这三个视频。可能不是句句都是干货，但是其切入的视角比较接近语义的本质，看完收益颇丰。

1 GPT-3

1.1 Token, Embedding Matrix

首先，所有文本的基本单位是——token。token 可以是一个词，或者几个字母，whatever。为了后面方便叙述，我们就将 token 认为是一个词。

这些 token 首先要转为 embedding ，也就是一个高维向量。具体转换方式很简单，GPT-3 内部有一个 embedding matrix ，可以理解为词表。有 50257 列（也就是有 50257 个 token ），每一个列向量的维度是 12288 维。所以一段句子丢进来，它们的 embedding 就查 embedding matrix 即可。

注意，这个 embedding matrix 是可学习的。

​​

将 token 编码为高维向量的意义是什么呢？在 \mathbb{R}^n 中，我们可以通过做点积来衡量两个向量之间的相似度，因此可以借助这一点来量化语义关系。这就是将其编码为高维向量的好处。

那向量的维度多少比较好呢？如果我们要求严格正交，那么 \mathbb{R}^n 空间里只能容忍 n 个方向的语义。但是如果我们只要将限制放宽一点点，即 $89\degree \sim 91\degree$ 都允许，那么 Johnson-Lindenstrauss lemma 指出：“在空间中塞进几乎垂直的向量，随维数增加指数增长”。

所以 GPT-3 的 12288 维，其实已经可以把至少 e^{12288} 个独立的语义编码进去了。语义多种多样，例如 “男性”，“人类”，“复数”，“好吃”，“喜欢大声说话的东西”。

因此这个原因可能解释了，为什么模型性能随规模扩大而显著提升。

1.2 Unembedding Matrix

每个 token 经过查表都有了自己的 embedding ，然后这些 embedding 经过多头注意力、MLP 等层，最终得到了结合上下文语义信息的 embedding，我们叫做 embedding' 。这些 embedding' 分别与 unembedding matrix 相乘的向量 softmax 后，就得到了每个 token 出现的概率。结合下图来加深理解：

​​

可以看到 unembedding matrix 有 50257 行，因为每行要与 embedding' 做点积得到一个标量，这个标量代表这一行对应的 token 出现的可能性数值。然后有 12288 列是因为要与 embedding' 的维度匹配。

同样的，这个 unembedding matrix 是可以训练的。

补充一下，人们喜欢把 softmax 前的列向量称为 logits ，softmax 后的列向量称为 probabilities 。

1.3 Attention Pattern（注意力模式）

​​

在一个 attention pattern 里，每个 embedding ，有三个东西：Q, K, V。

Q 和 K 是负责匹配信息的东西。Q 负责发出一个 “注意模式捕捉信号” ，K 负责回应一个 “注意模式捕捉信号” 。即第 i 个 embedding 与其余 embedding 在当前注意模式的 logits，就是 K_j \cdot Q_i, j \in [1, n]。

经过这样的处理，每个 embedding 都得到了与其余 embedding 在当前注意模式的 logits （一个列向量）。然后每列 softmax 一下就得到了概率，也可以理解为与其余 embedding 关联的一个权重。

如何得到 Q 和 K 呢？通过将 Q 矩阵和 K 矩阵与 embedding 相乘即可。Q/K 矩阵的维度是 $128 \times 12288$ 。

那 V 是干嘛的呢？V 可以理解为在当前注意力模式下，embedding 的语义信息。也是通过将 V 矩阵与 embedding 相乘得到。

所以对于第 i 个 embedding 来说，其在当前注意力模式下，其它 embedding 对其产生的贡献就是 \sum_{j=1}^n \text{softmax}(K_j \cdot Q_i) \cdot V_j 。那么最终输出的结果就是：E_i + \sum_{j=1}^n \text{softmax}(K_j \cdot Q_i) \cdot V_j 。

下面的图可以帮助很好的解释：

​​

V 矩阵的维度多大呢？首先列维度肯定是 12288 ，因为要跟 embedding 的维度匹配，然后行维度 GPT-3 选择的是 12288 ，即经过 V 矩阵变换后前后维度不变。

当然了这里有个小技巧，因为 $12288 \times 12288$ 这个矩阵太大了，所以可以将其分解为两个矩阵相乘：

​​

这个技巧就是低秩分解（lower rank，LoRA），很多 paper 喜欢把右边那个矩阵叫做 V 矩阵，左边的叫做 输出矩阵 。whatever they say, 我们清楚就好。

同样的，Q/K/V 矩阵是可以学习的。

1.4 Multi-head Attention Pattern

前面我在 Section 1.4 里一直提到一个词，注意力模式。什么叫一个注意力模式？我举个例子，前面的 Q,K,V 都是在某个注意力模式大前提下的，不同注意力模式的 Q,K,V 不同。

比如可能有一个注意力模式叫：“捕捉上下文颜色之间的关系”。或者，“捕捉上下文语法词性之间的关系”。或者，“捕捉上下文感情色彩之间的关系”。

所以如何捕捉上下文？就需要充分的捕捉上下文中蕴含的多种注意力模式。因此，GPT-3 中有 96 层 attention pattern 。显然不同注意力机制之间是互不干扰的，所以这 96 个 attention pattern 是并行运行的。

最终可以得到 \Delta E_i^{(1)}, \Delta E_i^{(2)}, \cdots, \Delta E_i^{(96)}，那么结合了所有注意力模式的上下文信息后，第 i 个 embedding 的最终输出就是：E_i + \Delta E_i^{(1)} + \Delta E_i^{(2)} + \cdots + \Delta E_i^{(96)} 。

1.5 Attention Pattern 里的掩码 mask trick

​​

在训练中，我们希望一句话能重复利用，榨干其训练价值。所以就是每一个词从前往后都去做预测，然后与正确答案去做对比（这其实就揭示了 GPT-3 是如何训练的了）。

这就要求，每个 embedding 不能只能接收它前面的 embedding 的信息。所以在用 Q,K 去算那个 logits 矩阵的时候，就要把那个矩阵变为一个 “上三角" 矩阵，即把下三角区域都赋值为 -\infty ，这样 softmax 之后权重矩阵就变为一个上三角矩阵了。也就是每个 embedding 不会受到后面 embedding 的影响了。

1.6 MLP

MLP 与 Multi-Head Attention 的关系是什么？看下面这张图：

​​

这张图里的 Attention 表示 Multi-Head Attention ，简略画了。所以在一个 transformer block 里，其实就是先经过 multi-head attention ，然后经过 MLP 。

（从上面的动图可以发现，对于所有 embeddings ，它们经过 MLP 是独立的。意思是只需要有一套 MLP 参数，那么每个 embedding 即可并行运行经过 MLP ）

根据科学家的研究，大模型学到的 “经验” ，似乎都存储在 MLP 的参数中。这似乎比较合理，因为 GPT-3 的 \frac23 的参数量都在 MLP 里，只有 \frac13 的参数量在 Attention 里。所以感觉 Attention 专注于处理上下文信息捕捉，而常识信息结合则需要通过经过 MLP 获得。

不过，DL 的可解释性一直不是一个确定的事情，上面只是一种观点。所以在 blog 里我就不打算花大篇幅去讲如何从语义变换, 存储的角度去理解 MLP 了。我只负责告诉你 MLP 的结构。

MLP 其实就干了三件事：Linear -> ReLU -> Linear：

​​

第一个 Linear ，就是有一个矩阵和一个 bias 列向量，如下图。GPT-3 的这个矩阵有 49152 行，列维度为 12288 。通过这个操作，embedding 的维度会由 12288 升高到 49152 。

​​

ReLU 就不说了。激活一下而已。

第二个 Linear ，也是一个矩阵和一个 bias 列向量，如下图。但是要把维度降回来，所以这个矩阵的维度是前边那个矩阵的逆。

​​

用公式来总结 embedding 经过 MLP 的变化过程，如下：

1.7 Positional Embedding

其实这一节应该放在 Section 1.1 那里讲的，不过我觉得这块内容不影响大局观，我们先要建立其整个 transformer block 的整体概念，再去细化各种细节。而 PE 就是其中一个细节。

位置编码肯定是很重要的，因为 “一只猫在桌子上” ，和，“一个桌子在猫上” 显然表达的意思不同。那么如何区别这两种意思呢？那么就要通过位置来区分。

首先先说一个最简单的想法，假设一段文本有 k 个 token(embedding) 。那么第一个 embedding 就给它每个元素加上 1，第二个 embedding 就给它每个元素加上 2 ... ...

这样可以，但是随着位置的增加，加的值会越来越大。这就是造成训练不稳定，因为我们通常希望神经网络里的数不要太大，最好在 [-1, 1] 之间。

那么，有人说，那假设文本有 k 个 embedding ，那么我就设置增量为 \frac{1}{k} 呗。然后第一个 embedding 给它每个元素加上 \frac{1}{k} \cdot 1 ，第二个 embedding 给它每个元素加上 \frac{1}{k} \cdot 2 ... ...

但是这样的想法不好，因为没有保持 “绝对位置” 。也就是，位置编码的效果应该是（某个位置，某个分量）的 PE 值应该是确定的。也就是只要确定了位置和第几个分量，那么就能唯一确认一个固定的 PE 值。这样的度量才是有效的。

但是刚才的想法，同一个位置，所加的数会随着文本 embedding 数量而变动，这就是不好的。

因此，我们总结出 PE 需要遵循的几个规则：

1. 取值范围最好在 [-1, 1]
2. 已知位置和第几个分量，那么就能唯一确认一个固定的 PE 值

那么，transformer 提出一种编码法则，可以满足上面俩条件。下面我直接给出编码法则：

​​

t 代表位置，d 代表 embedding 的维度，k 代表 embedding 里的第几个分量。

可以发现，当我们固定某个分量位置，随着位置的增大，它其实就是个周期性函数去取值。

当我们固定位置，随着分量位置的增大，它的函数的频率在降低。

可以理解为，每个分量上就是一个周期函数，但是函数频率随着分量位置增大而减小，此时脑中想象有 d 个频率不同的三角函数图像以 d \times 1 的形状排好，然后位置就是在 x 轴上取一个点，然后竖着对这 d \times 1 排好的图像切一刀，取到的纵坐标分别就是每个分量上的 PE 值。

2 Transformer Architecture

​​

其实把 Section 1 认真看完，transformer 的精髓就已经学到了。但是，论文里这张广为流传的图我们还是要来过一下。

首先我们要明白，这个图出自论文 “Attention is All You Need” ，而这篇论文的提出是为了解决机器翻译问题的，而不是像 GPT 那样的生成问题。所以这个架构乍一看好像无法下手分析。下面我来解读这个结构。

左下角的 inputs ，是要翻译的文本，比如是英文。

经过查表变 embedding 后，加上 PE 编码，经过 transformer block ，简单的过一个 multi-head attention 和 MLP 后，具有上下文语义的 embedding' 进去 decoder 结构。（Norm 就是对所有 embedding' 归一化因此而已）

右下角的 outputs 是翻译出来的文本，比如是中文。然后 outputs 是一个一个蹦出来的，最开始 outputs 的初始值是 <\s> ，这就是为什么右下角写了个 shifted right ，代表俺要开始翻译了！

所以当前已经生成的 outputs 从右下角进去，同样经过 embedding 和 PE 后，经过一个 Masked Multi-Head Attention 因为翻译后的文本的每个 token 只能利用前面的信息。然后就要注意了，接着的是一个叫 cross attention 的东西。

这个 cross attention 跟 self attention 的区别，就是 self attention 是全部利用自己的信息去捕捉上下文关系，而 cross attention 是捕捉别人的上下文关系！

实现方法就是把 K, V 换为别人的，Q 仍然是自己。这样利用 Q,K 算出来的 probabilities 矩阵就是每个 embedding 与别人的 embedding 之间的关联权重矩阵。这样最终算出来的 \underbrace{\sum_{j=1}^n \text{softmax}(K_j^{\text{别人}} \cdot Q_i) \cdot V_j^{\text{别人}}}_{\Delta E_i} 就是利用了别人上下文信息的向量。

OK，因为你要做翻译嘛，所以你肯定除了利用自己已经生成的 outputs 去做预测，还需要利用原文信息辅助，这就是为什么先进一个 masked multi-head self attention 再接一个 multi-head cross attention 的原因。

好啦，最终融合了自身上下文和原文上下文的 embedding' ，经过一个 MLP ，最终经过一个 Linear （就是 Unembedding Matrix），然后 softmax 一下即可得到 probabilities 啦！

讲实话，这个 cross attention 的模块还挺有启发性的。不觉得他扩展性很强吗？本质上就是融合其它信息的一种手段。