Arxiv 2307 | Retentive Network: A Successor to Transformer for Large Language Models

Retentive Network: A Successor to Transformer for Large Language Models

• 论文：https://arxiv.org/abs/2307.08621
• 代码：https://github.com/microsoft/unilm/tree/master/retnet

本文从序列建模的角度，构建了一种类似Transformer且更加高效的结构。在语言任务上展现出了良好的效率和性能。

• 利用类似于Transformer的并行组件实现了对于GPU并行能力的利用。
• 利用循环机制确保了$O(1)$级别的存储和计算复杂度。
• 利用分块循环策略从而执行有效的长序列建模。

实际中，并行编码每个局部的块来加速计算，同时循环编码全局块来节省显存。

序列建模

对于输入的长度为$N$的文本嵌入序列，由于其本身信息的前后依赖关系和因果关系的需求，所以本文是从循环模型的角度开始构建模型的。

基础的迭代形式：

对于第n次迭代的输入$X_n$，有

$$Q_n=X_n\cdot W_Q,K_n=X_n\cdot W_K,V_n=X_n\cdot W_V\in {\mathbb{R}}^{1\times d}$$

将序列建模认为成通过状态$S_n$，将$V(n)$映射为$O(n)$**的过程。**于是可以得到下式：

$$S_n=A{s}_{n-1}+K_n^{\top }{V}_n=A^{n-1}{K}_1^{\top }{V}_1+A^{n-2}{K}_2^{\top }{V}_2+\cdots +K_n^{\top }{V}_n=\sum _{m=1}^n{A}^{n-m}{K}_m^{\top }{V}_m$$

这里的$A\in {\mathbb{R}}^{d\times d}$描述了各个位置之间的相对关系。

$O_n=Q_n{S}_n={\sum }_{m=1}^n{Q}_n{A}^{n-m}{K}_m^{\top }{V}_m,Q_n\in {\mathbb{R}}^{1\times d}$

Parallel Retention

通过设置一个特殊的矩阵$A$，将其对角化处理获得$A=\Lambda (\gamma e^{i\theta }){\Lambda }^{-1}$，这里的两个矩阵$\Lambda$由于在公式中紧邻$Q_n,K_n$，所以可以将其合并到二者各自的权重矩阵$W_Q,W_K$中一同随着网络去学习，从而上式可以改写：

$$O_n=Q_n{S}_n=\sum _{m=1}^n{Q}_n(\gamma e^{i\theta }{)}^{n-m}{K}_m^{\top }{V}_m=\sum _{m=1}^n[Q_n(\gamma e^{i\theta }{)}^n][K_m(\gamma e^{i\theta }{)}^{-m}{]}^{\top }{V}_m=\sum _{m=1}^n{\gamma }^{n-m}(Q_n{e}^{in\theta })(K_m{e}^{im\theta }{)}^{†}{V}_m$$

这里将指数与转置融合和获得共轭转置。这里的复数系数实际上可以看做是一种位置嵌入，由于这里的计算反映出了与位置n和m的关联，所以可以认为是一种相对位置关系的表示。

由于这里$Q,K$索引上的独立性，所以很容易改为并行的基于矩阵运算的结构。将复数矩阵系数极其共轭形式分别合并到$Q,K$计算过程中，从而可以得到：

$$\begin{gathered} Q=(X{W}_Q)\odot \Theta ,K=(X{W}_K)\odot \bar{\Theta },V=X{W}_V,{\Theta }_n=e^{in\theta } \\ {D}_{nm}={\gamma }^{n-m}\text{ if }n\ge m\text{ else }0 \end{gathered}$$

从而得到整体模块的计算过程：

$$Retention(X)=(Q{K}^{\top }\odot D)V,D\in {\mathbb{R}}^{N\times N}$$

def ParallelRetention(
    q, # bsz ∗ num_head ∗ len ∗ qk_dim
    k, # bsz ∗ num_head ∗ len ∗ qk_dim
    v, # bsz ∗ num_head ∗ len ∗ v_dim
    decay_mask # num_head ∗ len ∗ len
 ):
     retention = q @ k.transpose(−1, −2)
     retention = retention ∗ decay_mask
     output = retention @ v
     output = group_norm(output)
     return output

这一形式实际上与Transformer的带mask的计算形式非常类似。

这里由于有$Q{K}^{\top }$，使用了三种归一化方式来提升数值精度，这些归一化策略实际上都是在GN输入上乘以了一个常数，而由于GN本身的尺度不变性，所以必不会影响GN的输出和反向的梯度。

• 使用特征维度归一化$Q{K}^{\top }/\sqrt{d}$
• 设置 D = { D n m ∑ i = 1 n D n i } D = \{\frac{D_{nm}}{\sqrt{\sum^n_{i=1}D_{ni}}}\} D={∑i=1n​Dni​ ​Dnm​​}
• 假定$R=Q{K}^{\top }\odot D$，设置 R = { R n m max ⁡ ( ∣ ∑ i = 1 n R n i ∣ , 1 ) } R = \{ \frac{R_{nm}}{\max(|\sum^{n}_{i=1} R_{ni}|, 1)} \} R={max(∣∑i=1n​Rni​∣,1)Rnm​​}

Recurrent Retention

但是，如果从序列形式的角度来看，前面的最一开始的建模过程也可以改写成另外一种类似于RNN的形式。
先将状态参数写成迭代形式：

$$S_n=\gamma S_{n-1}+K_n^{\top }{V}_n\in {\mathbb{R}}^{d\times d}$$

最终可以得到整体迭代计算过程：

R e t e n t i o n ( X n ) = Q n S n , n ∈ { 1 , … , N } Retention(X_n) = Q_n S_n, n \in \{1,\dots,N\} Retention(Xn​)=Qn​Sn​,n∈{1,…,N}

def RecurrentRetention(
    q, k, v, # bsz ∗ num_head ∗ len ∗ qkv_dim
    past_kv, # bsz ∗ num_head ∗ qk_dim ∗ v_dim
    decay # num_head ∗ 1 ∗ 1
):
    current_kv = decay ∗ past_kv + k.unsqueeze(−1) ∗ v.unsqueeze(−2)
    output = torch.sum(q.unsqueeze(−1) ∗ current_kv, dim=−2)
    output = group_norm(output)
    return output, current_kv

实际上这里的形式与线性Attention先计算KV的思路颇有相通之处。

Chunkwise Recurrent Retention

作者也提出了一种将上述两种形式进行混合的形式，通过将序列划分为连续的块，块内部执行并行形式的处理，块之间执行循环处理，实际的，对于第$i$个块，处理形式如下：

$$Retention(X_{[i]})=\underset{块内并行}{\underbrace{(Q_{[i]}{K}_{[i]}^{\top }\odot D)V_{[i]}}}+\underset{块间循环}{\underbrace{(Q_{[i]}{S}_i)\odot \xi }},{\xi }_{ij}={\gamma }^{i+1}$$

def ChunkwiseRetention(
    q, k, v, # bsz ∗ num_head ∗ chunk_size ∗ qkv_dim
    past_kv, # bsz ∗ num_head ∗ qk_dim ∗ v_dim
    decay_mask, # num_head ∗ chunk_size ∗ chunk_size
    chunk_decay, # num_head ∗ 1 ∗ 1
    inner_decay, # num_head ∗ chunk_size
):
    retention = q @ k.transpose(−1, −2)
    retention = retention ∗ decay_mask
    inner_retention = retention @ v
    cross_retention = (q @ past_kv) ∗ inner_decay
    retention = inner_retention + cross_retention
    output = group_norm(retention)
    current_kv = chunk_decay ∗ past_kv + k.transpose(−1, −2) @ v
    return output, current_kv

Gated Multi-Scale Retention

引入“头”的概念，为不同的头使用不同的变换权重，同时为不同的头分配不同的$\gamma$。

由于不同的头引入了不同的参数$\gamma$，所以实际的方差统计量会有所差异。所以这里使用GroupNorm独立归一化不同的头。

实际效果

仅在文本任务上进行了实验。