Diary

转码Day 22：感冒身体不适，今日看文档

lc：31

感冒了，一天擤鼻涕用了大半包抽纸，很难受，老是想打哈欠打不出来，而且头晕口干舌燥的，还有好几个口腔溃疡没好，也不知道过年前能不能好。

今日学习内容

去挖了一下苏剑林大佬的blog，学到了很多东西。

在RoPE中频率的计算公式为θi=b^(−2i/d)，底数b默认值为10000. 目前Long Context的主流做法之一是，先在b=10000上用短文本预训练，然后调大b并在长文本微调， “该过程给人的感觉是：调大b完全是因为“先短后长”的训练策略，如果一直都用长文本训练似乎就没必要调大b了？”

我之前的疑问就是，为什么不直接用大b来支持长文本训练？而是要先小后大呢？

论文《Base of RoPE Bounds Context Length》

“我们期望RoPE能具备两个理想性质，以达到更好的效果：1、远程衰减，即位置相近的Token平均来说获得更多的注意力；2、语义聚合，即语义相似的Token平均来说获得更多的注意力。” $\begin{aligned}&\mathbb{E}_{\boldsymbol{q},\boldsymbol{k},\boldsymbol{\varepsilon}}\left[\boldsymbol{q}^\top\boldsymbol{R}_{n-m}(\boldsymbol{q}+\boldsymbol{\varepsilon})-\boldsymbol{q}^\top\boldsymbol{R}_{n-m}\boldsymbol{k}\right]\\&\mathrm{=}=\mathbb{E}_q[q^\top\mathcal{R}_{n-m}q]-\mathbb{E}_{q,k}[q^\top\mathcal{R}_{n-m}k]\\&=\mathbb{E}_q[\boldsymbol{q}^\top\boldsymbol{R}_{n-m}\boldsymbol{q}]-\mathbb{E}_q[\boldsymbol{q}]^\top\boldsymbol{R}_{n-m}\mathbb{E}_k[\boldsymbol{k}\\&=\mathbb{E}_q\left[\boldsymbol{q}^\top\mathcal{R}_{n-m}\boldsymbol{q}\right]-\mu^2\mathbf{1}^\top\mathcal{R}_{n-m}\mathbf{1}\\&=\mathbb{E}_q\left[\sum_{i=0}^{d/2-1}(q_{2i}^2+q_{2i+1}^2)\cos(n-m)\theta_i\right]-\sum_{i=0}^{d/2-1}2\mu^2\cos(n-m)\theta_i\\&=\sum_{i=0}^{d/2-1}2(\mu^2+\sigma^2)\cos(n-m)\theta_i-\sum_{i=0}^{d/2-1}2\mu^2\cos(n-m)\theta_i\\&=\sum_{i=0}^{d/2-1}2\sigma^2\cos(n-m)\theta_i\end{aligned}$

作者从理论上证明了一个新的性质：

随着 token 相对距离增加，模型区分“相关 token”和“随机 token”的能力会逐渐衰减。

形式化地说： $\sum_{i=0}^{d/2-1}\cos m\theta_i\geq0,\quad m\in\{0,1,2,\cdots,L-1\}$

当 Bm,θ<0 时，模型甚至会：

  • 对随机 token 的注意力 ≥ 对相关 token 的注意力

  • 导致长距离信息无法被正确检索

RoPE 的 base 决定“有效上下文长度”

论文最重要的结论：

对于期望的上下文长度 L，RoPE 的 base 存在一个绝对下界。

$b^*=\inf\left\{\begin{array}{cc}b&\end{array}\right|f_b(m)\triangleq\sum_{i=0}^{d/2-1}\cos mb^{-2i/d}\geq0,m\in\{0,1,2,\cdots,L-1\}$

小 base 会产生“伪长上下文能力” 模型只是倾向于关注近距离 token，等价于“有效感受野缩小”。

核心结论：上下文长度增加 → 所需 base 呈指数增长 论文指出：

  • 避免 OOD 旋转角度只是必要条件

  • 但不是充分条件

理论上只要枚举的b足够多，那么对于任意L都可以找出最小的b。然而这里有个精度问题，原论文最大的L计算到了106，b至少要枚举到108，如果枚举间隔小，那么计算成本非常大，如果枚举间隔大，那么可能漏掉很多解。

LLAMA3，训练长度为8192，但RoPE的底数选择了500000（5e5）

缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA 梳理了从MHA、MQA、GQA到MLA的演变历程 MQA、GQA、MLA，都是围绕“如何减少KV Cache同时尽可能地保证效果”这个主题发展而来的产物 在GPU上部署模型的原则是：能一张卡部署的，就不要跨多张卡；能一台机部署的，就不要跨多台机。 对于MQA：这是最极限的情况，所有Q都共享了相同k v导致Attention的参数量减少了将近一半，而为了模型总参数量的不变，通常会相应地增大FFN/GLU的规模，这能弥补一部分效果损失。 目前看来大部分任务的损失都比较有限，且MQA的支持者相信这部分损失可以通过进一步训练来弥补回。

GQA:有人担心MQA对KV Cache的压缩太严重，以至于会影响模型的学习效率以及最终效果。为此，一个MHA与MQA之间的过渡版本GQA应运而生

MHA、MQA、GQA的铺垫后，迎来了MLA。 博客中说：MLA的本质改进不是低秩投影，而是低秩投影之后的工作。他认为GQA其实也是低秩投影。 $q_t^{(s)}k_i^{(s)\top}=\left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\right)\left(\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\right)\top=\boldsymbol{x}_t\left(\boldsymbol{W}_q^{(s)}\boldsymbol{W}_k^{(s)\top}\right)\boldsymbol{c}_i^\top$ ”意味着推理阶段，我们可以将$W_q^{(s)}W_k^{(s)\top}$合并起来作为Q的投影矩阵，那么$c_i$则取代了原本的$k_\mathrm{i}$ ,同理，在$o_t$后面我们还有一个投影矩阵，于是$v_i^{(s)}=c_iW_v^{(s)}$的$W_v^{(s)}$也可以吸收到后面的投影矩阵中去，于是等效地$v_\mathrm{i}$也可以用$c_\mathrm{i}$代替，也就是说此时KV Cache只需要存下所有的$c_\mathrm{i}$就行，而不至于存下所有的$k_i^{(s)}$、$v_i^{(s)}$。注意到$c_\mathrm{i}$跟$^{(s)}$无关，也就是说是所有头共享的，即MLA在推理阶段它可以恒等变换为一个MQA。“

看起来似乎非常好，但是上面的MLA有一个难以绕开的缺陷——不兼容RoPE（旋转位置编码） $q_i^{(s)}=\boldsymbol{x}_i\boldsymbol{W}_q^{(s)}\boldsymbol{R}_i\quad,\quad\boldsymbol{k}_i^{(s)}=\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\boldsymbol{R}_i\boldsymbol{q}_t^{(s)}\boldsymbol{k}_i^{(s)\top}=\left(\boldsymbol{x}_t\boldsymbol{W}_q^{(s)}\mathcal{R}_t\right)\left(\boldsymbol{c}_i\boldsymbol{W}_k^{(s)}\mathcal{R}_i\right)^\top=\boldsymbol{x}_t\left(\boldsymbol{W}_q^{(s)}\mathcal{R}_{t-i}\boldsymbol{W}_k^{(s)\top}\right)\boldsymbol{c}_i^\top$ $W_q^{(s)}R_{t-i}W_k^{(s)\top}$ 无法合并为一个固定的投影矩阵了（跟位置差t−i相关），从而MLA的想法若是结合RoPE实现的话，就无法利用恒等变换减少KV Cache，背离了设计缘由。 最后发布的MLA，采取了一种混合的方法——每个Attention Head的Q、K新增dr个维度用来添加RoPE

RoPE 影响的是内积：

$$Q_i^TR(i-j)K_j$$

当你把它限制在一个子空间：

$$Q_\text{ml}^TK_\text{ml}+Q_\text{rope}^TR(i-j)K_\text{rope}$$

这相当于： 内容相似度 + 位置相似度 这在表达上完全充分。

$q_{i}^{(s)}=\left[x_{i}W_{qc}^{(s)},x_{i}W_{qr}^{(s)}R_{i}\right]\in\mathbb{R}^{d_{k}+d_{r}}k_{i}^{(s)}=\left[c_{i}W_{kc}^{(s)},x_{i}W_{kr}^{\cancel{(s)}}R_{i}\right]\in\mathbb{R}^{d_{k}+d_{r}}$

K新增的维度每个Head共享，只需要Q的rope投影矩阵有head specific即可保证训练有效，且精度差不多，这样可以减少kv cache。 没有RoPE的维度就可以重复压缩的操作，在推理时KV Cache只需要存ci，新增的带RoPE的维度就可以用来补充位置信息，并且由于所有Head共享，所以也就只有在K Cache这里增加了dr个维度，原论文取了dr=dk/2=64，相比原本的dc=512，增加的幅度不大。Q 端新增 dr​ 维只带来每步计算/带宽开销，不增加推理 KV cache（因为不缓存 Q）。

最后有一个细节，就是MLA的最终版本，昨天看V3报告也看到了。 将Q的输入也改为了低秩投影形式，主要是为了减少训练期间参数量和相应的梯度所占显存 实际上MLA在解码阶段做的这个转换，虽然能有效减少KV Cache，但其解码的计算量是增加的。 还能提高推理效率是因为瓶颈不在解码，是带宽瓶颈和显存瓶颈 LLM架构参数一般满足h×dk=d，DeepSeek-V2不一样，它dk=128,d=5120，但h=128，是一般设置的3倍！这是因为MLA的KV Cache大小跟h无关，增大h只会增加计算量和提升模型能力，但不会增加KV Cache，所以不会带来速度瓶颈。

Transformer升级之路：20、MLA好在哪里?（上） 该实验初步结论 1、增大head_dims收益最大； 2、Partial RoPE对Loss也有一定帮助；[21年提出的只对部分维度加RoPE可能比全维度RoPE好,但现在仍然不是主流做法](Reason for doing partial rotary embedding? · Issue #40 · lucidrains/x-transformers) 3、KV-Shared应该也有一定作用。

Transformer升级之路：20、MLA好在哪里?（下）

给出断言：”在相同训练成本和推理成本下，MLA可能是效果最好的Full Attention变体。“ ”MLA之所以能够表现出色，有一个非常大的前提，那就是部分旋转的Partial RoPE效果不逊色于甚至可能优于完全体的RoPE。“ 个人理解所谓加dr维度用于rope编码也是partial rope的一种。 总的来说，在Partial RoPE的背景下，MLA几乎可以说是一个非常难以超越的Attention变体。

参考文献

  1. Jianlin Su. 缓存与效果的极限拉扯：从MHA、MQA、GQA到MLA. 科学空间, 2024.

  2. Jianlin Su. Transformer升级之路：20、MLA好在哪里?（上）. 科学空间, 2024.

  3. Jianlin Su. Transformer升级之路：20、MLA好在哪里?（下）. 科学空间, 2024.