文章

Transformer 笔记

Transformer 笔记

RNN 介绍

RNN 实际完成的就是从一个序列到另一个序列。对于自然语言,序列的单位就是一个词,序列中的某个词 t 有一个隐藏状态 $h_t$ ,它是由 t 本身和前一个词的隐藏状态 $h_{t-1}$ 作为参数来生成的,所以该隐藏状态能够包含之前所有词的信息(上下文)。

缺点:

  1. 较难进行并行化处理,无法充分利用 GPU。
  2. 当上下文信息很长是,隐藏状态需要占用大量内存空间。

架构

自然语言翻译一般使用 encoder-decoder 模型。

encoder 就是将一个序列 $(x_1,x_2,…,x_n)$ 转为计算机可以识别的 Tensor 序列 $ z = (z_1,z_2,…,z_n)$。比如 $x_t$ 是一个英文单词,$z_t$ 就是和该单词关联的一个向量。

decoder 就是将 encoder 的输出作为其输入,转为 $(y_1,y_2,…,y_m)$ 序列。长度可能和输入序列不同,比如英文翻译为中文。

这样的模型也是 auto-regressive(自回归) 的。每一次的输出都会使用之前的输出作为输入。

alt text

本架构图中,左边是 encoder ,右边是 decoder。注意到右边下面的 output 作为输入说明它是 auto-regressive 的。Nx 表示方框内的同样的层有 N 个。“Add & Norm” 层有一个输入是跳过它之前的层的,这个就是残差连接(residual connection)。

Encoder

简化的表示形式是 $LayerNorm(x + Sublayer(x))$。比如第一部分,Sublayer 就是 Multi-Head Attiontion 层,加上 x 表示是残差连接需要原始数据作为输入,这两个参数作为输入传递到 (Add & Norm) 层,上面一部分也是同理。

其中的 Normalization(标准化) 使用的是 LayerNorm 而不是 BatchNorm。

Normalization(标准化)是让一组数据减去其均值再除以标准差得到的,得到的结果就是均值为 0 ,标准差为 1 的一组数据。通过这种方式让每一层的输出都相对标准化,类似于应用间进行数据通信的 API。

BatchNorm 是将一个 batch 内不同样本间的同一个的特征进行标准化,而 LayerNorm 是将一个 batch 内同一个样本的所有特征进行标准化。

Decoder

decoder 相比于 Encoder 增加了一个 Masked Multi-Head Attention 层,该层用于实现 auto-regressive ,在推理阶段会用于处理之前的输出结果来作为输入。

在训练时,Decoder 的输入是训练集的数据(teacher forcing),而不是 Encoder 的输出,我们希望让其模拟在推理时将产生的输出作为输入的过程,所以增加一个 Mask 去遮盖后面的部分,避免其看到完整的训练集数据。

关于当前大模型使用的变种架构的情况如下,基本上都偏向于 Decoder-only 架构:

alt text

Attention

注意力函数需要一个 query 和 一组 key-value 作为输入,它会根据 query 和 key 的相似度作为对应的 value 的权重。不同的注意力函数的区别就是计算相似度的相似函数

打个比方:

学生要借书:

  • query: 要查询的书籍的种类(物理学,入门,中文)
  • key:书的种类(<生物学,入门>,<物理学,中文>,<化学>)
  • value:书籍(《生物学入门》,《物理学中文导论》,《化学导论》)
  • 得到的权重:0.4,0.6,0

最后这位学生借走了《物理学中文导论》,当然此时<物理学>这个 query 的重要程度比较大,他本身就是学物理的,即使有一本生物书是入门且中文的,他也不会考虑借那本书。

在上面的模型图中, 下面的两个 Attention 层有三个输入,也就是 query、key、value,在 self-attention (自注意力)机制中这三个参数其实是原始输入序列的三份相同的复制。此时对于 query 中的一项,也就是原始序列的一个词,计算注意力函数时它和 key 中相同位置的自身相似度肯定最高,权重也最高,然后如果序列中还有其他相似度较高的词,也会被考虑进来(较高权重),最后每个 value 被附加对应的权重,实现了为序列中的每一个词附加同一个序列中其他它感兴趣的部分的信息。

公式如下:

\[Attention(Q,K,V)=softmax({QK^T\over \sqrt{d_k}})V\]

权重指的其实就是“correlation(相关性)”,Attention 计算权重的方式就是点乘,也就是将 query 和 key 矩阵做点乘($QK^T$)。从本质上来说点乘是计算两个矩阵 correlation 的一种方式,还有其他很多方式,只不过 Transformer 综合考虑选择了点乘。这种关系和它们之间的距离形成反比,也就是距离越大,联系也会越小,所以除以了 d,也就是距离。softmax 将权重进行了归一化,让其分布在 (0,1) 之间且和为 1,方便计算。

架构图中上面的 Attention 层,它的 key ,value 来自于 Encoder 层,query 来自于下面的 Masked Multi-Head Attention 层。Encoder 和 Decoder 沟通的桥梁就是这个层。

Multi-Head Attention

将 query 、key 和 value 都投影到各自的 h 个低维空间中,再进行 Attention 操作,可以增加可学习的参数数量(增加了 3*h)。投影的概念就是降低维度,比如三维物体的三视图就是降低一个维度的投影,通过 3 个二维图来展现三维物体。

通过在不同的维度分别计算 Attention 后再进行组合,有什么好处?首先以三视图举例子,在三维物体中的一个点,投影到三个平面上,会出现 3 个点,原来我们是去三维空间找两点(类比 query 和 key)的距离信息(类比 attention 计算 value 的权重的过程),现在是在三个平面上分别去找这 6 个点的距离信息($query_x$和$key_x$,$query_y$和$key_y$,$query_z$和$key_z$),而且我们会为每个平面分配权重,这样假设 y 平面上的两个点($query_y$和$key_y$)距离很远,而在其他平面上距离较近,但因为我们并不关心 y 平面也就是其权重较小,那么这样计算出来的两个点间的最终距离就会很小,符合我们的预期。换个说法就是分成多个通道去找特征。

Position-wise Feed-Forward Networks

Feed-Forward 层实际就是一个 MLP,它对输入序列的每个词单独进行 MLP 的计算。因为其输入是 attention 层的输出,而通过 attention 层的处理,每个词都已经被附加了整个原始序列中它感兴趣的部分(与之相关的部分),所以不需要像 RNN 一样把前一个词的 MLP 计算结果作为输入,只要每个词单独输入就行,没有任何依赖。

现有的 LLM 主要参数都在 Feed-Forward 层中,大模型之所以强大,就是将大量的知识都保存到了这个层中。而计算量主要集中在 Attention 层。当然不能忽略大参数量带来的内存拷贝问题,CPU或GPU 从内存读取数据到 cache 处理也是非常耗时的一个过程,所以 FFN 即使计算量不大,在训练和推理中的耗时也是较高的。随着大模型参数的增加,数据传输时间已经成为实际瓶颈。

speculative decoding

而且因为推理时只能一个一个 token 生成,不能并行处理,进一步降低了速度。为了提高推理速度,有一种思路是利用一个小模型先推理一遍,因为参数小,推理也很快,但质量不太行,这时候再让大模型进行 check,这个 check 过程和训练很像,都是已经有完整的输入后进行的,可以并行处理提高效率,check 主要做的事就是检查每个 token 和自己预测的是否相同,发现错误时就主动从错误的地方再重新生成一遍。因为可以并行,理想情况下,大模型 check 过程和生成一个 token 的时间就是一样的。

如果运气好,小模型输出就是正确的,这就大幅提高了效率。

如果小模型输出到中间开始后面全错了,大模型也会从错误的地方开始重新生成,然后让小模型从这里继续,不断交替,总时间有很大概率也是减少的。

alt text

  • 绿色:小模型正确的
  • 红色:小模型错误的
  • 蓝色:大模型

当然这个小模型其实可以塞到大模型中,成为其独立的 heads,这样达到的效果也是一样的。

这种方式我们可以称为传统的 speculative decoding,后面出现了 MTP 技术,从训练阶段就考虑到了多 token 预测,效果也会更好。

Embeddings and Softmax

Embedding 就是将输入数据映射成计算机可以进行计算的向量,这里固定每个词映射成 512 长度的向量。

Positional Encoding

attention 机制带来的一个问题是其注意力信息是位置无关的,因为权重信息中没有位置信息,这也导致在翻译过程中将输入的词任意打乱,其注意力信息也是相同的,即使语义已经发生了很大的变化。

此时该模型在输入中加入了位置信息来解决该问题。

应用:预测

Transformer 可以用来做 token 预测,此时只需要 Decoder 就行,即 Decoder-only

我们为输入

I love you

添加一个 <s> 标签(right shift 一个 token):

<s> I love you

将输出设置为原始的输入:

I love you

然后就能从 <s> 推理出 I,从 I 推理出 love。得益于 Attention 机制的并行性,训练过程中,一个句子的每个单词的推理可以并行执行(因为训练中输入是确定的,不是从上一次的输出得到的),然后统一计算 loss 并回归。

为了实现预测,我们不能将输入中该单词之后的单词信息提供给注意力机制,比如 love 的注意力只能看到 I<s> ,但不能看到 you,所以训练时在处理每个单词的注意力时会给句子添加对应的掩码。

推理和训练过程不同,必须一个一个 token 生成,生成完的 token 作为下一次的输入,所以并行性并不如训练过程。比如推理时由 I 预测了 have,然后由 I have 预测了 a ,这样下去导致整个句子和训练过程就完全不同了,一个 token 预测错,后面就全错了,这也是预测这个应用的一个问题。

应用:翻译

Transformer 可以用来做机器翻译,这也是 Google 提出 transformer 时想要其解决的主要问题。这会利用到 Encoder-Decoder 架构。

参考

本文由作者按照 CC BY-SA 4.0 进行授权