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类），就是传统认知中标准的注意力

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

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

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

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

线性注意力

原始出处：Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

注意力计算可抽象成如下形式，sim表示可以使用任何的相似度函数，不一定非要softmax，可以是类似于多项式注意力或RBF核注意力等

$V_i^{^{\prime}}=\frac{\sum_{j=1}^{N}\operatorname{sim}\left(Q_i,K_j\right)V_j}{\sum_{j=1}^{N}\operatorname{sim}\left(Q_i,K_j\right)}.$

但相似度函数也是有要求的，需要施加唯一约束sim(·)非负。基于这个条件，给定这样一个内核 $\phi(x)$ ，可以将上式重写为

$V_i = \frac{\sum_{j=1}^N \phi(Q_i)^T \phi(K_j) V_j}{\sum_{j=1}^N \phi(Q_i)^T \phi(K_j)}$

利用矩阵乘法结合律得到

$V_i = \frac{ \phi(Q_i)^T \sum_{j=1}^N\phi(K_j) V_j^T}{ \phi(Q_i)^T \sum_{j=1}^N\phi(K_j)}$

为简化理解可写成如下形式

$\operatorname{Attention}(Q,K,V)=(\phi(Q)\phi(K)^T)V=\phi(Q)(\phi(K)^TV)$

深层理解：每个时间步的K和V可以提前计算并作为单个向量存储下来，在推理生成时直接用Q乘以每个时间步的KV，简单情况下KV cache的缓存向量都变少了，可以像RNN一样每次预测的时间都几乎是恒定的

注意：此时使用的一般都是绝对位置编码，Q、K矩阵没有乘以参数矩阵

此时 $\phi(K)^TV$ 的复杂度是 $O(d^2)$ ，所以整体复杂度变成了 $O(nd^2)$ ，随着序列长度n线性增长，此时就是线性注意力了。

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

$O = Δ^{-1} * (Q * K^T * V)$

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

因果模型存在的问题

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

$O=(QK^T\odot M)V$

$o_t=\sum_{j=1}^t(q_t^Tk_j)v_j$

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

累加求和操作的限制

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

$(\phi(Q)\phi(K)^T)V=\phi(Q)(\phi(K)^TV)$

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

在这里插入图片描述

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

$K^{T}= \begin{bmatrix} k_{1}^T & k_{2}^T & k_{3}^T & k_{4}^T \\ \end{bmatrix}= \begin{bmatrix} k_{11} & k_{21} & k_{31} & k_{41} \\ k_{12} & k_{22} & k_{32} & k_{42} \\ \vdots & \vdots & \vdots & \vdots \\ k_{1d} & k_{2d} & k_{3d} & k_{4d}\\ \end{bmatrix}$

$\begin{bmatrix} v_{1} \\ v_{2} \\ v_{3} \\ v_{4} \\ \end{bmatrix}= \begin{bmatrix} 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{bmatrix}$

$K^{T}V= \begin{bmatrix} k_{1}^Tv_1 + k_{2}^Tv_2 + k_{3}^Tv_3 + k_{4}^Tv_4 \\ \end{bmatrix}= \begin{bmatrix} [K^{T}V]_{1} & [K^{T}V]_{2} & ... & [K^{T}V]_{d} \end{bmatrix}$

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

$q_3K^{T}V= q_3 \begin{bmatrix} [K^{T}V]_{1} & [K^{T}V]_{2} & ... & [K^{T}V]_{d} \end{bmatrix}= \begin{bmatrix} q_3[K^{T}V]_{1} & q_3[K^{T}V]_{2} & ... & q_3[K^{T}V]_{d} \end{bmatrix}$

也可以使用以下代码测试

import torch
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(q3, kT @ v)
print('result', result)

在这里插入图片描述

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

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

$o_1 = q_1(k_1^Tv_1)$

$o_2 = q_2(k_1^Tv_1+k_2^Tv_2)$

$o_3 = q_3(k_1^Tv_1+k_2^Tv_2+k_3^Tv_3)$

$o_4 = q_4(k_1^Tv_1+k_2^Tv_2+k_3^Tv_3+k_4^Tv_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激活函数代替。即先变成了
$\operatorname{Attention}(Q,K,V)=(\phi(Q)\phi(K)^T\odot M)V$
然后还是先左乘计算，并没有解决线性注意力的根本问题，但是借鉴了flash attention中的硬件加速。

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

Lightning Attention-2

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

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

$o_2 = q_2(k_1^Tv_1+k_2^Tv_2)$

$o_3 = q_3(k_1^Tv_1+k_2^Tv_2+k_3^Tv_3)$

$o_4 = q_4(k_1^Tv_1+k_2^Tv_2+k_3^Tv_3+k_4^Tv_4)$

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

公式中也可以发现，当前step之前的k和v是可以相乘的，比如 $q_3$ 在计算时，可以将 $k_1^Tv_1+k_2^Tv_2+k_3^Tv_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^tV$ ，因为之前的kv是可以双向注意力的。然后将之前的kv结果缓存下来并更新，用于下一个step计算。

下图是Lightning Attention-2的结构图， $\lambda$ 是它的模型所使用的位置编码，忽略即可。
在这里插入图片描述
以下是前向传播和反向传播流程。

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