Retentive Network: A Successor to Transformer
for Large Language Models
保留网络：大型语言模型的继任者——Transformer

As shown in Figure 3(a), the retention layer is defined as:

where $\overline{\Theta}$ is the complex conjugate of $\Theta$, and $D\in\mathbb{R}^{|x|\times|x|}$ combines causal masking and exponential decay along relative distance as one matrix. Similar to self-attention, the parallel representation enables us to train the models with GPUs efficiently.

As shown in Figure 3(b), the proposed mechanism can also be written as recurrent neural networks (RNNs), which is favorable for inference. For the $n$-th timestep, we recurrently obtain the output as:

where $Q,K,V,\gamma$ are the same as in Equation (5).

A hybrid form of parallel representation and recurrent representation is available to accelerate training, especially for long sequences. We divide the input sequences into chunks. Within each chunk, we follow the parallel representation (Equation (5)) to conduct computation. In contrast, cross-chunk information is passed following the recurrent representation (Equation (LABEL:eq:ret:recurrent)). Specifically, let $B$ denote the chunk length. We compute the retention output of the $i$-th chunk via:

where ${[i]}$ indicates the $i$-th chunk, i.e., $x_{[i]}=[x_{(i-1)B+1},\cdots,x_{iB}]$.

We utilize the scale-invariant nature of $\mathrm{GroupNorm}$ to improve the numerical precision of retention layers. Specifically, multiplying a scalar value within $\mathrm{GroupNorm}$ does not affect outputs and backward gradients, i.e., $\mathrm{GroupNorm}(\alpha*\mathrm{head}_{i})=\mathrm{GroupNorm}(\mathrm{head}_{i})$. We implement three normalization factors in Equation (5). First, we normalize $QK^{\intercal}$ as $\nicefrac{{QK^{\intercal}}}{{\sqrt{d}}}$. Second, we replace $D$ with $\tilde{D}_{nm}=\nicefrac{{D_{nm}}}{{\sqrt{\sum_{i=1}^{n}D_{ni}}}}$. Third, let $R$ denote the retention scores $R=QK^{\intercal}\odot D$, we normalize it as $\tilde{R}_{nm}=\nicefrac{{R_{nm}}}{{\max(|\sum_{i=1}^{n}R_{ni}|,1)}}$. Then the retention output becomes $\mathrm{Retention}(X)=\tilde{R}V$. The above tricks do not affect the final results while stabilizing the numerical flow of both forward and backward passes, because of the scale-invariant property.

We use the parallel (Equation (5)) and chunkwise recurrent (Equation (7)) representations during the training process. The parallelization within sequences or chunks efficiently utilizes GPUs to accelerate computation. More favorably, chunkwise recurrence is especially useful for long-sequence training, which is efficient in terms of both FLOPs and memory consumption.

The recurrent representation (Equation (LABEL:eq:ret:recurrent)) is employed during the inference, which nicely fits autoregressive decoding. The $O(1)$ complexity reduces memory and inference latency while achieving equivalent results.

The parallel representation of retention shares similar spirits as Transformers [34]. The most related Transformer variant is Lex Transformer [29] which implements xPos as position embeddings. As described in Equation (3), the derivation of retention aligns with xPos. In comparison with attention, retention removes $\mathrm{softmax}$ and enables recurrent formulation, which significantly benefits inference.

Unlike Equation (2), if $Q_{n}$ and $K_{n}$ are content-unaware, the formulation can be degenerated to S4 [11], where $O=(QK^{\intercal},QAK^{\intercal},..,QA^{|x|-1}K^{\intercal})*V$.

The variants typically use various kernels $\nicefrac{{\phi(q_{i})\phi(k_{j})}}{{\sum_{n=1}^{|x|}\phi(q_{i})\phi(k_{n})}}$ to replace the $\mathrm{softmax}$ function. However, linear attention struggles to effectively encode position information, rendering the models less performant. Besides, we reexamine sequence modeling from scratch, rather than aiming at approximating $\mathrm{softmax}$.

Attention Free Transformer (AFT) simplifies dot-product attention to element-wise operations and moves $\mathrm{softmax}$ to key vectors. RWKV replaces AFT’s position embeddings with exponential decay and runs the models recurrently for training and inference. In comparison, retention preserves high-dimensional states to encode sequence information, which contributes to expressive ability and better performance.

Compared with relative position embedding methods proposed for Transformers, Equation 3 presents a similar formulation as xPos [29] and RoPE [31].

As shown in Equation (8), the retention layer uses Sub-LayerNorm [37] to normalize outputs. Because the multi-scale modeling leads to different variances for the heads, we replace the original LayerNorm with GroupNorm.

We re-allocate the parameters in MSR and FFN for fair comparisons. Let $d$ denote $d_{\text{model}}$ for simplicity here. In Transformers, there are about $4d^{2}$ parameters in self-attention where $W_{Q},W_{K},W_{V},W_{O}\in\mathbb{R}^{d\times d}$, and $8d^{2}$ parameters in FFN where the intermediate dimension is $4d$. In comparison, RetNet has $8d^{2}$ parameters in retention, where $W_{Q},W_{K}\in\mathbb{R}^{d\times d},W_{G},W_{V}\in\mathbb{R}^{d\times 2d},W_{O}\in\mathbb{R}^{2d\times d}$. Notice that the head dimension of $V$ is twice $Q,K$. The widened dimension is projected back to $d$ by $W_{O}$. In order to keep the parameter number the same as Transformer, the FFN intermediate dimension in RetNet is $2d$. Meanwhile, we set the head dimension to $256$ in our experiments, i.e., $256$ for queries and keys, and $512$ for values. For fair comparison, we keep $\mathbf{\gamma}$ identical among different model sizes, where $\mathbf{\gamma}=1-e^{\mathrm{linspace}(\log\nicefrac{{1}}{{32}},\log\nicefrac{{1}}{{512}},h)}\in\mathbb{R}^{h}$ instead of the default value in Equation (8).

As shown in Table 2, we train language models with various sizes (i.e., 1.3B, 2.7B, and 6.7B) from scratch. The training corpus is a curated compilation of The Pile [10], C4 [9], and The Stack [18]. We append the <bos> token to indicate the start of a sequence¹¹1We find that appending the <bos> token at the beginning benefits training stability and performance.. The training batch size is 4M tokens with 2048 maximal length. We train the models with 100B tokens, i.e., 25k steps. We use the AdamW [21] optimizer with $\beta_{1}=0.9,\beta_{2}=0.98$, and weight decay is set to $0.05$. The number of warmup steps is 375 with linear learning rate decay. The parameters are initialized following DeepNet [36] to guarantee training stability. The implementation is based on TorchScale [23]. We train the models with 512 AMD MI200 GPUs.

As shown in Figure 5, we report perplexity on the validation set for the language models based on Transformer and RetNet. We present the scaling curves with three model sizes, i.e., 1.3B, 2.7B, and 6.7B. RetNet achieves comparable results with Transformers. More importantly, the results indicate that RetNet is favorable regarding size scaling. Besides performance, the RetNet training is quite stable in our experiments. Experimental results show that RetNet is a strong competitor to Transformer for large language models. Empirically, we find that RetNet starts to outperform Transformer when the model size is larger than 2B. We also summarize the language modeling results with different context lengths in Appendix B.

We also compare the language models on a wide range of downstream tasks. We evaluate zero-shot and 4-shot learning with the 6.7B models. As shown in Table 3, the datasets include HellaSwag (HS) [39], BoolQ [6], COPA [38], PIQA [3], Winograd, Winogrande [20], and StoryCloze (SC) [22]. The accuracy numbers are consistent with language modeling perplexity presented in Figure 5. RetNet achieves comparable performance with Transformer on zero-shot and in-context learning settings.

As shown in Figure 6(a), the memory cost of Transformer increases linearly due to KV caches. In contrast, the memory consumption of RetNet remains consistent even for long sequences, requiring much less GPU memory to host RetNet. The additional memory consumption of RetNet is almost negligible (i.e., about 3%) while the model weights occupy 97%.

As presented in Figure 6(b), the throughput of Transformer drops along with the decoding length increases. In comparison, RetNet has higher and length-invariant throughput during decoding, by utilizing the recurrent representation of retention.

Latency is an important metric in deployment, which greatly affects user experience. We report decoding latency in Figure 6(c). Experimental results show that increasing batch size renders Transformer’s latency larger. Moreover, the latency of Transformers grows faster with longer input. In order to make latency acceptable, we have to restrict the batch size, which harms the overall inference throughput of Transformers. By contrast, RetNet’s decoding latency outperforms Transformers and keeps almost the same across different batch sizes and input lengths.

We ablate the $\mathrm{swish}$ gate and $\mathrm{GroupNorm}$ as described in Equation 8. Table 6 shows that the above two components improve the final performance. Firstly, the gating module is essential for enhancing non-linearity and improving model capability. Notice that we use the same parameter allocation as Transformers after removing the gate. Secondly, group normalization in retention balances the variances of multi-head outputs, which improves training stability and language modeling results.

Equation 8 shows that we use different $\mathbf{\gamma}$ as the decay rates for the retention heads. In the ablation studies, we examine removing $\gamma$ decay (i.e., “$-$ $\gamma$ decay”) and applying the same decay rate across heads (i.e., “$-$ multi-scale decay”). Specifically, ablating $\gamma$ decay is equivalent to $\gamma=1$. In the second setting, we set $\gamma=127/128$ for all heads. Table 6 indicates that both the decay mechanism and using multiple decay rates can improve the language modeling performance.

From the recurrent perspective of LABEL:eq:rnn, the head dimension implies the memory capacity of hidden states. In the ablation study, we reduce the default head dimension from $256$ to $64$, i.e., $64$ for queries and keys, and $128$ for values. We keep the hidden dimension $d_{\text{model}}$ the same so the number of heads increases. Experimental results in Table 6 show that the larger head dimension achieves better performance.