
【斯坦福CS231n 学习笔记5】Attention! Transformer!
学习了Attention与Transformer。
RNN 与 Seq 2 Seq 模型的局限性
在机器翻译等 Sequence-to-Sequence 任务中,传统的做法是使用一个 RNN 编码器和一个 RNN 解码器,编码器处理完整个输入序列后,需要将所有信息压缩进最后一个隐藏状态 ,并且将其作为一个全局的语义摘要,传递给解码器,这个摘要也被称为 Context Vector c,编码器再以这个 context vector 作为初始条件,逐步生成输出序列:。
但是问题就在于,这个 context vector c 长度实际上是固定的,无论输入序列是 5 个词,还是 50 个词,还是 1000 个词,context vector c 的维度是固定的,这意味着编码器在被迫使用一个固定长度的实向量,去承载一个长度可变、结构复杂、语义分布不均匀的序列。当输入序列很短时,假设还勉强成立,但是一旦序列变长,编码器就需要权衡哪些信息必须被保留,哪些信息可以被压缩、弱化甚至丢弃,哪些早期输入对输出可能有用,但是在时间上传播路径又比较长。而在反向传播时,长序列末端的误差也难以有效传导回序列开头的 Token。
从数学和优化的角度看,实际上是一个信息瓶颈问题,有限的容量不可能容纳无限的信息,因此这个建模假设本身就是有问题的。
解决的方法就是,编码器在生成第 个输出词时重新查看整个输入序列。
Attention
从 RNN 到 Attention
不同的输出应该关注不同的输入,如图中的案例:“we see the sky” → “vediamo il cielo”,当解码器需要生成 vediamo 时,最关心的应该是输入中的“we see”,当其要输出“cielo”的时候,关注点应该转向“sky”。输出序列的不同位置,对输入序列的关注分布本来就应该不同。而 RNN Seq2Seq 的固定 context vector 则将这种差异性给抹平了。
要关注不同的输入,就不能只用 ,而是保留所有的 ,在标准的 Seq2Seq 中,编码器的中间状态实际上一直存在,只是从前选择性忽略了它们,而只保留了最后一个,Attention 的第一步,就是保留所有的编码器隐藏层序列,作为一个可查询的记忆库。那么在解码器第 步的时候,应该关注哪一个 就成了一个问题。这就是 alignment 问题,Attention 的技术核心,就是是用一个可学习的方式来解决“对齐”。
在 Bahdanau Attention 中,定义了一个对齐打分函数 ,这里的 通常只是一个小型前馈网络,例如线性层加非线性,其作用是输出一个标量,用于衡量当前的解码器状态与第 个输入位置的匹配程度,这个标量可以称之为“相对相关性分数”。
在有了这一组分数 之后,需要得到一个可以解释为“关注程度”的量,并且其在所有输入位置上加起来和为 1,故对 做 softmax:
这一步将任意实数都映射到了 ,而后通过引入竞争关系得到一个可解释的注意力分布。Context vector 也不再是一个常量,而是由解码器的状态动态决定的一个加权后的视图:
解码器的更新公式也变成了:
解码器不再接受一个固定的摘要,而是主动决定自己此刻需要输入哪些信息。
Attention 没有任何显式监督,所有的对齐结构都是通过最终翻译损失反向传播自动学出来的,并且完全可微分、可学习。
通用Attention 计算结构
Attention 实际上并不依赖于 RNN 的递归结构本身,本质上是一个查询检索过程,真正依赖的只有三样东西:
- Query(Q):当前关注的向量(例如解码器当前的状态)。
- Key(K):被查询的向量集合的索引特征(例如编码器的各个隐藏状态)。
- Value(V):被查询的向量集合的实际内容(通常与 Key 相同,即编码器的隐藏状态)。
假设给定一个向量 ,以及一组向量 ,想要从中读出一条信息,Attention 可以描述为一个三阶段过程:
- 相关性打分:,这里的 可以是任何可微函数,如点积、MLP、线性变换后的内积。
- 归一化:softmax,, 这一步将相关性变为了注意力分布,并且引入了竞争关系,不是所有输入都同等重要。
- 加权聚合:,输出 是对输入集合的一个内容感知的摘要,而不是简单平均。
从建模的角度来看,这个 Attention 算子已经完全不再关心输入是不是序列、输入有无时间顺序、输入的模态等,而是根据 query 的需要,从一组向量中提取最相关的信息,在这个通用形式下,query 可以是任何东西,这也是为什么现在 Attention 能够被移植到许多任务中。
Attention Layer
在通用 Attention 计算结构中,已经将 Attention 抽象为一个三步算子:相关性打分、归一化、加权聚合。在概念上已经十分完备,但是在工程上还面临具体的相似度函数应该如何选择、输入向量在 Attention 中所扮演的角色、计算在高维空间中是否数值稳定三个问题。
在Bahdanau Attention 中,对齐分数是通过一个小 MLP 计算出来的,这在表达能力上没有问题,但是计算效率不高,且不利于大规模矩阵并行,因此提出了直接用向量内积来衡量相似度,对齐分数被简化为:
这意味着 Attention 完全可以被写成矩阵乘法,从而天然适配 GPU。
在通用 Attention 结构中,输入向量 同时承担了用于参加相似度计算和用来作为被聚合的内容的角色,而在这里将这两种角色显式分开,引入 Key(K)和 Value(V),其中 Key 是用于被匹配的向量,而 Value 是被读取的向量在形式上,通过线性映射分别得到:,,与此同时,Query 也通过线性映射得到:。其中 Query 决定当前在找什么,Key 决定能被如何匹配,Value 决定最终提供什么信息。现在得到了一个较好的 Attention 形式:
但是当 Key 和 Query 的维度 较大时,点积 的数值分布会随着维度增长而显著变大,这直接导致 softmax 的输入幅值过大,从而进入饱和区,导致梯度消失。
因此还需要引入一个修正项 ,使得 ,这样在统计意义上能够将点积的方差控制在一个稳定的范围内,防止 softman 过早进入极端分布,并且保证梯度在训练早期是可用的。
最终得到的 Attention 形式如下:
可以看到其完全由矩阵运算构成,可以并行计算所有的 Query,不依赖于序列递归结构。
Self-Attention
在 Self-Attention 中,不再区分 Query、Key 和 Value,而是直接对同一组输入做三次线性映射:、 、。现在每一个 token 既可以作为 Query 去查询其他 token,也可以作为 Key/Value,被其他 token 读取,即序列中的每一个位置,都可以直接与所有其它位置建立联系。
Attention 的计算仍然是
但是此处的 不再表示当前时间步与过去时间步之间的关系,而是表示任意 token i 与任意 token j 之间的相关性,可以将其理解为一个 token-token 关系矩阵。
在这个 Self-Attention 中,任意两个 token 之间的交互,路径长度恒为 1,这意味着长程依赖不再是难点,序列依赖只影响计算量,而不影响信息传播能力,模型的表达能力来自于关系建模,而非记忆长度。
但是 Self-Attention 本身不包含序列的顺序信息,如果打乱输入 的行顺序,输出也仅仅是对应行的顺序变化,数值本身并不变,导致其无法区分“A hit B”和 “B hit A”,该性质也被称为置换不变性(Permutation Invariant)。为了解决置换不变性带来的语序丢失问题,必须在输入嵌入中显式加入位置信息,直接将位置向量 加入到输入 上:,位置向量可以是预设的或可学习的参数。
Masked Self-Attention Layer
在用于序列生成任务(如机器翻译或文本生成)的训练阶段,为了并行计算,通常会将整个标准答案一次性输入给解码器,如果不做任何限制,当解码器在 时刻尝试预测 时刻的输出时,Self-Attention 机制就允许其“看到”输入序列 之后的所有位置,导致模型学不到任何东西。
解决方案就是在计算注意力分数时引入一个掩码矩阵,将 时刻之后的所有位置的注意力分数设置为 。
其中 在上三角区域为 ,其余为 0, 经过 后, 变为 0。这意味着模型在时刻 只能够关注 极其之前的时刻。
Transformer
Multiheaded Self-Attention Layer
在 Self-Attention 中,所有的关系都是在同一个表示子空间里建模的,即 Q、K、V 都是一次线性投影,所有的 token-token 关系,都通过同一个注意力分布来表达。但是在自然语言(或序列)中的关系本来就是多种多样的,例如语法关系、语义相似性、指代关系、局部和全局依赖,如果将这些关系全都放进一个 Attention 权重矩阵中,模型要么变得表达混乱,要么需要在不同的关系之间做出妥协。
所以不能只做一次 Attention,而是需要并行地做多次 Attention,但是这里的多次并不是重复同一件事,而是在不同的表示子空间中使用不同的 QKV 投影各自独立地学习注意力模式。
假设有 个 attention heads,对于同一个输入 , 我们做如下操作:对第 个 head:,,,然后分别计算 ,每个 head 的维度都更小,且具有独立的注意力分布,这样能够拆分表示空间并且关注完全不同的 token 关系。
而在得到多个 heads 之后,Transformer 并不是将其简单平均,而是将其拼接:
这是因为不同的 attention heads 是并列的信息通道,并不存在竞争关系,这能够让 transformer 同时建模多种依赖结构,而不需要人为指定哪一种关系更加重要。
Self-Attention is Four Matrix Multiplies
这里主要提出了 Self-Attention 模块中计算量最大、最核心的四个步骤。
分别是 QKV 投影、计算相似度、数值加权以及输出投影,在没有优化的情况下,时间复杂度和空间复杂度都是 ,而 2022 年提出的 Flash Attention 通过 IO-Awareness并行计算相似度和数值加权,使得空间复杂度变为了 ,最终加快了计算。
Three Ways of Processing Sequences
到目前为止已经学习了 RNN、CNN 和 Self-Attention 三种处理序列的方法,其各自拥有优缺点:
- RNN:内存与计算效率高。无论序列多长,它只需要一套参数反复使用,计算量和内存占用通常是 (线性复杂度),但是无法并行,需要计算隐藏状态。
- CNN:能够完全并行,但是对于长序列的支持较差,且需要堆叠多层以增加感受野。
- Self-Attention:长序列支持友好,每个输出直接依赖于全部输入,且计算高度并行,但是计算较为昂贵,需要 的时间复杂度与至少 的空间复杂度。
最终得到的结论就是:Attention is All You Need!
Transformer Block
输入与输出都是一组向量,所有的输入向量通过(Multiheaded) Self-Attention 来进行信息交换;之后是一个残差连接,这能够让梯度在反向传播时无损通过,防止深层网络退化;再然后是一个层归一化,用于稳定每层神经元的分布,加速收敛;在 Self-Attention 完成信息的聚合之后,MLP 负责信息的加工,由一个简单的前馈神经网络(通常包含两层线性变换和一个激活函数),独立地作用于每一个位置的向量,之后又是一个残差连接,最后进行层归一化输出内容。
一个 Transformer 实际上就是若干个 Transformer Block 的堆叠,随着堆叠的数量越来越多,模型的能力也越来越强(例如 GPT 系列,从一开始的 12 blocks 到 GPT-3 的 96 blocks)。
Transformer应用
Transformers for Language Modeling(LLM)
语言模型的训练目标是给定前面的 token:,预测下一个 token , 对于一个语言模型,最开始是通过一个嵌入矩阵将离散的词汇映射成连续空间向量中的点,在模型内部使用 masked attention 来处理,将模型约束在过去的输入,通过多个 Transformer block 最终在投影矩阵中将上下文感知的 token 表示投影回词表空间,而后再街上 softmax 来得到概率分布。
Vision Transformers (ViT)
Lecture 8 举的案例是一个图像分类任务。由于图像是一个规则的 2D 网络加上通道结构,在形式上并不是一个 token 序列,无法直接将其输入给 transformer,因此 Vision Transformers 所做的事情主要是预处理部分,将图片分割为若干个 patches,将每一个 patch 作为一个视觉 token,每个 patch 是一个 维的向量,然后通过 flatten 进行对齐,作为输入向量输入到 Transformer 模块中,需要注意的是,还需要通过位置编码来补充图片的结构信息,让 Self-Attention 知道图片的空间关系。最后的输出是每个 patch 一个向量,ViT 对所有 patch 的输出向量做平均,得到一个全局表示,再通过线性层来预测类别。
Tweaking Transformers
Pre-Norm Transformer
在原始的版本中,归一化层是在残差连接之后的,也成为 Post-Norm,这样容易导致梯度在反向传播中爆炸或消失,导致训练很难收敛,而 Pre-Norm 所做的改进就是将归一化层放到 Self-Attention 与 MLP 之前,这样可以使其训练更加稳定。
RMSNorm
原始的 LayerNorm 是先减去均值,再除以标准差,对数值重新进行了中心化。但是 RMSNorm 认为这种重新中心化是不必要的,只需要调整数据的缩放即可,结果是计算量更小,速度更快,且效果与 LayerNorm 持平甚至要更好:
SwiGLU
原始的 Transformer 使用 ReLU 或 GELU 激活函数,SwiGLU 结合了 Swish 激活函数和 GLU (门控线性单元)。通过引入门控机制,不仅对数据做非线性变换,还多学习了一个门来控制多少信息可以通过,虽然参数量略微增加,但是能够显著提升模型的推理性能和收敛速度。
Mixture of Experts(MOE)
这是实现超大参数规模的关键技术,假设模型有 E 个专家,但在处理每一个 Token 时,并不激活所有的专家,而是通过一个路由,只选择其中的 A 个最擅长的专家来计算,这能够极大解耦参数量和计算量,模型可以有超大的参数,但是在推理时,只需要相对较少的计算量。
参考如下:
https://cs231n.stanford.edu/index.html
