flash-linear-attention中的Chunkwise并行算法的理解

这里提一下，我维护的几三个记录个人学习笔记以及社区中其它大佬们的优秀博客链接的仓库都获得了不少star，感谢读者们的认可，我也会继续在开源社区多做贡献。github主页：https://github.com/BBuf ，欢迎来踩

0x0. 前言

我之前解读过causal linear attention的cuda实现，文章见：https://zhuanlan.zhihu.com/p/673896906 ，也是在评论区通过@sonta 了解到了flash-linear-attention的Chunkwise并行实现。上篇文章https://mp.weixin.qq.com/s/H6wWBxwIJNCzkIlH_uIuiw中说到后续想继续解析一下chunk_rwkv6的实现，chunk_rwkv6的实现思路仍然是沿用flash-linear-attention中的Chunkwise并行思路，由于之前没有认真看过这个Chunkwise的算法所以读起来有点困难，这里需要用普通并行以及RNN递归的视角去看待才能理解这个算法流程。这篇文章就从 Gated Linear Attention Transformers with Hardware-Efficient Training (https://arxiv.org/pdf/2312.06635) 这篇Paper对线性Attention的Chunwise并行讲解和伪代码入手深入理解下这个方法，另外我们也会在后面深入分析下代码的实现。这篇Paper的作者也是flash-linear-attention的作者。

0x1. Paper部分

Paper部分这里只关注Background里面和Linear Attention相关的两节。这里对其进行翻译和解读。

我们首先简要介绍一下线性注意力层的背景。对于符号表示，我们使用黑体大写字母表示矩阵（例如，S、Q），黑体小写字母表示向量（例如，$q_t$、$k_t$），斜体大写字母表示可学习的参数矩阵（例如，$W_K$）。通常我们使用相同的字母表示矩阵的行，例如，$q_t$ 表示矩阵 $Q$ 的第 $t$ 行。

2.1 并行和递归形式

标准的Transformers采用softmax注意力机制，该机制接受输入序列 $X\in {\mathbb{R}}^{L\times d}$（其中 $L$ 是长度， $d$ 是隐藏维度）并通过以下方式计算输出 $O\in {\mathbb{R}}^{L\times d}$：

$$Q,K,V=X{W}_Q,X{W}_K,X{W}_V,O=\text{softmax}\left((Q{K}^T)\odot M\right)V,$$

其中 $W_Q,W_K,W_V\in {\mathbb{R}}^{d\times d}$ 是可学习的矩阵， M ∈ { − ∞ , 1 } L × L M \in \{-\infty, 1\}^{L \times L} M∈{ −∞,1}L×L 是一个掩码，用于防止模型关注未来的token，即 $M_{ij}=1$ 当 $i\ge j$ 且$M_{ij}=-\infty$ 当 $i<j$。 （这里我们假设一个简单的单头注意力。）上述的并行注意力形式可以在给定完整输入 $X$ 的情况下并行计算 $O$，从而实现高效训练。然而，在推理过程中，Transformer必须使用以下递归形式：

$q_t,k_t,v_t=x_t{W}_Q,x_t{W}_K,x_t{W}_V$

$o_t=\frac{{\sum }_{i=1}^t\exp (q_t{k}_i^T)}{{\sum }_{i=1}^t\exp (q_t{k}_i^T)}$