大模型学习-基础篇（一）：Attention复杂度计算和PE介绍

系列文章目录

例如：第一章 大模型学习-基础知识

大模型学习-基础篇
大模型学习-实践篇

---

前言

此系列记录大模型学习的经验总结，包括且不限于：LLM相关的知识、理论、论文、部署、微调、量化等。

一、LLM是什么？

大规模语言模型是一种基于transformer架构的语言模型，可以完成自然语言处理中的各种任务，解决相关的问题。接下来简单介绍一下Transformer架构。

二、Transformer架构

Transformer是google的团队在2017年提出的一种seq2seq模型，模型结构主要包括encode模块和docoder模块。其中encoder模块包括多头注意力机制、前馈神经网络和残差连接，decoder模块则包括掩码多头注意力机制、多头注意力机制、前馈神经网络和残差连接。该框架提出后，经过各种改进和升级，目前已经成为LLM的基础架构，是我们学习LLM之前必须回顾的知识内容。关于transformer的详细介绍，可以参考该博客，作者使用直观的图片和公式，解释了transformer的各个模块的构成和功能。

这里重点介绍的两个关键知识点，multi head attention（MHA）和positional encoding（PE）。

MHA（Multi Head Attention）

Attention机制的核心在于Q,K,V三个向量序列的交互和融合，其中Q,K的交互给出了两两向量之间的某种相关度（权重），而最后的输出序列则是把V按照权重求和得到的。

$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$
时间复杂度：Transformers模型的时间复杂度主要取决于输入序列的长度N和模型中隐藏层的数量H。对于一个具有L个层的Transformer模型，其时间复杂度为$O(L{N}^{2}H)$，其中$N^2$来自于注意力机制的计算。因此，对于较长的输入序列和更深的模型，Transformer的时间复杂度可能会非常高。

Self Attention: $Q,K,V$同源，都是同一个$X$经过线性层变换后形成的矩阵。并且能够直接捕捉$X$中任意两个向量的关联，而且易于并行，这是Self Attention的优点。从理论上来讲，Self Attention的计算时间和显存占用量都是$O(n^2\ast d)$，其中$n$是$minibatch$中的$sequence$最大长度。

计算过程主要包括以下步骤：输入的线性映射、相似度计算、softmax和加权求和。
1）线性映射：两个矩阵$(n,d)\ast (d,d)$相乘，时间复杂度为$O(n{d}^2)$；
2）相似度计算$Q{K}^T$：两个矩阵$(n,d)\ast (d,n)$相乘，$O(n^{2}d)$；
3）$softmax(Q{K}^T)$：对每行进行$softmax$,每行的复杂度是$O(n)$，n行的时间复杂度$O(n^2)$；
4）加权求和$Q{K}^T\ast V$：两个矩阵$(n,n)\ast (n,d)$相乘，时间复杂度为$O(n^{2}d)$。
我们将各个结果相加以后，只保留主导项，那么时间复杂度为$O(n^2\ast d)$

MHA则是在$d$维子空间上分别进行自注意力机制的计算，假设共有$k$个头，每个头的维度数为$d_k$那么时间复杂度为：
$O(n\ast d^2)+O(k\ast n^2\ast d_k)+O(n^2)+O(n^2\ast d)$，我们只保留主导项的话，最终结果就是$O(k\ast n^2\ast d_k)$，其中$k\ast d_k=d$，所以最终结果还是$O(n^2\ast d)$。

---

如果更详细一点的话，实际上计算多头自注意力的方式是将形状为$(n,d)$的$Q,K,V$矩阵全部拆分为$(n,k,d_k)$，然后维度变换成$(k,n,d_k)$，再将$(k,n,d_k)$转置成$(k,d_k,n)$后与自身相乘，得到矩阵$(k,n,n)$，时间复杂度为$O(k\ast n\ast n\ast d_k)$。然后经过softmax的时间复杂度为$O(k\ast n^2)$，再乘以形状$(k,n,d_k)$的$V$矩阵，最后得到$(k,n,d_k)$，再将0维度和1维度交换，然后融合1维度和2维度，最后矩阵形状是$(n,d)$

---

如果有兴趣的话，可以列出这里面每一步计算的复杂度，然后求和。

正好介绍到这里，我对矩阵乘法的时间复杂度计算有一个经验总结。如果两个矩阵M和N相乘，最后得到的矩阵形状为(A,B,C)，那么时间复杂度就是ABC再乘以M和N相乘消掉的那一维度的长度N。
例如：$(M,N,H)\ast (M,H,N)=(M,N,N)$，消掉了$H$。所以时间复杂度为：$O(M\ast N\ast N\ast H)=O(M{N}^{2}H)$

Cross Attention：K/V不同源，可能来自于上一层的block输出，也可能来自同一层的其他并行block。

针对attention机制的复杂度进行优化，一直是学术界和企业界在共同努力的方向。尤其是在大语言模型中，由于上下文长度极大，并且存在不断变化的可能性，所以需要对attention机制进行极致的优化，从而节省计算开销，并且提高性能。接下来，我们对各种改进的attention机制进行简单的介绍：
（1）Atrous Self Attention（膨胀自注意力”、“空洞自注意力”、“带孔自注意力”）
原理：启发于“膨胀卷积（Atrous Convolution）”，对token间的相关性进行了约束，强行要求每个元素只跟它相对距离为k,2k,3k,…的元素关联，其中k>1，是预先设定的超参数。从下左的注意力矩阵看，就是强行要求相对距离不是k的倍数的注意力为0（白色是0）：

复杂度/性能：每个元素只跟大约$\frac{n}{k}$个元素算相关性，这样一来理想情况下运行效率和显存占用都变成了$O(\frac{n^2}{k})$。
效果：可能减少模型的过拟合，降低计算复杂度。但是会造成上下文信息缺失。

(2) Local Self Attention (局部自注意力)
原理：约束每个元素，使其只与自身的前后k个元素以及自身有关联。

复杂度/性能：是将相对距离超过k的注意力值都直接设为0。每个元素只跟$2k+1$个元素算相关性，这样一来理想情况下运行效率和显存占用都变成了$O((2k+1)n)\sim O(kn)$。
效果：显著降低计算性能，达到线性复杂度。但是直接丧失了长程的Token关联信息，削弱了模型捕捉超长上下文关系的能力。

（3）Sparse Self Attention （稀疏注意力）
原理：将局部注意力和空洞注意力结合结合起来，取长补短。理论上也可以学习到全局关联性，也省了显存。

复杂度/性能：将相对距离不超过k的注意力值保留，相对距离为k,2k,3k,…的注意力值都保留，其余设为0。如果我们把稀疏注意力想象成从n个token做点积，变成从n个token中取一个子集k做点积，那么$O(n^2\ast d)$中的一个$n$就会变成$k$，即时间复杂度变为：$O(k\ast d\ast n)$。根据我们对k的采样在$1到n$之间，可以算得总时间复杂度。

效果：是前两种注意力机制思路的一种折中办法，在性能和成本上取了平均。

苏神的代码实现：github链接。

PE（Positional Encoding）

位置编码是transforer中相当重要的一个模块，如果没有位置编码，那么由于Transformer的注意力机制是每个token都在进行超越上下文的计算，实际上每个token的位置都是一样的，那么模型将无法知道token之间的相对位置关系。因此我们需要设计一种编码，标记每个token的序列位置。而这个位置编码有两种设计思路：
（1）绝对位置编码：将位置信息以向量的形式，融入到句子token的嵌入表示中。
（2）相对位置编码：通过改进attention机制，实现一种能够感知相对位置关系的注意力机制。

绝对位置编码

（1）训练式：增加一个$(n,d)$的矩阵，作为句子token的位置编码信息，随着模型训练一起更新。
特点：简单有效，但是没有外推性（无法应对上下文长度不断变换/增加的场景），一种改进方案是层次分解。
（2）三角式（Sinusoidal位置编码）：
$$p(k,2i)=sin(\frac{k}{10000^{2i/d}})$$
$$p(k,2i+1)=cos(\frac{k}{10000^{2i/d}})$$
其中，$k$代表token在句子中的位置，2i代表对应token的偶数维度，2i+1代表对应token的奇数维度。由于正弦和余弦三角函数的和差化积公式，我们可以得到不同位置token之间的相对位置，这可以解释该编码提供的相对位置表示能力。另外，三角函数式位置编码的特点是有显式的生成规律，因此可以期望它存在一定的外推性。
（3）递归式：
RNN模型在结构上就自带了学习到位置信息的可能性，因为递归就意味着可以训练一个“数数”模型，因此，如果在输入后面先接一层RNN，然后再接Transformer，那么理论上就不需要加位置编码了。同理，我们也可以用RNN模型来学习一种绝对位置编码，比如从一个向量$p_0$出发，通过递归格式$P_{k+1}=f(P_k)$来得到各个位置的编码向量。
（4）相乘式：
默认选择相加是因为向量的相加具有比较鲜明的几何意义，但是对于深度学习模型来说，这种几何意义其实没有什么实际的价值。将“加”换成“乘”，也就是$x_k\otimes p_k$的方式，取代$x_k+p_k$的形式。

相对位置编码

（1）经典式
（2）XLNET式
（3）T5式
（4）DeBERTa式
（5）其他位置编码
这里的内容，此节暂不做介绍。因为后面我们将介绍RoPE旋转位置编码，这是目前大多数大模型使用的位置编码，我们需要重点理解和研究。此处强烈推荐苏神的公式解释，我看了有一种顿悟的感觉，把注意力机制一下子理解透了，耐心把每一行矩阵的变换看懂，你会突然发现自己对注意力机制的理解深入了很多。

总结

此章节介绍了Transformers库的结构和相关模块，主要参考了苏神的博客，transformers进阶之路，由于原博客较长，所以我这里精简了一下，并且补充了一些简单的内容。