Title: Retentive Network
Venue: NeurIPS 2024
Accepted: no
Keywords: Retentive Network

Model Architecture
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.

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 O (1) inference, which improves decoding throughput, latency, and GPU memory without sacrificing performance. The chunkwise recurrent representation facilitates efficient longsequence 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.

1 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 23 standard attention scores exp( q · k ) with kernels ϕ ( q ) · ϕ ( k ) , so that autoregressive inference can 24 be rewritten in a recurrent form. However, the modeling capability and performance are worse than 25 Transformers, which hinders the method's popularity. The second strand returns to recurrent models 26 for efficient inference while sacrificing training parallelism. As a remedy, element-wise operators [36] 27 are used for acceleration, however, representation capacity and performance are harmed. The third 28 line explores replacing attention with other mechanisms, such as S4 [20], and its variants [11, 38]. 29 None of the previous work can achieve strong performance and efficient inference at the same time 30 compared to Transformers. 31

In this work, we propose retentive networks (RetNet), achieving low-cost inference, efficient long32

sequence modeling, Transformer-comparable performance, and parallel model training simultane33

ously. Specifically, we introduce a multi-scale retention mechanism to substitute multi-head attention, 34 which has three computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent repre35 sentations. First, the parallel representation empowers training parallelism to utilize GPU devices 36 fully. Second, the recurrent representation enables efficient O (1) inference in terms of memory 37 and computation. The deployment cost and latency can be significantly reduced. Moreover, the 38 implementation is greatly simplified without key-value cache tricks. Third, the chunkwise recurrent 39

40

41

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 42 show that RetNet is consistently competitive in terms of both scaling curves and in-context learning. 43 Moreover, the inference cost of RetNet is length-invariant. For a 7B model and 8k sequence 44 length, RetNet decodes 8.4 × faster and saves 70% of memory than Transformers with key-value 45 caches. During training, RetNet also achieves 3 × acceleration than standard Transformer with 46 highly-optimized FlashAttention-2 [10]. Besides, RetNet's inference latency is insensitive to batch 47 size, allowing enormous throughput. The intriguing properties make RetNet a potential candidate to 48 replace Transformer for large language models. 49

50

51

52

53

54

55

56

57

58

59

60

61

62

2 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 |× d model , where d model is hidden dimension. Then we compute contextualized vector representations X l = RetNet l ( X l -1 ) , l ∈ [1 , L ] .

2.1 Retention

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:

where we map v n to the state vector s n , and then implement a linear transform to encode sequence 63 information recurrently. Next, we make the projection Q n , K n content-aware: 64

where W Q , W K ∈ R d × d are learnable matrices. 65

We diagonalize the matrix A = Λ( γe iθ )Λ -1 , where γ, θ ∈ R d . Then we obtain A n -m = 66 Λ( γe iθ ) n -m Λ -1 . By absorbing Λ into W Q and W K , we can rewrite Equation (1) as: 67

where Q n ( γe iθ ) n , K m ( γe iθ ) -m is known as xPos [45], i.e., a relative position embedding proposed 68 for Transformer. We further simplify γ as a scalar, Equation (3) becomes: 69

where † is the conjugate transpose. The formulation is easily parallelizable within training instances. 70

75

76

77

78

79

80

81

82

83

𝑋

𝐾

𝑄

𝑉

𝑂

GN

𝑄𝐾 ⊺ ⊙𝐷 𝑉

𝑋𝑛

𝑆𝑛-1

𝛾

𝐾𝑛

𝑉 𝑛

𝑄𝑛

𝑂𝑛

Recurrent

State

Input

Output

GN

Figure 1: RetNet has three equivalent computation paradigms, i.e., parallel, recurrent, and chunkwise recurrent representations. Given the same input, three paradigms obtain the same output. 'GN' is short for GroupNorm.

The Parallel Representation of Retention As shown in Figure 1a, the retention layer is defined as: 74

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:

where Q,K,V,γ are the same as in Equation (5).

The Chunkwise Recurrent Representation of Retention A hybrid form of parallel representation 84 and recurrent representation is available to accelerate training, especially for long sequences. We 85 divide the input sequences into chunks. Within each chunk, we follow the parallel representation 86 (Equation (5)) to conduct computation. In contrast, cross-chunk information is passed following the 87 recurrent representation (Equation (6)). Specifically, let B denote the chunk length. We compute the 88 retention output of the i -th chunk via: 89

where [ i ] indicates the i -th chunk, i.e., x [ i ] = [ x ( i -1) B +1 , · · · , x iB ] . The proof of the equivalence 90 between recurrent representation and chunkwise recurrent representation is described in Appendix B. 91

92

2.2 Gated Multi-Scale Retention

We use h = d model / d retention heads in each layer, where d is the head dimension. The heads use 93 different parameter matrices W Q , W K , W V ∈ R d × d . Moreover, m ultis cale r etention (MSR) assigns 94

different γ for each head. For simplicity, we set γ identical among different layers and keep them 95 fixed. In addition, we add a swish gate [23, 40] to increase the non-linearity of retention layers. 96 Formally, given input X , we define the layer as: 97

where W G , W O ∈ R d model × d model are learnable parameters, and GroupNorm [53] normalizes the 98 output of each head, following SubLN proposed in [43]. Notice that the heads use multiple γ scales, 99 which results in different variance statistics. So we normalize the head outputs separately. 100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

The pseudocode of retention is summarized in Figure 2.

Retention Score Normalization Weutilize 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 ˜ D nm = D nm / √ ∑ n i =1 D ni . Third, let R denote the retention scores R = QK ⊺ ⊙ D , we normalize it as ˜ R nm = R nm / max( ∑ n i =1 | R ni | , 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.

2.3 Overall Architecture of Retention Networks

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 embeddings X 0 = [ x 1 , · · · , x | x | ] ∈ R | x |× d model as the input and compute the model output X L :

where LN( · ) is LayerNorm [3]. The FFN part is computed as FFN( X ) = gelu( XW 1 ) W 2 , where 115 W 1 , W 2 are parameter matrices. 116

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

121

122

123

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.

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

3 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].

3.1 Comparison with Transformer Variants

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 12 d 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.

Fine-Grained Language Modeling Evaluation 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.

Knowledge-Intensive Tasks 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.

Table 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.

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

3.2 Language Modeling Evaluation with Various Model Sizes

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 AMDMI200 GPUs.

Figure 3 reports 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 in terms of size scaling. In addition to performance, 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.

3.3 Long-Context Evaluation

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.

Figure 4 reports the answer perplexity given various lengths of input document. It shows that both Transformer and RetNet+ perform better with longer input documents. The results indicate that the language models successfully utilize the long-distance context. Notice that the 12K and 16K results in Qasper are similar because the lengths of most documents are shorter than 16K. Moreover, RetNet+ obtains competitive results compared with Transformer for long-context modeling. Meanwhile, retention has better training and inference efficiency.

4096

8192

12288

16384

Context Length

1.5

1.6

1.7

Answer PPL

GovReport

4096

8192

12288

16384

Context Length

1.7

1.8

1.9

Answer PPL

Qasper

4096

8192

12288

16384

Context Length

3.25

3.50

3.75

4.00

4.25

Answer PPL

NarrativeQA

Transformer

3.4 Inference Cost

As shown in Figure 5, we compare memory cost, throughput, and latency of Transformer and RetNet 185 during inference. Transformers reuse KV caches of previously decoded tokens. RetNet uses the 186 recurrent representation as described in Equation (6). We evaluate the 6.7B model on the A100-80GB 187 GPU. Figure 5 shows that RetNet outperforms Transformer in terms of inference cost. 188

Memory As shown in Figure 5a, the memory cost of Transformer increases linearly due to KV 189 caches. In contrast, the memory consumption of RetNet remains consistent even for long sequences, 190

1.3B

2.7B

6.7B

6

8

10

12

14

16

18

Validation PPL

Transformer

RetNet

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

214

215

216

217

218

219

220

221

222

223

224

225

226

2048

3072

4096

5120

6144

7168

8192

Sequence Length

20

30

40

Model Weights

RetNet

Transformer

2048

3072

4096

5120

6144

7168

8192

Sequence Length

50

100

150

200

250

300

Throughput (wps)

RetNet

Transformer

2

4

6

8

100

200

300

Latency (ms)

Transformer (1024)

Transformer (2048)

Transformer (4096)

Transformer (8192)

RetNet (8192)

(a) GPU memory cost with varying sequence length.

(b) Inference throughput with varying sequence length.

(c) Inference latency with different batch sizes.

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%.

Throughput As presented in Figure 5b, 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 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.

3.5 Training Throughput

Figure 6 compares the training throughput of Transformer and RetNet, where the training sequence lengths range from 8192 to 65536. The model size is 3.5B, 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 train-

8192

16384

32768

65536

Sequence Length

4000

6000

8000

10000

12000

14000

Throughput (wps)

Transformer (FlashAttn-2)

RetNet (Triton)

RetNet (PyTorch)

ing 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.

3.6 Zero-Shot and Few-Shot Evaluation on Downstream Tasks

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.

3.7 Ablation Studies

We ablate various design choices of RetNet and report the language modeling results in Table 3. The evaluation settings and metrics are the same as in Section 3.1.

227

228

229

230

231

232

233

234

235

236

237

238

239

240

241

242

243

244

245

246

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.

3.8 Results on Vision Tasks

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:

where multi-scale decay is removed in bidirectional computation. Notice that we can compute 247 retention in different orders. Similar to linear attention [27], the Q ( K ⊺ V ) paradigm is an efficient 248 operator in bidirectional settings, especially for high-resolution images. 249

We perform experiments on ImageNet-1K classification [13], COCO object detection [32], and 250 ADE20K semantic segmentation [60]. We compare RetNet with DeiT [47] which is a well-tuned 251 vision Transformer. Besides, we follow [21] and plug in a depth-wise convolution in experiments. 252 We adopt the DeiT-M size, which has about 38M parameters. For ImageNet-1K image classification, 253

254

255

256

257

258

259

260

261

262

263

264

265

266

267

268

269

270

271

272

273

274

275

276

277

278

279

280

281

282

283

284

285

286

287

288

289

290

291

292

293

294

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.

4 Related Work

Numerous efforts are focused on reducing the quadratic complexity of attention mechanisms. Linear attention [27] uses various kernels ϕ ( q i ) ϕ ( k j ) / ∑ | x | n =1 ϕ ( q i ) ϕ ( k n ) 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. For training, RWKV's token-mixing complexity is O ( DN ) , and Mamba's complexity is O ( DHN ) with optimized CUDA kernels. Hyena's is O ( DN log N ) with Fast Fourier Transform acceleration. In comparison, the chunk-wise recurrent representation is O ( DN ( B + H )) , where B is the chunk size, and we usually set H = 256 , B ≤ 512 . However, chunk-wise computation is highly parallelized, enabling efficient hardware usage. For large model size (i.e., larger D ) or sequence length, the additional b + h has negligible effects. For inference, among the efficient architectures compared, Hyena has the same complexity (i.e., O ( N ) per step) as Transformer, while the others can perform O (1) decoding.

5 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.

295

296

297

298

299

300

301

302

303

304

305

306

307

308

309

310

311

312

313

References

J. Ainslie, J. Lee-Thorp, M. de Jong, Y. Zemlyanskiy, F. Lebrón, and S. Sanghai. GQA: Training generalized multi-query Transformer models from multi-head checkpoints. arXiv preprint arXiv:2305.13245 , 2023.

S. Arora, S. Eyuboglu, A. Timalsina, I. Johnson, M. Poli, J. Zou, A. Rudra, and C. Ré. Zoology: Measuring and improving recall in efficient language models. arXiv preprint arXiv:2312.04927 , 2023.

J. L. Ba, J. R. Kiros, and G. E. Hinton. Layer normalization. arXiv preprint arXiv:1607.06450 , 2016.

M. Beck, K. Pöppel, M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brandstetter, and S. Hochreiter. xLSTM: Extended long short-term memory. arXiv preprint arXiv:2405.04517 , 2024.

J. Berant, A. Chou, R. Frostig, and P. Liang. Semantic parsing on Freebase from questionanswer pairs. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing , pages 1533-1544, Seattle, Washington, USA, Oct. 2013. Association for Computational Linguistics.

Y. Bisk, R. Zellers, R. L. Bras, J. Gao, and Y. Choi. Piqa: Reasoning about physical commonsense in natural language. In Thirty-Fourth AAAI Conference on Artificial Intelligence , 2020.

A. Chowdhery, S. Narang, J. Devlin, M. Bosma, G. Mishra, A. Roberts, P. Barham, H. W. 314 Chung, C. Sutton, S. Gehrmann, P. Schuh, K. Shi, S. Tsvyashchenko, J. Maynez, A. B. Rao, 315

Barnes, Y. Tay, N. M. Shazeer, V. Prabhakaran, E. Reif, N. Du, B. C. Hutchinson, R. Pope, 316

Bradbury, J. Austin, M. Isard, G. Gur-Ari, P. Yin, T. Duke, A. Levskaya, S. Ghemawat, 317

Dev, H. Michalewski, X. García, V. Misra, K. Robinson, L. Fedus, D. Zhou, D. Ippolito, 318

Luan, H. Lim, B. Zoph, A. Spiridonov, R. Sepassi, D. Dohan, S. Agrawal, M. Omernick, 319

M. Dai, T. S. Pillai, M. Pellat, A. Lewkowycz, E. O. Moreira, R. Child, O. Polozov, K. Lee, 320

Zhou, X. Wang, B. Saeta, M. Díaz, O. Firat, M. Catasta, J. Wei, K. S. Meier-Hellstern, 321

Eck, J. Dean, S. Petrov, and N. Fiedel. PaLM: Scaling language modeling with pathways. 322 ArXiv , abs/2204.02311, 2022. 323

324

325

326

327

C. Clark, K. Lee, M.-W. Chang, T. Kwiatkowski, M. Collins, and K. Toutanova. BoolQ: Exploring the surprising difficulty of natural yes/no questions. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics , pages 2924-2936, 2019.

328

329

330

331

332

333

334

335

336

337

338

339

340

T. Computer. Redpajama-data: An open source recipe to reproduce llama training dataset, 2023. URL https://github.com/togethercomputer/RedPajama-Data .

T. Dao. FlashAttention-2: Faster attention with better parallelism and work partitioning. arXiv preprint arXiv:2307.08691 , 2023.

T. Dao, D. Y. Fu, K. K. Saab, A. W. Thomas, A. Rudra, and C. Ré. Hungry hungry hippos: Towards language modeling with state space models. arXiv preprint arXiv:2212.14052 , 2022.

S. De, S. L. Smith, A. Fernando, A. Botev, G. Cristian-Muraru, A. Gu, R. Haroun, L. Berrada, Y. Chen, S. Srinivasan, G. Desjardins, A. Doucet, D. Budden, Y. W. Teh, R. Pascanu, N. D. Freitas, and C. Gulcehre. Griffin: Mixing gated linear recurrences with local attention for efficient language models. 2024.

J. Deng, W. Dong, R. Socher, L.-J. Li, K. Li, and L. Fei-Fei. Imagenet: A large-scale hierarchical image database. In 2009 IEEE conference on computer vision and pattern recognition , pages 248-255. Ieee, 2009.

J. Dodge, A. Marasovi´ c, G. Ilharco, D. Groeneveld, M. Mitchell, and M. Gardner. Documenting 341 large webtext corpora: A case study on the colossal clean crawled corpus. In Conference on 342 Empirical Methods in Natural Language Processing , 2021. 343

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, 344 M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for 345 image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020. 346

347

348

349

L. Gao, S. Biderman, S. Black, L. Golding, T. Hoppe, C. Foster, J. Phang, H. He, A. Thite, N. Nabeshima, et al. The Pile: An 800GB dataset of diverse text for language modeling. arXiv preprint arXiv:2101.00027 , 2020.

350

351

352

353

354

355

356

357

358

359

360

361

362

363

364

365

366

367

368

369

370

371

372

373

374

375

376

377

378

379

380

381

382

383

384

385

L. Gao, J. Tow, B. Abbasi, S. Biderman, S. Black, A. DiPofi, C. Foster, L. Golding, J. Hsu, A. Le Noac'h, H. Li, K. McDonell, N. Muennighoff, C. Ociepa, J. Phang, L. Reynolds, H. Schoelkopf, A. Skowron, L. Sutawika, E. Tang, A. Thite, B. Wang, K. Wang, and A. Zou. A framework for few-shot language model evaluation, 12 2023. URL https://zenodo.org/ records/10256836 .

X. Geng and H. Liu. Openllama: An open reproduction of llama, May 2023. URL https: //github.com/openlm-research/open_llama .

A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 , 2023.

A. Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 , 2021.

D. Han, X. Pan, Y. Han, S. Song, and G. Huang. Flatten Transformer: Vision Transformer using focused linear attention. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 5961-5971, 2023.

K. He, G. Gkioxari, P. Dollár, and R. Girshick. Mask r-cnn. In Proceedings of the IEEE international conference on computer vision , pages 2961-2969, 2017.

D. Hendrycks and K. Gimpel. Gaussian error linear units (GELUs). arXiv: Learning , 2016.

D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt. Measuring massive multitask language understanding. arXiv preprint arXiv:2009.03300 , 2020.

S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural Computation , 9:1735-1780, Nov. 1997.

W. Hua, Z. Dai, H. Liu, and Q. Le. Transformer quality in linear time. In International Conference on Machine Learning , pages 9099-9117. PMLR, 2022.

A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention. In International Conference on Machine Learning , pages 5156-5165. PMLR, 2020.

T. Katsch. Gateloop: Fully data-controlled linear recurrence for sequence modeling. arXiv preprint arXiv:2311.01927 , 2023.

D. Kocetkov, R. Li, L. Ben Allal, J. Li, C. Mou, C. Muñoz Ferrandis, Y. Jernite, M. Mitchell, S. Hughes, T. Wolf, D. Bahdanau, L. von Werra, and H. de Vries. The Stack: 3TB of permissively licensed source code. Preprint , 2022.

H. Levesque, E. Davis, and L. Morgenstern. The winograd schema challenge. In Thirteenth International Conference on the Principles of Knowledge Representation and Reasoning , 2012.

O. Lieber, B. Lenz, H. Bata, G. Cohen, J. Osin, I. Dalmedigos, E. Safahi, S. Meirom, Y . Belinkov, S. Shalev-Shwartz, et al. Jamba: A hybrid Transformer-Mamba language model. arXiv preprint arXiv:2403.19887 , 2024.

T.-Y. Lin, M. Maire, S. Belongie, J. Hays, P. Perona, D. Ramanan, P. Dollár, and C. L. Zitnick. 386 Microsoft COCO: Common objects in context. In Computer Vision-ECCV 2014: 13th European 387 Conference, Zurich, Switzerland, September 6-12, 2014, Proceedings, Part V 13 , pages 740-755. 388 Springer, 2014. 389

I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference 390 on Learning Representations , 2019. 391

392

393

394

N. Mostafazadeh, M. Roth, A. Louis, N. Chambers, and J. Allen. Lsdsem 2017 shared task: The story cloze test. In Proceedings of the 2nd Workshop on Linking Models of Lexical, Sentential and Discourse-level Semantics , pages 46-51, 2017.

395

396

397

398

399

400

401

402

403

404

405

406

407

408

409

410

411

412

413

414

415

416

417

418

419

420

421

422

423

424

425

426

427

428

429

430

431

432

433

434

435

436

437

A. Orvieto, S. L. Smith, A. Gu, A. Fernando, C. Gulcehre, R. Pascanu, and S. De. Resurrecting recurrent neural networks for long sequences. ArXiv , abs/2303.06349, 2023.

B. Peng, E. Alcaide, Q. G. Anthony, A. Albalak, S. Arcadinho, H. Cao, X. Cheng, M. Chung, M. Grella, G. Kranthikiran, X. He, H. Hou, et al. RWKV: Reinventing RNNs for the Transformer era. ArXiv , abs/2305.13048, 2023.

H. Peng, N. Pappas, D. Yogatama, R. Schwartz, N. A. Smith, and L. Kong. Random feature attention. arXiv preprint arXiv:2103.02143 , 2021.

M. Poli, S. Massaroli, E. Nguyen, D. Y. Fu, T. Dao, S. Baccus, Y. Bengio, S. Ermon, and C. Ré. Hyena hierarchy: Towards larger convolutional language models. arXiv preprint arXiv:2302.10866 , 2023.

P. Rajpurkar, J. Zhang, K. Lopyrev, and P. Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing , pages 2383-2392, Austin, Texas, Nov. 2016. Association for Computational Linguistics. doi: 10.18653/v1/D16-1264.

P. Ramachandran, B. Zoph, and Q. V. Le. Swish: a self-gated activation function. arXiv: Neural and Evolutionary Computing , 2017.

U. Shaham, M. Ivgi, A. Efrat, J. Berant, and O. Levy. ZeroSCROLLS: A zero-shot benchmark for long text understanding. arXiv preprint arXiv:2305.14196 , 2023.

N. M. Shazeer. Fast Transformer decoding: One write-head is all you need. ArXiv , abs/1911.02150, 2019.

M. Shoeybi, M. Patwary, R. Puri, P. LeGresley, J. Casper, and B. Catanzaro. Megatron-LM: Training multi-billion parameter language models using model parallelism. arXiv preprint arXiv:1909.08053 , 2019.

J. Su, Y. Lu, S. Pan, B. Wen, and Y. Liu. Roformer: Enhanced transformer with rotary position embedding. arXiv preprint arXiv:2104.09864 , 2021.

Y. Sun, L. Dong, B. Patra, S. Ma, S. Huang, A. Benhaim, V. Chaudhary, X. Song, and F. Wei. A length-extrapolatable transformer. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 14590-14604, Toronto, Canada, July 2023. Association for Computational Linguistics.

P. Tillet and D. Cox. Triton: An intermediate language and compiler for tiled neural network computations. In Proceedings of the 3rd ACM SIGPLAN International Workshop on Machine Learning and Programming Languages , pages 10-19, 2019.

H. Touvron, M. Cord, M. Douze, F. Massa, A. Sablayrolles, and H. Jégou. Training data-efficient image transformers & distillation through attention. In International conference on machine learning , pages 10347-10357. PMLR, 2021.

H. Touvron, T. Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, et al. LLaMA: Open and efficient foundation language models. arXiv preprint arXiv:2302.13971 , 2023.

H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 , 2023.

J. Tow, M. Bellagente, D. Mahan, and C. Riquelme. StableLM 3B 4E1T. https://aka.ms/ StableLM-3B-4E1T , 2023.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and 438 I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 439 30: Annual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, 440 Long Beach, CA, USA , pages 6000-6010, 2017. 441

442

443

444

A. Wang, Y. Pruksachatkun, N. Nangia, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. SuperGLUE: A stickier benchmark for general-purpose language understanding systems. arXiv preprint arXiv:1905.00537 , 2019.

445

446

447

448

449

450

451

452

453

454

455

456

457

458

459

460

Y. Wu and K. He. Group normalization. In Proceedings of the European conference on computer vision (ECCV) , pages 3-19, 2018.

T. Xiao, Y. Liu, B. Zhou, Y . Jiang, and J. Sun. Unified perceptual parsing for scene understanding. In Proceedings of the European conference on computer vision (ECCV) , pages 418-434, 2018.

W. Xiong, J. Liu, I. Molybog, H. Zhang, P. Bhargava, R. Hou, L. Martin, R. Rungta, K. A. Sankararaman, B. Oguz, et al. Effective long-context scaling of foundation models. arXiv preprint arXiv:2309.16039 , 2023.

S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware-efficient training. arXiv preprint arXiv:2312.06635 , 2023.

R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. Hellaswag: Can a machine really finish your sentence? In Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics , 2019.

S. Zhai, W. Talbott, N. Srivastava, C. Huang, H. Goh, R. Zhang, and J. Susskind. An attention free transformer. arXiv preprint arXiv:2105.14103 , 2021.

B. Zhang and R. Sennrich. Root mean square layer normalization. Advances in Neural Information Processing Systems , 32, 2019.

B. Zhou, H. Zhao, X. Puig, T. Xiao, S. Fidler, A. Barriuso, and A. Torralba. Semantic 461 understanding of scenes through the ADE20k dataset. International Journal of Computer Vision , 462 127:302-321, 2019. 463

464

465

466

467

468

469

470

471

472

473

474

475

476

477

478

479

480

481

482

A Scaling Up Number of Training Tokens

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 TransformerRepro) for apple-to-apple comparison.

Our model RetNet+ follows the same configuration as in Section 3.3, which is a hybrid model. The model's hidden size is 3072, and the number of layers is 28. Without vocabulary embedding, the total number of parameters is 3.17B, which is between StableLM-3B-4E1T (2.7B) and OpenLLaMA-3Bv1 (3.19B). The batch size is 4M tokens. The training length is 4k. The learning rate is 3 . 2 × 10 -4 with 1000 warm-up steps and linear learning rate decay. The training corpus includes The Pile [16] and RedPajama [9]. TransformerRepro follows the exact same setting.

Table 5 reports accuracy numbers on the Harness-Eval benchmark [17]. We directly follow the evaluation protocol. The results show that RetNet+ achieves a performance comparable to TransformerRepro on language tasks. Notice that OpenLLaMA-3B-v1 and StableLM-3B use different learning rate schedules. The results of these two models are used for reference purposes.

B Equivalence Between Chunk-wise Recurrent Representation and Recurrent Representation

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:

where Γ i = γ n -i , ζ ij = γ B -m , and [ i ] indicates the i -th chunk, i.e., x [ i ] = [ x ( i -1) B +1 , · · · , x iB ] . 483 Then we write R n as a recurrent function and compute the retention output of the i -th chunk via: 484

485

486

487

488

489

490

491

492

493

494

495

496

497

498

499

500

501

502

503

504

505

506

507

508

509

510

Finally, we show that the chunkwise recurrent representation is equivalent to the other representations.

C Results with Different Context Lengths

As shown in Table 6, we report the results of language modeling with different context lengths. In order to make the numbers comparable, we use 2048 text chunks as evaluation data and only compute the perplexity for the last 128 tokens. Experimental results show that RetNet performs comparably with Transformer in different context lengths.

Table 6: Language modeling perplexity of RetNet and Transformer with different context length. The results show that RetNet has a consistent advantage across sequence length.

D Hyperparameters Used in Section 3.1

We 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 12 d 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.

E Hyperparameters Used in Section 3.2

Were-allocate the parameters in MSR and FFN for fair comparisons. Let d denote d model for simplicity here. In Transformers, there are about 4 d 2 parameters in self-attention where W Q , W K , W V , W O ∈ R d × d , and 8 d 2 parameters in FFN where the intermediate dimension is 4 d . In comparison, RetNet has 8 d 2 parameters in retention, where W Q , W K ∈ R d × d , W G , W V ∈ R d × 2 d , W O ∈ R 2 d × d . Notice that the head dimension of V is twice Q,K , similar to GAU [26]. 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 2 d . Meanwhile, we set the head dimension to 256 , i.e., 256 for

queries and keys, and 512 for values. For fair comparison, we keep γ identical among different model 511 sizes, where γ = 1 -e linspace(log 1 / 32 , log 1 / 512 ,h ) ∈ R h instead of the default value in Equation (8). 512

F Results on Open-Ended Generation Tasks 513

Table 9 presents one-shot performance on two open-ended question-answering tasks, including 514 SQUAD [39] and WebQS [5], with 6.7B models as follows. We report the recall metric in the table, 515 i.e., whether the answers are contained in the generated response. 516

G Inference Cost of Grouped-Query Retention

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.

517

518

519

520

521

522

523

524

525

526

527

528

529

530

531

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 532 LLaMA2 can run without out of memory. Table 10 shows that the inference cost of Transformers 533 increases with the sequence length. In contrast, RetNet is length-invariant. Moreover, RetNet-70B534 GQ2 achieves better latency, throughput, and GPU memory than LLaMA2-70B-2k/8k equipped 535

with grouped-query attention. Notice that the evaluation metrics are averaged over positions of 536 different sequence lengths for a fair comparison, rather than only considering the inference cost of 537 the maximum length. 538

Table 10: Inference cost of RetNet and LLaMA2-70B with difference batch size and length. LLaMA270B is equipped with grouped-query attention, reducing key/value heads by 8 × . '-GQ2' means grouped-query retention, which reduces half of key/value heads. '-2k' and '-8k' indicate sequence length for LLaMA2, while RetNet is length-invariant. RetNet is capable of large-batch inference and is favourable in terms of latency, throughput, and GPU memory.

H Hyperparameters Used in Section 3.8 539

540

541

542

543

544

545

546

547

548

549

550

551

552

553

554

555

556

557

558

559

560

561

562

563

564

565

566

567

568

569

570

571

572

573

574

575

576

577

578

579

580

581

582

583

584

585

586

587

588

589

590

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

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.

2. Limitations

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [NA]

Justification: There is no theoretical result in this paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: The experiment can be easily reproduced based on the model description, hyperparameter, and any well-known pre-training corpus.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

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.

6. 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.

7. 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).

698

699

700

701

702

703

704

705

706

707

708

709

710

711

712

713

714

715

716

717

718

719

720

721

722

723

724

725

726

727

728

729

730

731

732

733

734

735

736

737

738

739

740

741

742

743

744

745

746

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.

8. 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).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We follow the NeurIPS Code of Ethics in the research.

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).

10. 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.

747

748

749

750

751

752

753

754

755

756

757

758

759

760

761

762

763

764

765

766

767

768

769

770

771

772

773

774

775

776

777

778

779

780

781

782

783

784

785

786

787

788

789

790

791

792

793

794

795

796

797

798

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).

11. 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.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: The paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects. Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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. \nThe 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. \nIt 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.\n\nSTRENGTHS:\n1. The RetNet also shows competitive performance in language modeling and knowledge-intensive tasks compared to other Transformer variants and has the potential to replace Transformers for large language models.\n2. Achieves significantly better inference efficiency in terms of memory, speed, and latency.\n\nWEAKNESSES:\n1. 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.\n2. 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.\n\nQUESTIONS:\n1. 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? \n2. 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.\n\nLIMITATIONS:\nnone", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a linear attention model called RetNet for language modeling, which has a linear training complexity and constant inference complexity.\n\nSTRENGTHS:\n1. RetNet has both linear time complexity and constant inference memory complexity. \n2. RetNet has a chunk recurrent form which can be beneficial for speculative decoding.\n\nWEAKNESSES:\n1. 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].\n2. 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.\n3. 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.\n4. 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.\n5. 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.\n\n[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.\n\n[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\n\nQUESTIONS:\nCan you also add the results on MMLU and Qasper with the standard metrics besides perplexity?\n\nLIMITATIONS:\nNo, the authors should have a limitation section to point out the strong assumptions of their approximation of self-attention and relative position embedding.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-13"}

comments: 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.

STRENGTHS:
- The structure of RetNet is easy to understand and follow
- RetNet exhibits promising training and inference efficiency, and is able to scale up to 6B.
- Comprehensive evaluation on both language and vision tasks, highlighting its generalizability.

WEAKNESSES:
- 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:
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:
I didn't see serious problems.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-15