RoPE

最早的注意力架构，位置编码采用的是余弦编码，也就是加上一个三角函数，波长和token位置成正比。

也有一些比较简单的模型如GPT3之前，使用绝对位置编码，第i个token就加上一个固定值，表示第i个。

Llama使用了旋转位置编码，并且由于Llama是第一个性能能对齐闭源模型的开源模型，许多后来的开源和闭源模型都借鉴了它的设计，包括RoPE。

一个好的位置编码应该具有相对位置的特性，也就是两个不同token的词向量，计算余弦相似度应该只和他们的相对位置有关，和绝对位置以及其他值无关。

余弦编码和绝对位置编码计算余弦相似度，都不止包含相对位置信息，还包含别的信息。

旋转位置编码的思想是，把词向量看成向量，然后位置信息用向量旋转的角度表示，如下图，know原本是一个词向量，当他在1位置的时候，旋转一个单位，在3位置的时候旋转三个单位。旋转位置编码最大的特点就是他不是加到词向量上的，他是对词向量进行旋转。由于是旋转，所以做余弦相似度计算时天然算出来的就是两个词向量的夹角，这就满足了只和相对位置有关。

上面的旋转用的是二维词向量，如果更高维怎么办？可以把高维向量两两配对，看成多个二维向量，分别旋转。如下图一个词向量的1，2两个位置配对，3,4两个位置配对，以此类推。

那么为了实现这样的两两配对旋转，位置编码矩阵如下，按$2\times 2$分块，只有对角线上的块有效。当然这个矩阵形式是数学上的定义，实际实现上，由于这个矩阵大部分都是0，使用这个矩阵做计算是很浪费的，实际做法是转化成逐元素操作，给每个位置先乘上对应的三角函数，然后再线性组合。

线性层维度

从这一节开始进入模型超参数设置。

线性层会把词向量映射到一个更大的维度，做激活函数，再映射回来。这个更大维度具体是多少，则要看模型的设计。

这里其实存在规律，一般而言，会设计成
$$d_{ffn}=4d_{model}$$
可以看到大多数模型都在这个比例附近

实验表明过大的比例会导致loss反而上升，可能是出现了过拟合。并且过大的中间维度计算量也会成倍增长。权衡之下，选择2.6到4之间是最好的。

多头注意力头数

多头注意力是为了让不同的头学习不同特征引入的，如下图把原始qkv向量都分别作切分，得到多组，每组是一个头，单独计算注意力，最后再把所有头拼接起来，恢复到初始维度。如下图演示中$numheads\times headdim=modeldim$，这看起来是理所当然的，但实际上不是必须实现的。

实际上可以不相等，甚至可以远大于。只要最后再加一个线性层，把输出从$numheads\times headdim$映射到$modeldim$就行。

模型长宽比

定义是模型词向量维度比上模型层数，也就是模型宽度比上模型长度。可以发现大多数模型的比值都在100-150这个区间，可以看到更倾向于增加模型宽度而不是深度。

这样的设计首先考虑模型并行化的难度。如果模型变宽，每一层的计算量变大了，一般采用张量并行，把一层的张量切成多份给不同GPU，分别计算再拼起来。这是容易实现的。

如果模型变长，一般采用流水线并行，如下图，把模型看成一个流水线，不同层看成不同阶段，分配给不同GPU。流水线并行实际上更难实现，为了避免流水线气泡需要复杂的调度设计，而张量并行只要一个all-reduce就能实现。

从模型能力上来看，如下图，纵轴是loss，横轴是比值，最优的比值也是在100-150左右。这印证了前面的表格数据。

再看下图，左边两个标着DM，表示增加模型词向量维度的实验，右侧两个标着NL，表示增加模型层数的实验。可以看到增加DM，模型性能基本是线性增长的，而增加NL，模型性能的增长是越来越慢的，也就是增加宽度的边际效益更高。这解释了为什么最优情况下宽度比长度大得多。

词典大小

只处理一种语言的模型，词典大小基本3-5万之间，多语言模型，词典大小基本在10-25万之间。对此虽然没有实验，但这些公司的团队在实际训练时肯定也是做了计算量和表达能力的取舍。

正则化

如下图是主流模型的dropout和参数衰减比例。可以看到早期模型更多使用dropout，新模型更多使用参数衰减。

但现在的正则化的作用已经不是防止过拟合了，如下图，左侧第一张表示不同的参数衰减比例，loss是接近的。现在正则化的作用更多影响模型收敛速度和最终收敛位置。最右侧的图展示了固定学习率下，增加参数衰减会拖慢loss收敛，中间图则展示了在余弦退火学习率调度下，参数衰减能到达更低的收敛位置。

Z-loss

从这一小节开始进入训练稳定性。

训练稳定性当然也是重要的，如下图，蓝色虽然loss更低，但梯度和loss经常出现尖峰，这对训练不利，可能在训练途中出现梯度爆炸，这次训练就作废了。

训练中一个不稳定的来源是softmax，由于有指数项，可能会出现有的项突然变大，导致梯度尖峰。一个解决办法是，在loss项里增加一个${\log }^2(\sum e^x)$项，这样可以惩罚过大的指数。

有一个小问题是，前面我们在infra里讨论过flash attn的safe softmax，也是为了避免softmax的输出过大，给所有指数都先减去最大值，那为什么还需要zl-oss?答案是safe softmax只是为了防止计算过程中指数过大造成的数值溢出，最后计算出的结果和原始softmax是完全相等的，也就是说这只是一种数值技巧。

而z-loss解决的是输出值不均匀的问题，比如最终输出logits，有一个token概率接近于100%，其他位置都是0，此时交叉熵会极小，导致模型几乎停止更新，也就是梯度消失。同时这样的不均匀logits也可能导致反向传播时的梯度爆炸。这个情况safe softmax是不管的，必须依靠z-loss来调整各个token的概率。

QK norm

归一化是为了解决模型输出不稳定，已经在注意力和线性层使用了，而模型计算中另一个容易出现数值不稳定的地方是$Q{K}^T$相乘。于是一个自然的想法是，能不能给这里也加一个归一化，参考之前的经验，加的是前置归一化，如下图。

Logits soft-capping

出于和z-loss类似的动机，为了使logits输出更均匀，更稳定。另一个更温和的做法是，对logits输出在增加一层修改，先做一个整除，经过tanh，最后再乘上之前的除数。由于双曲正切是有上下界$[-1,1]$的，最后值域会被限制在$[-cap,cap]$之间，且经过tanh变得更平滑。

多头注意力变体

最后来介绍一些多头注意力的变体。

先来分析一下标准多头注意力MHA的计算强度，对于prefill分析如下图。只要k,b,n够大就能达到较高的计算强度，打满GPU。

分析一下具体计算，注意力得分计算，每一个batch是(n,d)*(d,n)，考虑到batch，总计算量$O(b{n}^{2}d)$，最后输出前有一个映射，前面讨论多头注意力头数时，提到过，头数乘上每个头的维度，是可以不等于模型词向量维度的，需要经过一个线性层映射，每个batch是一个(n,d)乘上一个(d,d)，复杂度$O(bn{d}^2)$。由于这里规定了$n\le d$，所以计算复杂度的主项是$O(bn{d}^2)$

对于decode阶段如下图。每次一个token的q和前面所有token的kv计算，前面token的kv每次都一样，可以保存下来，直接读取。

计算量上仍然是$O(bn{d}^2)$，因为decode虽然每次只生成一个token，但是总计算量还是一样的。内存访问量则变多了，因为每次生成一个token，都需要读取原来的全部kv cache，这样原来的内存访问量，除了激活值$X$只读一次，其余部分都需要乘上序列长度$n$。由于$X$不乘$n$，小了一个量级，内存复杂度上可以忽略这一项。

访存量变大了一个量级，计算量不变，计算强度变差了很多，现在需要打batch,大dim，小序列长度n才能实现较高的计算强度。

可以看到瓶颈在于内存复杂度，那么一个自然的想法是，可以减少内存消耗。MQA就是这样做的，而且很暴力，q仍然保持多头，kv只保留一个头，用这一个kv头广播到所有q头计算注意力。那么计算复杂仍然不变。

内存复杂度，需要保存的kv张量的d维度变小了，变成k（每个头的维度）。并且由于kv项变小了，这次不忽略激活值X了，内存占用为bnd。投影矩阵仍然不变，复杂度$O(n{d}^2)$

整体计算强度如下图最后一行，在原来的$\frac{n}{d}$项的分母上增加了一个头数h。计算强度更大了。并且增加了一个$\frac{1}{d}$

MQA虽然大幅减少内存，但对模型能力的损失也是很大的。于是折中的做法是GQA，几个q head一组，共享一个kv head，这样减少了内存，也保留了一部分表达能力。

具体效果如下图，第一张图横轴是推理耗时，纵轴是模型表现。可以看到GQA可以同时达到小参数MHA的推理速度，以及大参数MHA的性能。而MQA虽然推理最快，但没有比GQA快太多，且表现明显不如大参数MHA和GQA。

第二张图展示了GQA分组数对推理速度的影响，可以看到分组数只要稍微减少，从64减到32，推理速度就能大幅提升，减到16已经接近于MQA的速度。

所以GQA是个很好的权衡设计，既有接近于MQA的推理速度，又保留了接近MHA的模型能力。

稀疏注意力、滑动窗口注意力

对注意力的另一种优化方式是，每个q都要和全部kv计算，这太慢了，我们可以选择性让q只和一部分kv计算，这和前面的GQA不同，GQA优化的是内存，但计算量是不变的。稀疏注意力直接优化计算复杂度，减少了计算量。

如下图，每行对应一个q，每列表示一个token能看到的kv，蓝色表示能看到，灰色表示看不到，深蓝表示当前预测出的token。第一个是原始注意力，考虑因果掩码，每个token能看到他前面的全部token。

第二张图每个token只能看到前面三个token，这类似于一个滑动窗口。还有一个很浅的蓝色块，也是能看到的，窗口左侧每四个位置有一个。相当于每四个位置做一个总结，只看这个总结。

第三张图类似，也是滑动窗口+每四个位置一个总结，但总结token的位置是固定的，窗口的位置也是固定的。

这是一次注意力计算内部的，稀疏注意力还可以在注意力模块间稀疏。如下图，每四个注意力block，前三个是稀疏注意力，比如SWA(sliding window attn滑动窗口注意力)，第四个block是全局注意力。并且SWA有RoPE编码，全局注意力没有位置编码。这个位置编码的设计是，让中间的有位置编码的局部注意力层，根据位置编码处理局部信息，让没位置编码的全局注意力层不受距离的影响，只要有重要信息都记住，充当局部注意力之间的全局桥梁。

预告

注意力变体是个很广阔的领域，这次介绍的都是原始注意力的魔改，还有很多和原始注意力结构完全不同，但定位是注意力层的其他注意力模块，下一讲会详细讲。另外下一讲也会讲线性层的一个重大优化，混合专家架构（MoE）