Paper: Retentive Network: A Successor to Transformer for Large Language Models (ICLR-2024)

Reviewer comment (Experiments):
The main weakness with this paper is overclaiming and lack of citations, which can be misleading for readers. For example, the claim in Figure 1 that "RetNet makes the 'impossible triangle' possible" is an absolute overclaim because the paper lacks validation with larger models and comparison with open-source Transformer models.

Author response:
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.

Reviewer comment (Evaluation):
The authors claim that RWKV and Linear Attention perform poorly, but according to [1], [2], their performance can be on par with Transformers.

Author response:
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.

Reviewer comment (Novelty):
In Section 2, 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 [2], [3].

Author response:
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.

Reviewer comment (Novelty):
Figure 1 is an absolute overclaim because RWKV[1] and H3[2] have already demonstrated models at the billion-scale level that can achieve performance comparable to Transformers, with parallel training and constant inference. The suggestion is to remove this figure as it could mislead readers.

Author response:
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.

Reviewer comment (Presentation):
The description of RWKV in Table 1 is completely wrong. According to RWKV[1] and the description in [2], RWKV can indeed be computed in parallel. Additionally, according to the RWKV paper, its performance is comparable to Transformers, so describing its performance as ✔ is also inaccurate. Overall, Table 1 is highly misleading and can affect the authors' judgment of model performance. The suggestion is to reorganize this table accordingly.

Author response:
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.

Reviewer comment (Novelty):
The form of Equation 1 is similar to RFA-GATE presented in [7], but the authors did not cite these articles throughout the paper.

Author response:
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.

Reviewer comment (Novelty):
The form of GroupNorm in Equation 8 is consistent with the NormAttention proposed in [5], but the author does not cite it at all.

Author response:
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.

Reviewer comment (Theory):
There is a mistake in the description of the Linear Attention section in Section 2.4. Firstly, Linear Attention refers to the use of the Right-product trick to reduce complexity, but it does not necessarily imply approximation. For example, [4], [5], and [6] do not involve an approximation approach like softmax.

Author response:
“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”.

Reviewer comment (Experiments):
Table 2 lacks comparison with open-source models. Firstly, the configuration of the Transformer is not mentioned, whether it is based on GPT2 architecture or Llama architecture. Additionally, there is no information provided regarding the parameter count or training data. There is also no comparison with open-source models such as Bloom, Pythia, GPT-Neo, or RWKV. Comparing with these open-source models would allow readers to better understand the performance level of the proposed model.

Author response:
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.

Reviewer comment (Experiments):
There is a lack of ablation analysis on the architectural design, such as why the dimension of $W_v$ is chosen as d * 2d. Additionally, an ablation for adding head dimension should be included in Table 5.

Author response:
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.

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

Author response:
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).

Reviewer comment (Theory):
Are there any theoretical proof that retention is more capable than full attention?

Author response:
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.

Reviewer comment (Evaluation):
The paper is not clear why in Figure 3, RetNet is more effective in the large model regime. According to some prior work [1][2], two model architectures should not cross over when scaling the model up in log scale.

Author response:
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.

Reviewer comment (Experiments):
Results are only provided on classification and summarization tasks, not generative tasks like translation and question answering. This shows limited generalization of the model architecture.

Author response:
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

Reviewer comment (Theory):
Why would RetNet be better than full attention based transformer in terms of quality? Full attention should be more capable than recurrent networks.

Author response:
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.

Reviewer comment (Experiments):
Do you have results on NLG tasks? Refer to GPT3 paper and report some numbers on generative tasks.

Author response:
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

Reviewer comment (Evaluation):
The model scaling curve in Figure 3 looks odd to me. Why would the two lines cross over? It does not align with other empirical findings in early papers.

Author response:
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.

Reviewer comment (Theory):
Table 1 is mis-leading. If I understand correctly, the recurrent state $s_n$ in Eq (1) is in the shape of $d\times d$, where $d$ is the model dimension and is very large in practice (sometimes even larger than $N$). In S4 or other SSMs, the shape of the recurrent hidden state is $h\times d$ with relatively small $h$, e.g. $h=32$. However, in Table 1 the authors claimed the inference cost of RetNet is $O(1)$.

Author response:
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.

Reviewer comment (Experiments):
Table 3 is unclear. If I understand correctly, the parallel representation in Eq(5) is used for model training. However, the complexity of Eq(5) is almost the same with standard attention mechanism, with only difference that it does not need to compute softmax. Then why RetNet is faster and using less-memory than Flash Attention? Flash Attention also leveraged the block-wise computation of the attention matrix, which is similar to the chunkwise representation of RetNet. In addition, which version of Flash Attention was used in these experiments? Flash Attention v2 has been significantly optimized.

Author response:
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.

Reviewer comment (Experiments):
The experimental setting is unconvincing. Though the authors scaled-up RetNet to 13B, the total number of training tokens are only 100B. It is well-known that Transformer is data-eager, and it is unfair to compare with Transformer with relatively small training data. Moreover, the improved version of Transformer in Llama (with RMSNorm and SwiGLU) has achieved significantly better results than the standard Transformer. But all the comparisons in this paper are with standard Transformer.

Author response:
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.

Reviewer comment (Experiments):
In Table 1, the authors claimed that only Transformer and RetNet are achieving "good performance". However, there are no direct comparison of RetNet with other models in large-scale setting. The results in Table 4 were conducted with a small 200M model size and only 5B tokens.

Author response:
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.