Retentive Network

Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 11:20:13

Abstract: In this work, we propose Retentive Network (RetNet) as a foundation architecture for large language models, simultaneously achieving training parallelism, low-cost inference, and good performance. We theoretically derive the connection between recurrence and attention. Then we propose the retention mechanism for sequence modeling, which supports three computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent. Specifically, the parallel representation allows for training parallelism. The recurrent representation enables low-cost 
 inference, which improves decoding throughput, latency, and GPU memory without sacrificing performance. The chunkwise recurrent representation facilitates efficient long-sequence modeling with linear complexity, where each chunk is encoded parallelly while recurrently summarizing the chunks. Experimental results on language modeling show that RetNet achieves favorable scaling results, parallel training, low-cost deployment, and efficient inference.

Main paper: 
 introduction: 
 Transformer [51] has become the de facto architecture for large language models, which was initially proposed to overcome the sequential training issue of recurrent models [25] . However, training parallelism of Transformers is at the cost of inefficient inference, because of the O(N ) complexity per step and memory-bound key-value cache [42] , which renders Transformers unfriendly to deployment.The growing sequence length increases GPU memory consumption as well as latency and reduces inference speed. Numerous efforts have continued to develop the next-generation architecture, aiming at retaining training parallelism and competitive performance as Transformers while having efficient O(1) inference. It is challenging to achieve the above goals simultaneously.There have been three main strands of research. First, linearized attention [27, 37] approximates standard attention scores exp(q • k) with kernels ϕ(q) • ϕ(k), so that autoregressive inference can be rewritten in a recurrent form. However, the modeling capability and performance are worse than Transformers, which hinders the method's popularity. The second strand returns to recurrent models for efficient inference while sacrificing training parallelism. As a remedy, element-wise operators [36] are used for acceleration, however, representation capacity and performance are harmed. The third line explores replacing attention with other mechanisms, such as S4 [20] , and its variants [11, 38] .None of the previous work can achieve strong performance and efficient inference at the same time compared to Transformers.In this work, we propose retentive networks (RetNet), achieving low-cost inference, efficient longsequence modeling, Transformer-comparable performance, and parallel model training simultaneously. Specifically, we introduce a multi-scale retention mechanism to substitute multi-head attention, which has three computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent representations. First, the parallel representation empowers training parallelism to utilize GPU devices fully. Second, the recurrent representation enables efficient O(1) inference in terms of memory and computation. The deployment cost and latency can be significantly reduced. Moreover, the implementation is greatly simplified without key-value cache tricks. Third, the chunkwise recurrent Submitted to 38th Conference on Neural Information Processing Systems (NeurIPS 2024). Do not distribute.representation can perform efficient long-sequence modeling. We parallelly encode each local block for computation speed while recurrently encoding the global blocks to save GPU memory.We compare RetNet with Transformer and its variants. Experimental results on language modeling show that RetNet is consistently competitive in terms of both scaling curves and in-context learning.Moreover, the inference cost of RetNet is length-invariant. For a 7B model and 8k sequence length, RetNet decodes 8.4× faster and saves 70% of memory than Transformers with key-value caches. During training, RetNet also achieves 3× acceleration than standard Transformer with highly-optimized FlashAttention-2 [10] . Besides, RetNet's inference latency is insensitive to batch size, allowing enormous throughput. The intriguing properties make RetNet a potential candidate to replace Transformer for large language models. 
retentive network: 
 Retentive network (RetNet) is stacked with L identical blocks, which follows a similar layout (i.e., residual connection, and pre-LayerNorm) as in Transformer [51] . Each RetNet block contains two modules: a multi-scale retention (MSR) module, and a feed-forward network (FFN) module. We introduce the MSR module in the following sections. Given an input sequence x = x 1 • • • x |x| , RetNet encodes the sequence in an autoregressive way. The input vectors{x i } |x| i=1 is first packed into X 0 = [x 1 , • • • , x |x| ] ∈ R |x|×dmodel, where d model is hidden dimension. Then we compute contextualized vector representations X l = RetNet l (X l-1 ), l ∈ [1, L].In this section, we introduce the retention mechanism that has a dual form of recurrence and parallelism. So we can train the models in a parallel way while recurrently conducting inference.Consider a sequence modeling problem that maps v(n) → o(n) through states s n . Let v n , o n denote v(n), o(n) for simplicity. We formulate the mapping in a recurrent manner:EQUATIONwhere we map v n to the state vector s n , and then implement a linear transform to encode sequence information recurrently. Next, we make the projection Q n , K n content-aware:EQUATIONwhere W Q , W K ∈ R d×d are learnable matrices.We diagonalize the matrix A = Λ(γe iθ )Λ -1 , where γ, θ ∈ R d . Then we obtain A n-m = Λ(γe iθ ) n-m Λ -1 . By absorbing Λ into W Q and W K , we can rewrite Equation (1) as:EQUATIONwhere Q n (γe iθ ) n , K m (γe iθ ) -m is known as xPos [45] , i.e., a relative position embedding proposed for Transformer. We further simplify γ as a scalar, Equation (3) becomes:EQUATION)where † is the conjugate transpose. The formulation is easily parallelizable within training instances.In summary, we start with recurrent modeling as shown in Equation ( 1), and then derive its parallel formulation in Equation ( 4). We consider the original mapping v(n) → o(n) as vectors and obtain the retention mechanism as follows. The Parallel Representation of Retention As shown in Figure 1a , the retention layer is defined as:Q = (XW Q ) ⊙ Θ, K = (XW K ) ⊙ Θ, V = XW V Θ n = e inθ , D nm = γ n-m , n ≥ m 0, n < m Retention(X) = (QK ⊺ ⊙ D)V (5)where D ∈ R |x|×|x| combines causal masking and exponential decay along relative distance as one matrix, and Θ is the complex conjugate of Θ. In practice, we map Q, K ∈ R d → C d/2 , add the complex position embedding Θ, then map them back to R d , following the implementation trick as in LLaMA [48, 44] . Similar to self-attention, the parallel representation enables us to train the models with GPUs efficiently.The Recurrent Representation of Retention As shown in Figure 1b , 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:EQUATIONwhere Q, K, V, γ are the same as in Equation (5) .The Chunkwise Recurrent Representation of Retention 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 ( 6)). Specifically, let B denote the chunk length. We compute the retention output of the i-th chunk via:EQUATIONwhere[i] indicates the i-th chunk, i.e., x [i] = [x (i-1)B+1 , • • • , x iB ].The proof of the equivalence between recurrent representation and chunkwise recurrent representation is described in Appendix B.We use h = dmodel /d retention heads in each layer, where d is the head dimension. The heads use 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.different parameter matrices W Q , W K , W V ∈ R d×d .transpose(-1, -2) @ v return output, current_kv Chunkwise retention combines the above advantages (i.e., parallel within each chunk and recurrent across chunks), which has linear memory complexity for long sequences. different γ for each head. For simplicity, we set γ identical among different layers and keep them fixed. In addition, we add a swish gate [23, 40] to increase the non-linearity of retention layers.Formally, given input X, we define the layer as:γ = 1 -2 -5-arange(0,h) ∈ R h head i = Retention(X, γ i ) Y = GroupNorm h (Concat(head 1 , • • • , head h )) MSR(X) = (swish(XW G ) ⊙ Y )W O (8)where W G , W O ∈ R dmodel×dmodel are learnable parameters, and GroupNorm [53] normalizes the output of each head, following SubLN proposed in [43] . Notice that the heads use multiple γ scales, which results in different variance statistics. So we normalize the head outputs separately.The pseudocode of retention is summarized in Figure 2 .Retention Score Normalization We utilize the scale-invariant nature of GroupNorm to improve the numerical precision of retention layers. Specifically, multiplying a scalar value within GroupNorm does not affect outputs and backward gradients, i.e., GroupNorm(α * head i ) = GroupNorm(head i ).We implement three normalization factors in Equation (5) . First, we normalize QK ⊺ as QK ⊺ / √ d.Second, we replace D with Dnm = Dnm / √ n i=1 Dni. Third, let R denote the retention scores R = QK ⊺ ⊙ D, we normalize it as Rnm = Rnm /max( n i=1 |Rni|,1). Then the retention output becomes Retention(X) = RV . 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.For an L-layer retention network, we stack multi-scale retention (MSR) and feed-forward network (FFN) to build the model. Formally, the input sequence {x i } |x| i=1 is transformed into vectors by a word embedding layer. We use the packed embeddingsX 0 = [x 1 , • • • , x |x| ] ∈ R |x|×dmodel as theinput and compute the model output X L :EQUATIONwhere LN(•) is LayerNorm [3] . The FFN part is computed as FFN(X) = gelu(XW 1 )W 2 , whereW 1 , W 2 are parameter matrices.Training 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.Inference The recurrent representation (Equation ( 6)) is employed during inference, which nicely fits autoregressive decoding. The O(1) complexity reduces memory and inference latency while achieving equivalent results.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: Limitations are discussed in the paper.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.• The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.• The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.• The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.• The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.• If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.• While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
experiments: 
 We perform language modeling experiments to evaluate RetNet. First, we present the scaling curves of Transformer and RetNet. Second, we follow the training settings of StableLM-4E1T [50] to compare with open-source Transformer models in downstream benchmarks. Moreover, for training and inference, we compare speed, memory consumption, and latency. The training corpus is a curated compilation of The Pile [16] , C4 [14] , and The Stack [29] .We compare RetNet with various efficient Transformer variants, including RWKV [36] , H3 [11] ,Hyena [38] , and Mamba [19] . We use LLaMA [48] architecture, including RMSNorm [59] and SwiGLU [40, 7] module, as the Transformer backbone, which shows better performance and stability.Consequently, other variants follow these settings. Specifically, Mamba does not have FFN layers so we only implement RMSNorm. For RetNet, the FFN intermediate dimension is 5 3 d and the value dimensions in W G , W V , W O are also 5 3 d, where the overall parameters are still 12d 2 . All models have 400M parameters with 24 layers and a hidden dimension of 1024. For H3, we set the head dimension to 8. For RWKV, we use the TimeMix module to substitute self-attention layers while keeping FFN layers consistent with other models for fair comparisons. We train the models with 40k steps with a batch size of 0.25M tokens.As shown in Table 1 , we first report the language modeling perplexity of validation sets. Besides the overall validation set, following [2] , we divide perplexity into "AR-Hit" and "First Occur". Specifically, AR-Hit contains the predicted tokens that are previously seen bigrams in the previous context, which evaluates the associative recall ability."First Occur" has the predicted tokens that can not be recalled from the context. Among various Transformer variants, RetNet outperforms previous methods on both "AR-Hit" and "First Occur" splits, which is important for real-world use cases.We also evaluate Massive Multitask Language Understanding (MMLU; [24] ) answer perplexity to evaluate models on knowledge-intensive tasks. We report the average perplexity of the correct answers, i.e., given input [Question, "Answer:", Correct Answer], we calculate the perplexity of the "Correct Answer" part. RetNet achieves competitive results among the architectures. 1 : Perplexity results on language modeling and MMLU [24] answers. We use the augmented Transformer architecture proposed in LLaMA [48] for reference. For language modeling, we report perplexity on both the overall validation set and fine-grained diagnosis sets [2] , i.e., "AR-Hit" evaluates the associative recall capability, and "First-Occur" indicates the regular language modeling performance. Besides, we evaluate the answer perplexity of MMLU subsets. We train language models with various sizes (i.e., 1.3B, 2.7B, and 6.7B) from scratch. The training batch size is 4M tokens with 2048 maximal length. We train the models with 25k steps. The detailed hyper-parameters are described in Appendix E. We train the models with 512 AMD MI200 GPUs. We evaluate long-context modeling on the ZeroSCROLLS [41] benchmark. We train a hybrid model of size 2.7B, RetNet+, which stacks the attention and retention layers. Specifically, we insert one attention layer after every 3 retention layers. We follow most configurations of the 2.7B model as in Section 3.2. We scale the number of training tokens to 420B tokens. The batch size is 4M tokens.We first train the model with 4K length and then extend the sequence length to 16K for the last 50B training tokens. The rotation base scaling [55] is used for length extension. As shown in Figure 5 , we compare memory cost, throughput, and latency of Transformer and RetNet during inference. Transformers reuse KV caches of previously decoded tokens. RetNet uses the recurrent representation as described in Equation ( 6). We evaluate the 6.7B model on the A100-80GB GPU. Figure 5 shows that RetNet outperforms Transformer in terms of inference cost.Memory As shown in Figure 5a , the memory cost of Transformer increases linearly due to KV caches. In contrast, the memory consumption of RetNet remains consistent even for long sequences, Latency Latency is an important metric in deployment that greatly affects the user experience.We report the decoding latency in Figure 5c . Experimental results show that increasing batch size renders the 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 stays almost the same across different batch sizes and input lengths. where the hidden dimension is 3072 and the layer size is 28. We use highly optimized FlashAttention-2 [10] for Transformers. In comparison, we implement chunk recurrent representation (Equation ( 7)) using Triton [46] , where the computation is both memory-friendly and computationally efficient. The chunk size is set to 256.We evaluate the results with eight Nvidia H100-80GB GPUs because FlashAttention-2 is highly optimized for H100 cards.Experimental results show that RetNet has higher training throughput than Transformers. The acceleration ratio increases as the sequence length is longer.When the training length is 64k, RetNet's throughput is about 3 times than Transformer's.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 2 , the datasets include HellaSwag (HS; [57] ), BoolQ [8] , COPA [52] , PIQA [6] , Winograd, Winogrande [30] , and StoryCloze (SC; [34] ).The accuracy numbers are consistent with language modeling perplexity presented in Figure 3 . RetNet achieves comparable performance with Transformer on zero-shot and in-context learning settings.We ablate various design choices of RetNet and report the language modeling results in Architecture We ablate the swish gate and GroupNorm as described in Equation (8) . Table 3 shows that the above two components improve performance. First, the gating module is essential for enhancing non-linearity and improving model capability. Notice that we use the same parameter allocation as in Transformers after removing the gate. Second, group normalization in retention balances the variances of multi-head outputs, which improves training stability and language modeling results.Multi-Scale Decay Equation ( 8) shows that we use different γ as the decay rates for the retention heads. In the ablation studies, we examine removing γ decay (i.e., "-γ decay") and applying the same decay rate across heads (i.e., "-multi-scale decay"). Specifically, ablating γ decay is equivalent to γ = 1. In the second setting, we set γ = 1 -2 -6.5 for all heads. Table 3 indicates that both the decay mechanism and using multiple decay rates can improve the language modeling performance.Head Dimension As indicated by the recurrent perspective of Equation ( 1), the head dimension implies the memory capacity of hidden states. In ablation, we reduce the default head dimension from 256 to 64, i.e., 64 for queries and keys, and ⌊ 5 3 ×64⌋ ≈ 108 for values. We keep the hidden dimension d model the same. Accordingly, we adjust the multi-scale decay as γ = 1 -2 -5-arange(0,h)/4 to keep the same decay range. Table 3 shows that the larger head dimension achieves better performance. Table 3 : Perplexity results on language modeling and MMLU [24] answers. For language modeling, we report perplexity on both the overall validation set and fine-grained diagnosis sets [2] , i.e., "AR-Hit" evaluates the associative recall capability, and "First-Occur" indicates the regular language modeling performance. Besides, we evaluate the answer perplexity of the MMLU subsets.We also compare RetNet with vision Transformers [15, 47] in Table 4 , where bidirectional encoders are evaluated. Unlike causal language models, the vision encoders do not require recurrent representations. Specifically, we use retention as follows:Q = (XW Q ) ⊙ Θ, K = (XW K ) ⊙ Θ, V = XW V Retention(X) = (QK ⊺ )V = Q(K ⊺ V )where multi-scale decay is removed in bidirectional computation. Notice that we can compute retention in different orders. Similar to linear attention [27] , the Q(K ⊺ V ) paradigm is an efficient operator in bidirectional settings, especially for high-resolution images.We perform experiments on ImageNet-1K classification [13] , COCO object detection [32] , and ADE20K semantic segmentation [60] . We compare RetNet with DeiT [47] which is a well-tuned vision Transformer. Besides, we follow [21] and plug in a depth-wise convolution in experiments.We we use AdamW [33] for 300 epochs, and 20 epochs of linear warm-up. The learning rate is 1 × 10 -3 , the batch size is 1024, and the weight decay is 0.05. For COCO object detection, we use Mask R-CNN [22] as the task head, and the above models pre-trained on ImageNet as the backbone with 3x schedules. In ADE20K experiments, we use UperNet [54] as the segmentation head. The detailed configuration can be found in Appendix H.Table 4 shows the results across various vision tasks. RetNet is competitive compared with DeiT.For classification and segmentation, RetNet is slightly better than DeiT, where RetNet achieves 0.81% accuracy improvement on ImageNet and 0.61% mIoU improvement on ADE20K. For object detection, the results are comparable.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA]Answer: [Yes] Justification: Code will be released in camera-ready version. All of the data we use is public-available.Guidelines:• The answer NA means that paper does not include experiments requiring code.• Please see the NeurIPS code and data submission guidelines (https://nips.cc/ public/guides/CodeSubmissionPolicy) for more details.• While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).• The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://nips.cc/public/guides/CodeSubmissionPolicy) for more details.• The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.• The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.• At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).• Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
related work: 
 Numerous efforts are focused on reducing the quadratic complexity of attention mechanisms. Linear attention [27] uses various kernels ϕ(qi)ϕ(kj ) / |x| n=1 ϕ(qi)ϕ(kn) to replace the softmax function. In contrast, we reexamine sequence modeling from scratch, rather than aiming at approximating softmax. AFT [58] simplifies dot-product attention to element-wise and moves softmax to key vectors. RWKV [36] 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.S4 [20] unifies convolution and recurrence format and achieves O(N log N ) training complexity leveraging the FFT kernel. Unlike Equation (2), if Q n and K n are content-unaware, the formulation can be degenerated to S4 [20] . Hyena [38] generates the convolution kernels, achieving sub-quadratic training efficiency but keeping O(N ) complexity in single-step inference. Recently, most related work has focused on modifying γ in Equation ( 6) as a data-dependent variable, such as Mamba [19] ,GLA [56] , Gateloop [28] , and xLSTM [4] . Another strand explores hybrid architectures [31, 12] that interleave the above components with attention layers.In addition, we discuss the training and inference efficiency of some related methods. Let D denote the hidden dimension, H the head dimension, and N the sequence length.  
conclusion: 
 We propose retentive networks (RetNet) for sequence modeling, which enables various representations, i.e., parallel, recurrent, and chunkwise recurrent. RetNet achieves significantly better inference efficiency (in terms of memory, speed, and latency), favorable training parallelization, and competitive performance compared with Transformers. The above advantages make RetNet an ideal successor to Transformers for large language models, especially considering the deployment benefits brought by the O(1) inference complexity. In the future, we are interested in deploying RetNet on various edge devices, such as mobile phones. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Hyperparameters are attached in the appendix.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.• The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [No] Justification: For large language models, the variance between different runs is negligible.Moreover, the evaluation pipeline is deterministic.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.• The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).• The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)• The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean.• It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.• For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).• If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes]Justification: The corresponding resources are stated in the paper.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.• The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.• The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: We work on fundamental research that has no direct societal impact.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.• Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.• The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.• If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA]Justification: The paper does not pose safety risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.• Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.• We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We carefully follow the licenses of open-source code, data, and models.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL.• The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.• If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: We scale up the number of training tokens to 350B for the 3B-size models. We compare with strong Transformer checkpoints including OpenLLaMA [18] and StableLM [50] . Moreover, we reproduce a Transformer language model (named Transformer Repro ) for apple-to-apple comparison.Our model RetNet+ follows the same configuration as in Section 3. with 1000 warm-up steps and linear learning rate decay. The training corpus includes The Pile [16] and RedPajama [9] . Transformer Repro follows the exact same setting. We illustrate the equivalence between the recurrent representation and the chunk-wise recurrent representation. Specifically, let B denote the chunk length. For the output O n , n can be divided as n = kB + r where B is the chunk size. Following Equation 6, we have:EQUATIONwhereΓ i = γ n-i , ζ ij = γ B-m , and [i] indicates the i-th chunk, i.e., x [i] = [x (i-1)B+1 , • • • , x iB ].Then we write R n as a recurrent function and compute the retention output of the i-th chunk via:R i =K ⊺ [i] (V [i] ⊙ ζ) + γ B R i-1 ζ ij = γ B-i , ξ ij = γ i Retention(X [i] ) = (Q [i] K ⊺ [i] ⊙ D)V [i] Inner-Chunk + (Q [i] ⊙ ξ)R i-1Cross-Chunk (11) Finally, we show that the chunkwise recurrent representation is equivalent to the other representations.As shown in D Hyperparameters Used in Section 3.1We use LLaMA [48] architecture, including RMSNorm [59] and SwiGLU [40, 7] module, as the Transformer backbone, which shows better performance and stability. The weights of word embedding and softmax projection are shared. Consequently, other variants follow these settings.For RetNet, the FFN intermediate dimension is 5 3 d and the value dimensions in W G , W V , W O are also 5 3 d, where the overall parameters are still 12d 2 .For H3, we set the head dimension to 8. For RWKV, we use the TimeMix module to substitute self-attention layers while keeping FFN layers consistent with other models for fair comparisons.For Mamba, we follow all the details in the paper [19] , where double-SSM layers are implemented instead of "SSM + SwiGLU". In addition to RetNet and Mamba, the FFN intermediate dimension is all 8 3 d. All models have 400M parameters, 24 layers, and a hidden dimension of 1024. We train the models with 40k steps and a batch size of 0.25M tokens.Layers Q , W K ∈ R d×d , W G , W V ∈ R d×2d , W O ∈ R 2d×d . Notice that the head dimension of V is twice Q, K,Table 9 presents one-shot performance on two open-ended question-answering tasks, including SQUAD [39] and WebQS [5] , with 6.7B models as follows. We report the recall metric in the table , i.e., whether the answers are contained in the generated response. We compare with grouped-query attention [1] and evaluate the method in the context of RetNet.Grouped-query attention makes a trade-off between performance and efficiency, which has been successfully verified in LLaMA2 34B/70B [49] . The method reduces the overhead of key/value cache during inference. Moreover, the performance of grouped-query attention is better than multi-query attention [42] , overcoming the quality degradation brought by using one-head key value.As shown in Table 10 , we compare the inference cost with grouped-query attention and apply the method for RetNet. For the LLaMA2 70B model, the number of key/value heads is reduced by 8×, where the query head number is 64 while the key/value head number is 8. For RetNet-70B, the parameter allocation is identical to LLaMA [48] , where the dimension is 8192, and the head number is 32 for RetNet. For RetNet-70B-GQ2, the key-value head number is 16, where grouped-query retention is applied. We run the inference with four A100 GPUs without quantization.When the batch size is 256, LLaMA2 runs out of memory while RetNet without group query still has a high throughput. When equipped with grouped-query retention, RetNet-70B achieves 38% acceleration and saves 30% memory.We evaluate LLaMA2 under 2k and 8k lengths separately. The batch size is reduced to 8 so that LLaMA2 can run without out of memory. Table 10 shows that the inference cost of Transformers increases with the sequence length. In contrast, RetNet is length-invariant. Moreover, RetNet-70B-GQ2 achieves better latency, throughput, and GPU memory than LLaMA2-70B-2k/8k equipped NeurIPS Paper Checklist Justification: The abstract and introduction is carefully written.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper.• The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.• The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.• It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 


Decision: The paper proposes the Retentive Network, which is a kind of linear attention network. All reviewers agree that the idea (if novel) is interesting. 

However, Reviewer maEf pointed out that the proposed approach highly resembles an EMNLP'22 paper [1]. In the author response, the author emphasized different roots of their design and some minor difference. However, this does not seem to warrant a main conference paper. Also, there lacks empirical comparison between the two models. 

The other two reviewers give acceptance-side scores (borderline acc & weak acc), but their original reviews did not factor in the EMNLP'22 paper. In the reviewer-AC discussion, the reviewers read the related paper and acknowledged the resemblance. 

The authors should honestly discuss and compare with [1] in their next version, and more importantly, reconsider the significance of their contributions.

[1] Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 7025–7041. 2022.

Reviews: review: paper_summary: paper_summary: This paper presents Retentive Network (RetNet), a family of efficient models that incorporate exponential decay within a linear attention-like structure. RetNet shares similarities with state-space models and linearized attention, enabling both training parallelism and O(1) inference cost. Additionally, RetNet supports chunk-wise parallel computation for efficient long-sequence training. Experimental results demonstrate RetNet achieves performance comparable to Transformers and outperforms other efficient variants on language modeling and vision tasks.
strength_weakness: strength_weakness: - Some experiments could be improved
- Some claims may be misleading
- RetNet's performance lags behind Transformers at smaller model scales, suggesting it might be more demanding in terms of capacity and compute resources for optimal performance. This trade-off should be carefully considered and analyzed.
questions: questions: 1. The inference results in Figure 5 start with 2048. What’s the inference speed for shorter sequences?
2. The claim that "None of the previous work can achieve strong performance and efficient inference at the same time compared with Transformers" is overly strong and potentially misleading. Recent advancements in efficient modeling, such as Mamba, have demonstrated better scaling properties than Transformers.
3. It would be great to include the training loss curve for both Transformer and RetNet.
limitations: limitations: I didn't see serious problems.
paper_summary: paper_summary: The authors propose a linear attention model called RetNet for language modeling, which has a linear training complexity and constant inference complexity.
strength_weakness: strength_weakness: 1. The authors introduce a new term called "Retention," but this is essentially the same as Linear Attention without the denominator, which has already been proposed in [1].
2. Lack of comparison with the baselines on open source pretraining data. All the training experiments are conducted on in-house data mixtures, which harms the reproducibility.
3. The paper doesn't compare RetNet with other linear attention model (such as GLA, RWKV, Mamba) on downstream tasks with standard metrics instead of perplexity. Table 2 only include RetNet and Transformer. The efficiency measurment of RetNet+ is absent.
4. The evaluation on MMLU/Qasper is using perplexity but not the widely-used accuracy/F1 metric. The perplexity results don't necessarily mean that the model can make correct choices for the samples in MMLU, and has less guidance for the model's downstream performance.
5. Missing citations: The authors should also cite [1] for the normalization after retention, and discuss the details of the triton implementation of RetNet and its difference from the implementation in the Flash Linear Attention [2] library.

[1] Zhen Qin, Xiaodong Han, Weixuan Sun, Dongxu Li, Lingpeng Kong, Nick Barnes, and Yiran Zhong. The devil in linear transformer. In Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 7025–7041, Abu Dhabi, United Arab Emirates, Dec. 2022. Association for Computational Linguistics.

[2] Yang, Songlin and Zhang, Yu. FLA: A Triton-Based Library for Hardware-Efficient Implementations of Linear Attention Mechanism. https://github.com/sustcsonglin/flash-linear-attention
questions: questions: Can you also add the results on MMLU and Qasper with the standard metrics besides perplexity?
limitations: limitations: No, the authors should have a limitation section to point out the strong assumptions of their approximation of self-attention and relative position embedding.
paper_summary: paper_summary: The paper proposes the Retentive Network (RetNet) as a foundation architecture for large language models. RetNet has a multi-scale retention mechanism with three computation paradigms: parallel, recurrent, and chunkwise recurrent. 
The retention mechanism starts with a recurrent modeling formulation and derives a parallel formulation. It maps input vectors to state vectors recurrently and implements a linear transform to encode sequence information. Then, it makes the projection content-aware by using learnable matrices. The retention layer is defined using these matrices and a complex position embedding, combining causal masking and exponential decay along relative distance. 
It achieves low-cost inference, efficient long-sequence modeling, comparable performance to Transformers, and parallel training. Experimental results show its superiority in language modeling, inference cost, and training throughput.
strength_weakness: strength_weakness: 1. The paper presents the scaling curves of RetNet and Transformer with model sizes ranging from 1.3B to 6.7B, concluding that RetNet is favorable in terms of size scaling and starts to outperform Transformer when the model size is larger than 2B. However, it does not provide a detailed explanation for this trend. Understanding the underlying reasons for this performance difference with increasing model size could provide more insights into the effectiveness of RetNet and its potential advantages over Transformer.
2. The use of $\gamma$ in the RetNet may appear somewhat heuristic. The paper assigns different $\gamma$ for each head in the multi-scale retention (MSR) module and keeps them fixed among different layers. While this approach is used to achieve certain effects, such as enhancing the non-linearity of the retention layers and improving the model's performance, the specific rationale for choosing these values and the potential impact on the model's behavior could be further explained.
questions: questions: 1. Can you provide more insights into why RetNet starts to outperform Transformer when the model size is larger than 2B? What specific characteristics or mechanisms of RetNet contribute to this improved performance at larger scales? Have you tried different learning rates to investigate their impact on the scaling of the RetNet model? If so, what were the results and how did they affect the model's performance and scalability? 
2. A more detailed discussion on the selection of $\gamma$ and its effects on the model's performance, as well as how it compares to other possible approaches or values, would be beneficial in providing a deeper understanding of the RetNet's functionality. Additionally, exploring the sensitivity of the model to different values of $\gamma$ or conducting experiments to justify the choice could strengthen the argument for using this particular heuristic.
limitations: limitations: none

Comments: comment: Thank you for your reply. I decided to maintain my original score.
Dear reviewers,

The discussion period will end soon. If you haven't responded to the authors' rebuttal, please do so and kick off the discussion. 

Best, 
SACa
Dear Reviewer,

Thank you for taking the time to review our submission.

We have carefully considered your feedback and submitted our responses and clarifications during the rebuttal phase. We highly value your insights and would appreciate any further discussion or questions you might have regarding our responses.

We understand that you have a busy schedule, but your additional feedback would be incredibly helpful for us to improve our work. If you could take a moment to review our rebuttal and provide any further comments, we would be very grateful.

Thank you again for your time and consideration.

Best regards