Retentive Network: A Successor to Transformer for Large Language Models

Q8: Table 2 lacks comparison with open-source models.

A8: We use Transformers with RoPE relative positions, which is a relatively strong baseline. Compared with LLaMA, there are only two differences: RMSNorm and SwiGLU. Specifically, RMSNorm typically improves stability instead of performance. Both RMSNorm and SwiGLU are orthogonal to our work. As long as the comparison settings are rigor and fair, the conclusions are solid. I totally understand you would like to compare different checkpoints, although they usually used different training data, data preprocessing pipelines, architecture configurations, number of training tokens, and hyperparameters. The paper focuses on fair comparisons rather than chasing for benchmarks.

Q9: The evaluation scope is too limited.

A9: The language modeling perplexity correlates very well with different tasks. Moreover, the language models have unified all tasks as generation. We report language modeling perplexity and various downstream task performance. We didn’t evaluate translation because our training corpus contains English data instead of multilingual data. We additionally evaluate one-shot performance on two open-ended question answering tasks with 7B models as follows. Notice that we report the recall metric in the table, i.e., whether the answers are contained in the generated response.

Dataset / Transformer / RetNet

Squad / 67.7 / 72.7

WebQS / 36.4 / 40.4

Q10: the evaluation datasets in Tables 4 and 5 are inconsistent with Table 3

A10: We divide our evaluation into two groups: large scale comparison with Transformers and small size with other baselines. Evaluating perplexity is a stable and predictable metric for language modeling. We mainly report end-task accuracy numbers for larger-size models while reporting perplexity for small-size models. As indicated in previous work ([2304.15004] Are Emergent Abilities of Large Language Models a Mirage?), perplexity is smoother and more robust than accuracy for small models.

Q11: There is a lack of ablation analysis on the architectural design.

A11: The dimension modification of $W_v$ and FFNs aim to align the same parameter count with Transformers. Besides, increasing head dimension is harmful for inference efficiency. So the current setting is more scientific for evaluation.

Q1: “Impossible Triangle” is an absolute overclaim because RWKV and H3 have already demonstrated models are comparable to Transformers

A1: The claim is fair enough. The “comparable performance” means that the models achieve similar results under the same setting (e.g., #parameters, and training corpus). For example, previous comparisons use Transformers with absolute position while the compared methods benefit from relative position modeling. Moreover, in H3 paper, the comparable results are in hybrid settings (i.e., combine H3 and Transformer layers), but we don’t add any Transformer layers. We conducted various controlled experiments (with matched #parameters and using the same training corpus) to compare different architectures. We are confident that the claim holds well. The experiments in Table 4 also show that previous methods still have a big gap.

Q2: RWKV can indeed be computed in parallel.

A2: We give a clear definition on “training parallelization” in the caption of Table 1, which is discussed from the sequential perspective. “∗”: whether the training implementation is sequentially parallelized, although RWKV uses channel-wise parallelism. As stated in A1, RWKV’s performance is actually not comparable with Transformers according to our experiments (i.e., same #parameters, same data, and with relative position modelings). So, the statement of RWKV in Table 1 is fair enough.

Q3: The form of Equation 1 is similar to RFA-GATE

A3: Eq. (1) and RFA-GATE are not the same. Linear Attention is a candidate for RFA-Gate, where they both have denominators. In contrast, Eq. (1) does not include $A$. We can add the citation as suggested.

Q4: The form of GroupNorm in Equation 8 is consistent with the NormAttention.

A4: GroupNorm is significantly better than LayerNorm for multi-scale retention, as different heads have different statistic information. The implementation is different from NormAttention here. We can discuss the difference in the paper.

Q5: There is a mistake in the description of the Linear Attention section in Section 2.4.

A5: “approximating softmax” means that previous methods also encode a “probability distribution” on different positions, which is valid in [4] and [6]. Moreover, “Linear Attention” is not “Linearized Attention”, it stands for the paper “Transformers are RNNs: Fast autoregressive transformers with linear attention”.

Q6: The statement "However, linear attention struggles to effectively encode position information, rendering the models less performant" should be supported by a reference since there could be various reasons for the poor performance.

A6: Previous methods use kernel methods on query and key, which brings difficulty on relative position such as RoPE, where RoPE produces negative values for denominators. Otherwise the attention scores tend to become too flatten. [5] uses block softmax attention as an ad-hoc solution to make attention local.

Q7: The discussion about MEGA and MEGA-chunk is missing, as they are similar to RetNet in utilizing the EMA technique.

A7: RetNet has no direct connection with EMA technique. Instead, EMA is a weaker version of S4, so it’s convincing to compare with H3, which is a stronger baseline.

Q1: The paper is not clear to the reviewer why these two forms in Figure 2 (a) and Figure 2 (b) are equivalent.

A1: The equivalency between parallel and recurrent form is discussed mathematically from Eq. (1) to Eq. (6), where Figure 2a is for Eq. (5), and Figure 2b is for Eq. (6).

Q2: Are there any theoretical proof that retention is more capable than full attention?

A2: The comparable capabilities are evaluated empirically in the submission. For retention, we allow negative values for “retention scores”. In comparison, full attention only allows positive values for attention scores. If we consider finite arithmetic precision in practice, retention potentially utilizes the capacity more sufficiently without wasting the negative numerical ranges. The claim “Full attention should be more capable than recurrent networks” is not trivial because attention is often worse than RNN or S4 under many tasks, such as learning formal language.

Q3: two model architectures should not cross over when scaling the model up in log scale.

A3: Here we are comparing two different architectures in one figure. The curves do not cross over for the same architecture. But it doesn’t hold anymore if we compare different methods. If the learning curve is $f(m)=\alpha m^{\beta_g}+\gamma$, there are three variants $\alpha, \beta, \gamma$. When $f_1(m) = f_2(m)$ and $\beta_{g1} > \beta_{g2}$, it’s possible that two lines are cross over. Our scaling curves are not contradictory with early papers.

Q4: Do you have results on generative tasks like translation and question answering?

A4: The language modeling perplexity correlates very well with different tasks. Moreover, the language models have unified all tasks as generation. We report language modeling perplexity and various downstream task performance. We didn’t evaluate translation because our training corpus contains English data instead of multilingual data. We additionally evaluate one-shot performance on two open-ended question answering tasks with 7B models as follows. Notice that we report the recall metric in the table, i.e., whether the answers are contained in the generated response.

Dataset / Transformer / RetNet

Squad / 67.7 / 72.7

WebQS / 36.4 / 40.4

Q1: The inference cost is not constant

A1: Eq. (1) is discussed under one-head settings. We still use multi-head in Eq. (8), and the shape of $s_n$ is $h\times d$, where $h$ is query and key's head dimension. So the inference cost will be $O(1)$. Our inference cost experiment in Figure 4 demonstrates that.

Q2: why RetNet is faster and using less-memory than Flash Attention?

A2: For the comparison with FlashAttention, we use an optimized FlashAttention, which is better than FlashAttention1 and is comparable with FlashAttention2 for long sequences. Since the publication time of FlashAttention2 is later than ICLR submission deadline, we didn't compare with the official one in the submission.

Q3: the total number of training tokens are only 100B. It’s not enough since Transformer is data-eager.

A3: According to the Chinchilla scaling law, 7B model’s converge speed slows down at $FLOPs=10^{21}$. The flops for 100B tokens are more than $(7\times 10^ 9) \times (100 \times 10^ 9) \times 6 > 4\times 10^{21}$. So 100B tokens are enough for evaluating model performance.

Q4: the improved version of Transformer in Llama (with RMSNorm and SwiGLU) has achieved significantly better results than the standard Transformer.

A4: The LLaMA improvement is on training stability and FFN capacity, which is useful regardless of attention or other token-mixing methods. RMSNorm and SwiGLU can also be applied to RetNet. The comparisons under standard Transformer are fair.

Q5: there are no direct comparison of RetNet with other models in large-scale setting.

A5: The conclusions of comparing with other architectures are consistent even with larger-scale training. Scaling up every method went beyond our computation budget. As the de facto architecture is Transformer now, we spent more resources on the comparisons with Transformers, where we ran the experiments in large-scale settings.

Q6: Missing citations.

A6: We can add the suggested paper “[1] Hua et al., Transformer Quality in Linear Time” in the revision.

1. "this is essentially a transformer without the softmax and an added time decay." However, doing that in a naive way can't get a well-performed model. Previous works also made a similar contribution, but their performance still has a big gap in Table 4. We utilize new designs to achieve comparable performance, including gated multi-scale retention and group normalization.
2. Figure 2a and 2b correspond to Equation 5 and 6. Mathematical explanation may help.
3. RetNet's first shortcoming is that it is not strictly comparable under small. size, e.g. 1.3B. Besides, we work on GPT settings so it can't migrate to bidirectional transformers. Regular tasks in NLP almost follow the pre-training performance.
4. The code will be open-source definitely after review.