MiniMax-01中Lightning Attention的由来（线性注意力进化史）

引言

MiniMax-01: Scaling Foundation Models with Lightning Attention表明自己是第一个将线性注意力应用到如此大规模的模型，他所使用的核心技术就是Lightning Attention。

那为什么线性注意力20年在文章Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention中就提出了，现在才出第一个线性注意力的大模型呢？

本文就从线性注意力机制入手，详细探讨其起源、存在的显著局限性，以及Lightning Attention的具体实现细节。

原始注意力

现在主流的有两类模型，一种是应用双向注意力的bert类模型，另一种是应用单向注意力的gpt类模型，他们所使用的注意力其实是有细微差别的。

• 双向注意力（bert类），就是传统认知中标准的注意力

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{Q{K}^T}{\sqrt{d_{}}})V$$

• 单向注意力（因果模型，gpt类），只能看到当前和前面的token，所有要在softmax之前乘上一个掩码矩阵，M为单向掩码矩阵

$$\mathrm{Attention}(Q,K,V)=\mathrm{softmax}(\frac{Q{K}^T\odot M}{\sqrt{d_{}}})V$$

其中Q、K、V每个矩阵的维度都是[n, d]，即[序列长度，隐层维度]，此时$Q{K}^T$的维度是[n, n]，所以整体复杂度是$O(n^{2}d)$。其中d是固定大小，$n^2$随着序列长度平方增加，就主导了整体的复杂度。

线性注意力

原始注意力中softmax的作用主要是引入非线性（取概率化再与V乘都是次要的），那就可以将其换成其他的非线性激活函数。
$$\mathrm{Attention}(Q,K,V)=(\varphi (Q)\varphi (K{)}^T)V$$
这里的$\varphi$代表所使用的激活函数，有很多种可以选择（论文常用的有1+elu）。这里的归一化就先省略掉了，有一些论文就将K矩阵的归一化放到分母上（或者说K矩阵归一化的逆）。

此时观察，使用softmax必须等$Q{K}^T$先计算完，而使用其他的激活函数只对单个Q或者K进行运算，不需要绑定$Q{K}^T$。所以就可以将左乘变成右乘
$$(\varphi (Q)\varphi (K{)}^T)V=\varphi (Q)(\varphi (K{)}^{T}V)$$
此时$\varphi (K{)}^{T}V$的复杂度是$O(d^2)$，所以整体复杂度变成了$O(n{d}^2)$，随着序列长度n线性增长，此时就是线性注意力了。

（可选）：通常线性注意力的公式还有如下形式

$$O={\Delta }^{-1}\ast (Q\ast K^T\ast V)$$

（可选）其中，Δ起到了归一化的作用。Δ的每个对角元素是$K^T\ast 1$的值,这反映了每个键向量的重要程度。将${\Delta }^{-1}$乘到结果上,就相当于对注意力输出进行了逆归一化。相当于只对K归一化，Q本身就是一个合适的查询向量，不需要归一化。

因果模型存在的问题

注意上面的线性注意力是类bert模型的情况下，并没有与掩码矩阵相乘，此时可以顺畅的先右乘来降低复杂度。但现在的大模型都是生成模型，使用的因果模型结构，都是单向注意力，就必须要乘以掩码矩阵，所以不能顺畅的右乘了。
左乘线性注意力公式如下，输出为O，每个step的输出为当前的$q_t$乘以前面的$k_j$，再乘以$v_j$累加求和。此时$Q{K}^T$可以正常进行矩阵运算，然后使用$\odot$（Hadamard Product）进行逐元素相乘，得到掩码后的矩阵。

$$O=(Q{K}^T\odot M)V$$

$$o_t=\sum _{j=1}^t(q_t^T{k}_j)v_j$$

此时注意，上面公式的运算涉及$\odot$，它不适用于矩阵乘法交换律和结合律，即无法$Q(K^T\odot MV)$。$\odot$是逐元素相乘，所以两个矩阵的维度必须相同，即使将M的位置放到前面，$K^{T}V$的维度是[d, d]，也无法与M逐元素相乘。

累加求和操作的限制

双向注意力模型（bert）中使用的线性注意力如下，可以先算KV

$$(\varphi (Q)\varphi (K{)}^T)V=\varphi (Q)(\varphi (K{)}^{T}V)$$

QKV的维度都为[n, d]，这里假设序列长度为4，双向和单向注意力如下图

• 双向注意力计算
  K和V的矩阵如下，得到的$K^{T}V$的维度是[d, d]

$$K^T=\left[\begin{array}{cccc}{k}_1^T& k_2^T& k_3^T& k_4^T\end{array}\right]=\left[\begin{array}{cccc}{k}_{11}& k_{21}& k_{31}& k_{41}\\ k_{12}& k_{22}& k_{32}& k_{42}\\ ⋮& ⋮& ⋮& ⋮\\ k_{1d}& k_{2d}& k_{3d}& k_{4d}\end{array}\right]$$

$$V=\left[\begin{array}{c}{v}_1\\ v_2\\ v_3\\ v_4\end{array}\right]=\left[\begin{array}{cccc}{v}_{11}& v_{12}& ...& v_{1d}\\ v_{21}& v_{22}& ...& v_{2d}\\ v_{31}& v_{32}& ...& v_{3d}\\ v_{41}& v_{42}& ...& v_{4d}\end{array}\right]$$

$$K^{T}V=\left[\begin{array}{c}{k}_1^T{v}_1+k_2^T{v}_2+k_3^T{v}_3+k_4^T{v}_4\end{array}\right]=\left[\begin{array}{c}[K^{T}V{]}_1\\ [K^{T}V{]}_2\\ ⋮\\ [K^{T}V{]}_d\end{array}\right]$$

此时计算$q_3$的注意力输出就可以使用以下方法。注意这是点积，q3是一个向量，$K^{T}V$是一个矩阵，向量在与矩阵点积的时候会进行广播拓展，复制成多份分别与矩阵中的向量点积。$[K^{T}V{]}_1$是一个向量，$q_3[K^{T}V{]}_1$点积后会得到一个值，所以$q_3{K}^{T}V$最终的结果是一个向量，长度为隐层维度d。

$$q_3{K}^{T}V=q_3\left[\begin{array}{c}[K^{T}V{]}_1\\ [K^{T}V{]}_2\\ ⋮\\ [K^{T}V{]}_d\end{array}\right]=\left[\begin{array}{c}{q}_3[K^{T}V{]}_1\\ q_3[K^{T}V{]}_2\\ ⋮\\ q_3[K^{T}V{]}_d\end{array}\right]$$

也可以使用以下代码测试

q3 = torch.tensor([1, 2, 3, 4, 5, 6])
print(q3)# [n, d] = [4, 6]
kT = torch.tensor([[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3], [4, 4, 4, 4],[5, 5, 5, 5],[6, 6, 6, 6]])
v = torch.tensor([[1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1]])print('kT @ v', kT @ v)
# q与(k.T @ v)的点积
result = torch.matmul(q, kT @ v)
print('result', result)

此时$K^{T}V$的结果是双向的，$k_3$的输出矩阵中使用了$v_4$，这样双向注意力就可以顺畅的右乘得到$K^{T}V$结果再与Q相乘，得到所有token的输出。

但因果模型的注意力是单向的，$K^{T}V$在计算的时候前面的K不能与后面的V相乘，所以只能一个一个算然后累加求和。

$$o_1=q_1(k_1^T{v}_1)$$

$$o_2=q_2(k_1^T{v}_1+k_2^T{v}_2)$$

$$o_3=q_3(k_1^T{v}_1+k_2^T{v}_2+k_3^T{v}_3)$$

$$o_4=q_4(k_1^T{v}_1+k_2^T{v}_2+k_3^T{v}_3+k_4^T{v}_4)$$

这样的累加操作无法进行高效的矩阵乘法，虽然计算复杂度降低了，但实际运算的效率并不高。

Lightning Attention

到这里可以引出MiniMax-01 中所使用的Lightning Attention了，但其实这个注意力有两个版本，MiniMax-01中所提到的就是是Lightning Attention-2，那咱们先看看第一个版本做了什么。

Lightning Attention-1

源自：TransNormerLLM: A Faster and Better Large Language Model with Improved TransNormer

Lightning Attention-1针对于原始注意力取消了softmax，使用Swish激活函数代替。即先变成了
$$\mathrm{Attention}(Q,K,V)=(\varphi (Q)\varphi (K{)}^T\odot M)V$$
然后还是先左乘计算，并没有解决线性注意力的根本问题，但是借鉴了flash attention中的硬件加速。

其前向和反向传播流程如下，就是将QKV切块，放到高速SRAM中去计算。虽然变快了，但此时的复杂度还是$O(n^{2}d)$。

Lightning Attention-2

源自：Lightning Attention-2: A Free Lunch for Handling Unlimited Sequence Lengths in Large Language Models

Lightning Attention-2解决了因果模型在计算单向注意力时，需要进行累加求和操作导致无法矩阵运算的情况，实现了单向注意力先计算右乘，成功将复杂度降为$O(n{d}^2)$。
$$o_1=q_1(k_1^T{v}_1)$$

$$o_2=q_2(k_1^T{v}_1+k_2^T{v}_2)$$

$$o_3=q_3(k_1^T{v}_1+k_2^T{v}_2+k_3^T{v}_3)$$

$$o_4=q_4(k_1^T{v}_1+k_2^T{v}_2+k_3^T{v}_3+k_4^T{v}_4)$$

再将这个累加求和公式拿过来，配合下图观察发现，之前的问题是每次计算$Q{K}^T$都在整个序列上计算，这样每次都是所有序列的token互相注意到。那如果在序列这个维度拆分成小份，比如图中右侧先计算$k_1$和$k_2$，然后用于$q_3$的计算就完全没有问题，$k_4$后面的就不计算了。这样就既能矩阵运算，又能符合单向掩码。

公式中也可以发现，当前step之前的k和v是可以相乘的，比如$q_3$在计算时，可以将$k_1^T{v}_1+k_2^T{v}_2+k_3^T{v}_3$使用矩阵操作运算。所以Lightning Attention-2将大矩阵拆开，类似flash attention拆成多个block。

这些 block 不能拆分成 n 份，这样block的意义就没有了，for循环计算反而更慢。所以每个 block 中会有多个时间步的token。

此时这些 block 就可以分为两类，一类是块内（intra block），一类是块间（inter block）。块内代表当前块 q 的序列下标和 kv 序列下标相同，块间即不同。

块内在计算$q_i$时直接矩阵右乘很容易算上$k_{i+1}{v}_{i+1}$，所以块内使用传统的左乘并与掩码矩阵相乘。块间计算时就可以先右乘计算$K^{t}V$，因为之前的kv是可以双向注意力的。然后将之前的kv结果缓存下来并更新，用于下一个step计算。

下图是Lightning Attention-2的结构图，$\lambda$是它的模型所使用的位置编码，忽略即可。

以下是前向传播和反向传播流程。


问题：M矩阵维度是[B, B]，相当于每一个块代表了多个序列步n，在对角线位置是1，那在这个块内前面的q就可以注意到后面的kv了

解答：M矩阵维度虽然是[B, B]，但只是这么切割，其内部值仍然是下三角。

备注

个人理解，若有不对请指出，谢谢。