2023Retnet论文解读

简介

微软亚洲研究院上线了大模型新架构的论文“Retentive Network: A Successor to Transformer for Large Language Models”，该基础架构采用了新的 Retention 机制来代替 Attention，向 Transformer 发起挑战！相对于 Transformer 架构的优势是同时具备:训练可并行、推理成本低和良好的性能，不可能三角。

论文中给出一个很形象的示意图，RetNet 在正中间表示同时具备三个优点，而其他的架构 Linear Transformer、Recurrent Network 和 Transformer 都只能同时具备其中两个优点。

实验结果

论文给出的 RetNet 和 Transformer 的对比实验结果:

当输入序列长度增加的时候，RetNet 的 GPU 显存占用一直是稳定的和权值差不多，而 Transformer 则是和输入长度成正比。

首先看红色线和紫色线，都是输入长度在 8192 下，RetNet 和 Transformer 推理延时的对比。
可以看到当 batch size 增加的时候， RetNet 的推理延时也还是很稳定，而 Transformer 的推理延时则是和 batch size 成正比。
而 Transformer 即使是输入长度缩小到 1024 ，推理延时也还是比 RetNet 要高。

LLM困惑度$ppl\in [1,\infty ]$上的实验也比之前的一些工作要好。

在67亿参数规模上的LLM实验效果也比Transformer好

训练效率的对比，看得出来去掉softmax之后有一定的提升。

RetNet 架构解读

与标准自注意力机制相比，保持机制有几大特点：
引入位置相关的指数衰减项取代 softmax，简化了计算，同时使前步的信息以衰减的形式保留下来。
引入复数空间表达位置信息，取代绝对或相对位置编码，容易转换为递归形式。
另外，保持机制使用多尺度的衰减率，增加了模型的表达能力，并利用 GroupNorm 的缩放不变性来提高 Retention 层的数值精度。
RetNet 架构和 Transformer 类似，也是堆叠 $L$层同样的模块，每个模块内部包含两个子模块：一个 multi-scale retention（MSR）和一个 feed-forward network (FFN)。
下面详细解读一下这个 retention 子模块。
首先给定一个输入序列 { x i } i = 1 ∣ x ∣ \left\{ x_i \right\}_{i=1}^{|x|} {xi​}i=1∣x∣​：
$$x=x_1\cdot \cdot \cdot x_{|x|}$$
其中 $|x|$表示序列的长度。然后输入序列首先经过 embedding 层得到词嵌入向量：
$$X^0=[x_1,\dots ,x_{|x|}]\in {\mathbb{R}}^{|x|\times d}$$
其中 $d$ 表示隐含层的维度。

Retention 机制

首先对给定输入词嵌入向量序列$X\in {\mathbb{R}}^{|x|\times d}$中的每个时间步 $n$ 的向量$X_n\in {\mathbb{R}}^{1\times d}$都乘以权值$w_v\in {\mathbb{R}}^{d\times d}$得到$v_n\in {\mathbb{R}}^{1\times d}$：
$$v_n=X_n\cdot w_V$$
然后同样有类似 Transformer 架构的 $Q$ 和 $K$ 的投影：
$$Q=X{W}_Q,K=X{W}_K$$
其中$W_Q\in {\mathbb{R}}^{d\times d}$,$W_K\in {\mathbb{R}}^{d\times d}$是需要学习的权值。
接着假设现在有一个序列建模的问题，通过状态$s_n\in {\mathbb{R}}^{d\times d}$将$v_n$映射为$o_n$向量。首先来看论文中给出的映射方式定义：
$$s_n=A{s}_{n-1}+K_n^T\cdot v_n$$
$$o_n=Q_n\cdot s_n=\sum _{m=1}^n{Q}_n{A}^{n-m}{K_m}^T{v}_m$$
其中$A\in {\mathbb{R}}^{d\times d}$是一个矩阵，$K_n\in {\mathbb{R}}^{1\times d}$表示时间步 $n$ 对应的 $K$ 投影，则$K^T{v}_n\in {\mathbb{R}}^{d\times d}$。同样$Q_n\in {\mathbb{R}}^{1\times d}$表示时间步 $n$ 对应的 $Q$ 投影。

那么上面公式中的 $o_n$计算公式是怎么得出来呢，下面详细解释一下，首先将$Q_n{s}_n$展开：
$$Q_n{s}_n=Q_n(A{s}_{n-1}+{K_n}^T{v}_n)$$
$$=Q_n(A(A{s}_{n-2}+{K_{n-1}}^T{v}_{n-1})+{K_n}^T{v}_n)$$
$$=Q_n(A^2{s}_{n-2}+A^1{K_{n-1}}^T{v}_{n-1}+A^0{K_n}^T{v}_n)$$
$$=Q_n(A^2(s_{n-3}+A^1{K_{n-2}}^T{v}_{n-2})+A^1{K_{n-1}}^T{v}_{n-1}+A^0{K_n}^T{v}_n)$$
$$=Q_n(A^3{s}_{n-3}+A^2{K_{n-2}}^T{v}_{n-2}+A^1{K_{n-1}}^T{v}_{n-1}+A^0{K_n}^T{v}_n)$$
其中 $A^0$表示单位矩阵(主对角线元素为1，其余元素为0的方阵)。然后我们假定 $s_0$ 为初始状态元素为全0的矩阵，则有：
$$s_1=A{s}_0+{K_1}^T{v}_1={K_1}^T{v}_1$$
再继续上述推导过程：
$$Q_n{s}_n=Q_n(A^3{s}_{n-3}+A^2{K_{n-2}}^T{v}_{n-2}+A^1{K_{n-1}}^T{v}_{n-1}+A^0{K_n}^T{v}_n)$$
$$=Q_n(A^{n-(n-3)}{s}_{n-3}+A^{n-(n-2)}{K_{n-2}}^T{v}_{n-2}+A^{n-(n-1)}{K_{n-1}}^T{v}_{n-1}+A^{n-(n-0)}{K_n}^T{v}_n)$$
所以根据上述推导过程和条件归纳可得：
$$Q_n{s}_n=Q_n(A^{n-1}{s}_1+A^{n-2}{K_2}^T{v}_2+...+A^{n-(n-2)}{K_{n-2}}^T{v}_{n-2}+A^{n-(n-1)}{K_{n-1}}^T{v}_{n-1}+A^{n-(n-0)}{K_n}^T{v}_n)$$
$$=Q_n(A^{n-1}{K_1}^T{v}_1+A^{n-2}{K_2}^T{v}_2+...+A^{n-(n-2)}{K_{n-2}}^T{v}_{n-2}+A^{n-(n-1)}{K_{n-1}}^T{v}_{n-1}+A^{n-(n-0)}{K_n}^T{v}_n)$$
$$=\sum _{m=1}^n{Q}_n{A}^{n-m}{K_m}^T{v}_m$$
然后我们来看一下 $A$ 矩阵是什么，论文中定义了 $A$ 是一个可对角化的矩阵，具体定义为：
$$A=\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}$$
其中$\gamma ,\theta \in {\mathbb{R}}^d$都是$d$维向量，$\Lambda$是可逆矩阵。
欧拉公式：
$$e^{ix}=cosx+i\cdot sinx$$
其中 $x$ 表示任意实数， $e$ 是自然对数的底数，$i$ 是复数中的虚数单位，也可以表示为实部$cosx$ ，虚部 $sinx$ 的一个复数，欧拉公式建立了指数函数、三角函数和复数之间的桥梁。
而这里 $\theta$ 是一个 $d$ 维向量:
$$\theta =[{\theta }_1,{\theta }_2,...,{\theta }_{d-1},{\theta }_d]$$
则 $e^{i\theta }$ 也就是将向量元素两两一组表示分别表示为复数的实部和虚部：
$$e^{i\theta }=[cos{\theta }_1,sin{\theta }_2,...,cos{\theta }_{d-1},sin{\theta }_d]$$
然后 $\gamma e^{i\theta }$ 就是一个对角矩阵，对角元素的值就对应将 $\gamma$ 和 $e^{i\theta }$ 转成复数向量相乘再将结果转回实数向量的结果。
现在我们知道了矩阵 $A$ 的构成就能得到：
$$A^{n-m}=(\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}{)}^{n-m}$$
这里因为 $\Lambda$ 是可逆矩阵则有性质
$$\Lambda {\Lambda }^{-1}={\Lambda }^{-1}\Lambda =I$$
其中 $I$为单位矩阵，则将 $n-m$ 次方展开：
$$A^{n-m}=\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}....\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}$$
就是 $n-m$ 个 $\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}$ 矩阵相乘，中间相邻的 ${\Lambda }^{-1}\Lambda$ 都消掉了，所以可得：
$$A^{n-m}=\Lambda (\gamma e^{i\theta }{)}^{n-m}{\Lambda }^{-1}$$
然后我们回到计算 $o_n$ 的公式：
$$o_n=\sum _{m=1}^n{Q}_n{A}^{n-m}{K_m}^T{v}_m$$
$$=\sum _{m=1}^n{Q}_n(\Lambda (\gamma e^{i\theta }{)}^{n-m}{\Lambda }^{-1}){K_m}^T{v}_m$$
$$=\sum _{m=1}^n{X}_n{W}_Q(\Lambda (\gamma e^{i\theta }{)}^{n-m}{\Lambda }^{-1})(X_m{W}_K{)}^T{v}_m$$
$$=\sum _{m=1}^n{X}_n{W}_Q(\Lambda (\gamma e^{i\theta }{)}^{n-m}{\Lambda }^{-1})X_m^T{W}_K^T{v}_m$$
接着论文中提出把 $\Lambda$ 吸收进 $W_Q$ 和 $W_K$ 也就是 $W_Q\Lambda$ 和 ${\Lambda }^{-1}{W}_K^T$ 分别用 $W_Q$ 和 $W_K^T$ 替代当作学习的权值，那么可得:
$$o_n=\sum _{m=1}^n{Q}_n(\gamma e^{i\theta }{)}^{n-m}{K}_m^T{v}_m$$
$$=\sum _{m=1}^n{Q}_n(\gamma e^{i\theta }{)}^n(\gamma e^{i\theta }{)}^{-m}{K}_m^T{v}_m$$
$$=\sum _{m=1}^n{Q}_n(\gamma e^{i\theta }{)}^n(K_m\gamma e^{i\theta }{)}^{-m}{)}^T{v}_m$$
$$=\sum _{m=1}^n{Q}_n({\gamma }^n{e}^{i\theta n})(K_m{\gamma }^{-m}{e}^{i\theta (-m)}{)}^T{v}_m$$
接着将公式简化，将 $\gamma$ 改为一个实数常量，那么可得：
$$o_n=\sum _{m=1}^n{Q}_n({\gamma }^n{e}^{i\theta n})(K_m{\gamma }^{-m}{e}^{i\theta (-m)}{)}^T{v}_m$$
$$=\sum _{m=1}^n{\gamma }^{n-m}(Q_n{e}^{i\theta n})(K_m{e}^{i\theta (-m)}{)}^T{v}_m$$
在继续推导前，先来仔细看一下 $e^{i\theta (-m)}$，借助欧拉公式展开：
$$e^{i\theta (-m)}=[cos-m{\theta }_1,sin-m{\theta }_2,...,cos-m{\theta }_{d-1},sin-m{\theta }_d]$$
三角函数性质：
$$cos(-\theta )=cos\theta$$
$$sin(-\theta )=-sin\theta$$
则有：
$$e^{i\theta (-m)}=[cosm{\theta }_1,-sinm{\theta }_2,...,cosm{\theta }_{d-1},-sinm{\theta }_d]$$
转为复数形式表示就是：
$$e^{i\theta (-m)}=[cosm{\theta }_1,-isinm{\theta }_2,...,cosm{\theta }_{d-1},-isinm{\theta }_d]$$
刚好就对应 $e^{i\theta m}$ 的共轭
$$e^{i\theta m}=[cosm{\theta }_1,+isinm{\theta }_2,...,cosm{\theta }_{d-1},+isinm{\theta }_d]$$
所以可得：
$$o_n=\sum _{m=1}^n{\gamma }^{n-m}(Q_n{e}^{i\theta n})(K_m{e}^{i\theta m}{)}^T{v}_m$$
$$=\sum _{m=1}^n{\gamma }^{n-m}(Q_n{e}^{i\theta n})(K_m{e}^{i\theta m}{)}^{†}{v}_m$$
其中 $†$ 表示共轭转置操作。

Retention 的训练并行表示

首先回顾单个时间步 $n$ 的输出 $o_n$ 的计算公式如下：
$$o_n=\sum _{m=1}^n{\gamma }^{n-m}(Q_n{e}^{i\theta n})(K_m{e}^{i\theta m}{)}^{†}{v}_m$$
而所有时间步的输出是可以并行计算的，用矩阵形式表达如下：
$$((Q\odot \Theta )(K\odot \bar{\Theta }{)}^T\odot D)V$$
其中 $V\in {\mathbb{R}}^{|x|\times d}$，而 $\odot$ 表示两个矩阵逐元素相乘， $Q\in {\mathbb{R}}^{|x|\times d}$和 $K\in {\mathbb{R}}^{|x|\times d}$ 每一行对应一个时间步的 $q$ 和 $k$ 向量。
而 $\Theta \in {\mathbb{R}}^{|x|\times d}$ 每一行对应向量$e^{i\theta n},n=1,...,|x|$ 。 $\bar{\Theta }\in {\mathbb{R}}^{|x|\times d}$就是对应 $\Theta$ 矩阵的共轭，也就是将 $\Theta$ 矩阵每一行改为复数的共轭形式。
而 $D\in {\mathbb{R}}^{|x|\times |x|}$ 矩阵是一个下三角矩阵，其中第 $n$ 行第 $m$ 列的元素计算方式：
$$D_{nm}={\gamma }^{n-m},n>=m$$
$$D_{nm}=0,n<m$$
标准Transformers是： $softmax(Q{K}^T/\sqrt{d_k})V$
而RetNet是：$(Q{K}^T\odot D)V$
除去和尺度、稀疏化等相关的softmax之外，只差了一个$D$，可以看到，这个$D$就是上面公式，合并记忆节点的线性加权权重系数。
相当于一个逐渐衰减的不可训练的相对位置编码。

Retention 的推理循环表示

推理阶段的循环表示论文中定义如下:
$$S_n=\gamma S_{n-1}+{K_n}^T{V}_n$$
$$Retention(X_n)=Q_n{S}_n$$
单个时间步 $n$ 的输出 $o_n$ 的计算公式：
$$o_n=\sum _{m=1}^n{\gamma }^{n-m}(Q_n{e}^{i\theta n})(K_m{e}^{i\theta m}{)}^{†}{v}_m$$
$$=Q_n{e}^{i\theta n}(\sum _{m=1}^n{\gamma }^{n-m}(K_m{e}^{i\theta m}{)}^{†}{v}_m)$$
$$=Q_n{e}^{i\theta n}({\gamma }^{n-n}(K_n{e}^{i\theta n}{)}^{†}{v}_n+(\sum _{m=1}^{n-1}{\gamma }^{n-m}(K_m{e}^{i\theta m}{)}^{†}{v}_m)$$
$$=Q_n{e}^{i\theta n}((K_n{e}^{i\theta n}{)}^{†}{v}_n+(\sum _{m=1}^{n-1}{\gamma }^{n-m}(K_m{e}^{i\theta m}{)}^{†}{v}_m)$$
$$=Q_n{e}^{i\theta n}((K_n{e}^{i\theta n}{)}^{†}{v}_n+\gamma (K_{n-1}{e}^{i\theta (n-1)}{)}^{†}{v}_n+\sum _{m=1}^{n-2}{\gamma }^{n-m-1}(K_m{e}^{i\theta m}{)}^{†}{v}_m)$$
上述公式最后一步和推理阶段循环表示公式中各个元素的对应关系是：
$$Q_n=Q_n{e}^{i\theta n}$$
$$S_{n-1}=(K_{n-1}{e}^{i\theta (n-1)}{)}^{†}{v}_n+\sum _{m=1}^{n-2}{\gamma }^{n-m-1}(K_m{e}^{i\theta m}{)}^{†}{v}_m$$
$${K_n}^T{V}_n=(K_n{e}^{i\theta n}{)}^{†}{v}_n$$
对应论文中的图示：

图中的 $GN$ 表示 GroupNorm。
普通Transformer在解码过程中都有Key和Value，所以解码需要$O(n)$复杂度。
而RetNet解码只需要$O(1)$复杂度，这是因为其将前序节点的表示做成了一个向量，类似于RNN一样循环更新，用Key和Value计算出一个向量之后，去更新权重向量$S$,相当于RNN中的记忆向量。而输出就是Query与$S$的计算结果。
标准RNN是当前输入$X_t$和记忆向量$S_{t-1}$接起来过全连接层，即$S_t=f(U\cdot X_t+W\cdot S_{t-1})$，RetNet是用一个实数$\gamma$做权重，进行线性加权。
可以看到在推理阶段，RetNet 在计算当前时间步 $n$ 的输出 $O_n$ 只依赖于上一个时间步产出的状态矩阵 $S_{n-1}$。
其实就是把计算顺序改了一下，先计算的 $K_n$ 和 $V_n$ 的相乘然后一直累加到状态矩阵 $S_n$ 上，最后再和 $Q_n$ 相乘。
而不是像 Transformer 架构那样，每个时间步的计算要先算 $Q_n$ 和前面所有时间步的 $K$ 相乘得到 attention 权值再和 $V$ 相乘求和，这样就需要一直保留历史的 $K$ 和 $V$。

Gated Multi-Scale Retention

然后 RetNet 每一层中的 Retention 子模块其实也是分了 $h$ 个头，每个头用不同的 $W_Q,W_K,W_V\in {\mathbb{R}}^{d\times d}$ 参数，同时每个头都采用不同的 $\gamma$ 常量，这也是 Multi-Scale Retention 名称的来由。
则对输入 $X$， MSR 层的输出是：
$$\gamma =1-2^{-5-arange(0,h)}\in {\mathbb{R}}^h$$
$$hea{d}_i=Retention(X,{\gamma }_i)$$
$$Y=GroupNor{m}_h(Concat(hea{d}_1,...,hea{d}_h))$$
$$MSR(X)=(swish(X{W}_G)\odot Y)W_O$$
其中$W_G,W_O\in {\mathbb{R}}^{d\ast h\times d\ast h}$ ， $swish$ 是激活函数用来生成门控阈值，还有由于每个头均采用不同的 $\gamma$，所以每个头的输出要单独做 normalize 之后再 concat。