大语言模型之十一 Transformer后继者Retentive Networks (RetNet)

在《大语言模型之四-LlaMA-2从模型到应用》的LLama-2推理图中可以看到，在输入“你好！”时，是串行进行的，即先输入“你”这个token，然后是“好”，再然后是“！”token，前一个token需要保留前面的k和v矩阵，这就意味着随着输入sequence length的增长，需要的内存也会快速增长，计算量也会快速增长。这也显示了Transformer尽管在模型训练的时候并发（相比RNN）性能好，且模型的效果也好，但是推理的时候效率就比较低。

RetNet特点

微软提出的RetNet在训练并发、模型效果以及推理效率上都取得了不错的效果。下图是其paper中关于模型性能和推理效率和Transformer的对比情况。
在这里插入图片描述
其官方paper宣称，其实验数据显示，在语言建模任务上：
RetNet 可以达到与 Transformer 相当的困惑度（perplexity）
推理速度达8.4倍
推理算法延迟降低90%
内存占用减少70%
具有良好的扩展性
实现以上改进是因为RetNet 在 Transformer 的基础上，使用多尺度保持（Retention）机制替代了标准的自注意力机制。

与标准自注意力机制相比，保持机制有几大特点：

引入位置相关的指数衰减项取代 softmax，简化了计算，同时使前步的信息以衰减的形式保留下来。
引入复数空间表达位置信息，取代绝对或相对位置编码，容易转换为递归形式。
保持机制使用多尺度的衰减率，增加了模型的表达能力，并利用 GroupNorm 的缩放不变性来提高 Retention 层的数值精度。

训练并发

因为Transformer采用了self-attention机制，每一阶段的输出都可以用Q,K,V进行并发处理，这大大提高了GPU利用率，提高了训练效率。RNN网络的好处是推理的效率高（相比开篇提的KV历史是不需要保留的，这使得计算量和内存都极大减少了），内存复杂度低O(1)。
RetNet的巧妙之处在于，训练的时候依然类似Transformer的并发结构（如图3左边），而在推理的时候则可以采用RNN的结构（如图3右边)。即实现了parallel training, recurrent/chunk-wise inference。

RetNet改进点

RetNet相比于Transformer主要有两点改进：
1.引入multi-scale retention替代了multi-head attention；
2.RetNet可以用三种（parallel/recurrent/chunk-wise ）方式实现，公式和结果上是相等的，因而可以在训练和推理的时候选择最为高效的方式实现；chunk-wise可以更高效的处理长sequence的序列情况。
在这里插入图片描述

并行训练

在这里插入图片描述
RetNet摒弃了softmax操作，引入了基于D矩阵的Hadamard积（对应元素相乘），然后是GroupNorm操作。Transformer中softmax为输入序列中的每个token提供了相对的注意力权重，有助于模型学习和保留长期依赖关系。之前也有一些研究是舍弃softmax运算，但是会降低模型性能。
具体来说Transformer中的softmax主要实现了两个目标：
1.对不同的时间步长采用不同的方式加权，这有助于模型注意力放在序列中应该感兴趣的部分，这是相比RNN而言最重要的贡献。RetNet论文中的D矩阵实现注意力机制，D矩阵是一个因果注意力矩阵，即序列的当前输入只能看到过去的信息，D-矩阵假设最近的时间步长比过去的时间步长更重要，因此采用了指数衰减权重。因此，softmax足够灵活，可以对不同的步骤做不同的权重估计，而D-矩阵以固定的预定义方式（指数衰减）权衡所有步骤。最终paper里的结果显示RetNet效果是比Transformer好的。
2.引入非线性，当没有softmax的时候， $QK^T$ 就是一种仿射变换（即从一种空间的标识变为另一种空间的表示)，再多层的注意力堆叠也依然是一种仿射变换，非线性是通过GroupNorm的方式实现的，为什么是GroupNorm实现非线性，论文中并没有提及，似乎是盲测多种结构在GroupNorm时效果是最好的。

Transformer和RetNet的异同

这里我将博客《大语言模型之四-LlaMA-2从模型到应用》中LlaMA-2的推理流程图展示在这里了，该图的详细说明见博客。
在这里插入图片描述
图中第一步是用 $W_q, W_v, W_k$ 得到 $Q, V, 和 K$ ，即 $Q=XW_q$ ， $K=XW_k$ ， $V=XW_v$ ，然后通过softmax得到Attention score
$(\frac{Q K^T}{\sqrt{d_k}})V (1)$
由于RetNet在循环网络结构和并行网络结构都可以得到一样的结果，所以论文的作者现在循环结构中实现激励的“保留”模块，然后将“保留”模块向量化，因此得到的第n个序列输入的结果(保留注意力)为：
$o_n=\sum_{m=1}^n(Q_nA^{n-m}K_m^T)v_m, Q_n \in \mathbb{R}^{1\times d} (2)$
从上面公式1和2可以看到RetNet网络使用pos矩阵 $A^{n-m}$ 替换掉了Transformer中的softmax的Attention部分（非线性部分这里还没替换掉)。Retention score的想法是这样的，首先通过循环网络状态 $\mathbf s(n)$ 做映射 $\rightarrow o(n)$ ，即
$\mathbf s_n = A \mathbf s_{n-1} +K_n^Tv_n, A \in \mathbb R^{d\times d}, K_n \in \mathbb R^{1 \times d} (3)$
然后使用线性变换递归对序列进行编码：
$o_n=Q_n \mathbf s_n=\sum_{m=1}^n(Q_nA^{n-m}K_m^T)v_m, Q_n \in \mathbb{R}^{1\times d}(4)$
根据原论文的公式3，可以将位置分为共轭的两个部分。
$\sum_{m=1}^n(Q_n(\gamma e^{i \theta})^{n})(K_m(\gamma e^{i\theta})^{-m})^Tv_m, \gamma, \theta \in \mathbb R^d (5)$
其中 $Q_n(\gamma e^{i \theta})^{n}$ 和 $K_m(\gamma e^{i\theta})^{-m}$ 是位置矩阵，可以采用论文所述的xPos，为了简化上述的方程，可以用标量值替代 $\gamma$ ，这样在训练的时候可以采用并发的方式训练：
在这里插入图片描述
这样可以看到在RetNet时候，在得到Q、K以及V的方法和原始的Transformer一样是可以并发进行的， $e^{in\theta}$ 对应的位置信息逐点相乘即可。最后一步的Retention score使用到的D矩阵也是可以提前计算的，因为它只是一个相对位置嵌入+因果掩码。

在这里插入图片描述
图 3 RetNet的两种实现方式
公式6对应的就是图3左边的实现方式，在推理的时候采用循环网络的架构，即图3中右边的实现方式，公式如下7所示。
$S_n=\gamma S_{n-1}+K_n^TV_n , Retention(X_n)=Q_nS_n, n=1,\cdots, |x| （7）$
其中的Q、K、V以及 $\gamma$ 的作用和意义和上式6是一样的。

RetNet并行计算过程

假设仅有“你好”这两个token输入序列，长度记为N，Embedding dim，D=3，则得到的QKV是矩阵的维度是NxD，假设初始的矩阵为：
$Q=\begin{bmatrix} 1 & 2 & 1 \\ 3 & 2 &3 \end{bmatrix},K=\begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 &6 \end{bmatrix},V=\begin{bmatrix} 5 & 4 & 3 \\ 2 & 1 &0 \end{bmatrix}$
第一步：计算 $QK^T$
$QK^T=\begin{bmatrix} 1 & 2 & 1 \\ 3 & 2 &3 \end{bmatrix} \begin{bmatrix} 1 & 4 \\ 2 & 5 \\ 3 & 6 \end{bmatrix}=\begin{bmatrix} 8 & 20 \\ 16 & 40 \end{bmatrix} （8）$
第二步：计算 $QK^T$ 和 $D$ 矩阵的Hadamard积（对应元素相乘)
$QK^T\odot D=\begin{bmatrix} 8 & 20 \\ 16 & 40 \end{bmatrix} \odot \begin{bmatrix} 1 & 0 \\ 0.25 & 1 \end{bmatrix}= \begin{bmatrix} 8 & 4 \\ 0& 40 \end{bmatrix}（9）$
其中 $D=\begin{bmatrix} \gamma^{1-0} & \gamma^{1-1} \\ \gamma^{2-0} & \gamma^{2-1} \end{bmatrix}$ ，当 $\gamma=0.5$ 时可以得到式9。
第三步和V相乘：
$(QK^T \odot D)V = \begin{bmatrix} 8 & 4 \\ 0& 40 \end{bmatrix} \begin{bmatrix} 5 & 4 & 3\\ 2& 1 & 0 \end{bmatrix} = \begin{bmatrix} 40 & 32 & 24\\ 100& 56 & 12 \end{bmatrix} (10)$
这样就得到了两个token输入时的最终上下文Embedding结果。

RetNet循环网络计算过程

图3的右侧所示过程。这里的Q、K、V和上面的并行计算并不是一样的，这里有下标n，这表明了是第n个输入token对应的矩阵，因而是一个1xD维矩阵(上一小节是NxD维)，另外一个区别是含有当前token之前时间和位置信息的状态S，当前状态和前一个状态使用指数衰减因子，如公式7所示。
在计算的时候先是KV不再是QK相乘，图3中可以看到，和并行计算一样，初始的矩阵为
$Q=\begin{bmatrix} 1 & 2 & 1 \\ 3 & 2 &3 \end{bmatrix},K=\begin{bmatrix} 1 & 2 & 3 \\ 4 & 5 &6 \end{bmatrix},V=\begin{bmatrix} 5 & 4 & 3 \\ 2 & 1 &0 \end{bmatrix}$
第一步计算 $K_1^T\otimes V_1$
$K_1^T V_1 = \begin{bmatrix} 1 \\ 2 \\ 3 \end{bmatrix} \begin{bmatrix} 5 & 4 & 3\end{bmatrix} = \begin{bmatrix} 5 & 4 & 3\\ 10& 8 & 6 \\ 15 & 12 &9 \end{bmatrix} (11)$
第二步计算 $S_1$
因为在输入“你”token之前并没有其他token输入，因而 $S_0$ 是不存在的，因而先前没有状态叠加到“你”这个token对应的状态的。
$S_1=\gamma^{0}S_0 + K_1^TV_1=\begin{bmatrix} 5 & 4 & 3\\ 10& 8 & 6 \\ 15 & 12 &9 \end{bmatrix} (12)$
第三步将Q和 $S_1$ 相乘得到最终的Attention score。
$Q_1 \otimes S_1 = \begin{bmatrix} 5 & 4 & 3\\ 10& 8 & 6 \\ 15 & 12 &9 \end{bmatrix} \otimes \begin{bmatrix} 1 \\ 2 \\ 1 \end{bmatrix} = \begin{bmatrix} 5 & 4 & 3\\ 20& 16 & 12 \\ 15 & 12 &9 \end{bmatrix}_{sum} = \begin{bmatrix} 40 & 32 & 24\end{bmatrix} (13)$
这和公式10最终结果的第一行是一样的。
第四步计算 $K_2^TV_2$ ，和第一步类似，得到：
$K_2^T V_2 = \begin{bmatrix} 4 \\ 5 \\ 6 \end{bmatrix} \begin{bmatrix} 2 & 1 & 0\end{bmatrix} = \begin{bmatrix} 8 & 4 & 0\\ 10& 5 & 0 \\ 12 & 6 & 0 \end{bmatrix} (14)$
第五步计算 $S_2$ ，
$S_2=\gamma^{2}S_1 + K_2^TV_2=0.5^2\begin{bmatrix} 5 & 4 & 3\\ 20& 16 & 12 \\ 15 & 12 & 9 \end{bmatrix} + \begin{bmatrix} 8 & 4 & 0\\ 10& 5 & 0 \\ 12 & 6 & 0 \end{bmatrix} = \begin{bmatrix} 9.25 & 0.25 & 0.75\\ 12.5& 7 & 1.5 \\ 15.5 & 9 & 2.25 \end{bmatrix} (15)$
第六步计算最终RetNet score。
$Q_2 \otimes S_2 = \begin{bmatrix} 9.25 & 0.25 & 0.75\\ 12.5& 7 & 1.5 \\ 15.5 & 9 & 2.25 \end{bmatrix} \otimes \begin{bmatrix} 3 \\ 2 \\ 3 \end{bmatrix}(16) =\begin{bmatrix} 27.25 & 15 & 2.25 \\ 25& 14 & 3 \\ 47.25 & 27 & 6.75 \end{bmatrix}_{sum} = \begin{bmatrix} 100 & 56 & 12\end{bmatrix} (16)$
这样公式16的结果正好等于公式10，从公式13和16结果看和并行计算的结果是一样的，但是这中间并没有太多的存储空间需求，也没有非线性softmax计算。