Source-linked AI summary
Retentive Network: A Successor to Transformer for Large Language Models
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, Furu Wei
TL;DR
Existing sequence-modeling approaches had not simultaneously delivered training parallelism, strong performance, and low inference cost. RetNet introduces multi-scale retention with parallel, recurrent, and chunkwise recurrent forms, and reports competitive language modeling alongside substantially more efficient training and inference. The paper concludes that these properties make RetNet a strong successor to Transformer for large language models.
Problem
Previous approaches did not simultaneously achieve training parallelism, good performance, and low inference cost compared with Transformers.
Method
RetNet replaces multi-head attention with multi-scale retention supporting parallel, recurrent, and chunkwise recurrent computation.
Results
RetNet is competitive with Transformers in language modeling while reducing reported training and inference costs.
Takeaways & Limitations
RetNet combines parallel training, O(1) inference, and efficient long-sequence modeling in one architecture for large language models.
Abstract
from arXiv · showhide
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 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. The intriguing properties make RetNet a strong successor to Transformer for large language models. Code will be available at https://aka.ms/retnet.
1 Introduction
RetNet is proposed to resolve the trade-off among training parallelism, model performance, and inference cost. Its retention mechanism supports these properties simultaneously, with experiments reporting strong efficiency and competitive language-modeling performance.
- Results: RetNet combines training parallelism, Transformer-comparable performance, and low-cost inference in the reported language-modeling experiments.The paper presents this combination as making the “impossible triangle” possible.
- Motivation: Prior approaches improved inference efficiency or explored alternative sequence mechanisms, but none matched Transformers across the full trade-off.Linearized attention and recurrent approaches could reduce inference costs, while other mechanisms explored Transformer replacements; their modeling capacity or performance remained limited.
- Approach: RetNet uses multi-scale retention with parallel, recurrent, and chunkwise recurrent computation paradigms.The parallel form supports training, the recurrent form supports O(1) inference, and the chunkwise form supports efficient long-sequence modeling.
- Results: 8.4× faster decoding and 70% lower memory are reported for a 7B model at 8k sequence length versus Transformers with key-value caches.These figures concern inference decoding and memory, respectively.
2 Retentive Networks
RetNet derives a retention mechanism with parallel, recurrent, and chunkwise recurrent forms, combining GPU-friendly training with efficient autoregressive inference and long-sequence modeling. Its architecture uses multi-scale retention, GroupNorm, and gating within Transformer-like residual blocks.
- Overall architecture: The network stacks L identical residual pre-LayerNorm blocks, each containing multi-scale retention and a feed-forward network.Inputs are packed into hidden-dimensional matrices and transformed through successive RetNet layers.
- Retention mechanism: RetNet derives retention’s recurrent and parallel forms, enabling parallel training while recurrently conducting inference.The recurrent and parallel representations are mathematically connected and support different computation settings.
- Retention mechanism: Chunkwise recurrence computes each chunk in parallel while passing cross-chunk information recurrently.This hybrid representation is designed to accelerate training, especially for long sequences.
- Gated multi-scale retention: Multi-scale retention uses multiple retention heads with different fixed γ scales, GroupNorm, and a swish gate to increase layer non-linearity.GroupNorm separately normalizes head outputs because multi-scale heads have different variance statistics.
- Computation paradigms: RetNet’s recurrent representation supports O(1) inference complexity, reducing memory and inference latency for autoregressive decoding.The recurrent representation is used during inference and achieves equivalent results according to the paper.
- Relation to prior methods: Retention removes softmax relative to attention and preserves high-dimensional states, supporting recurrent formulation and expressive sequence modeling.The paper connects retention to xPos-like relative position encoding and contrasts it with element-wise recurrent alternatives.
3 Experiments
Experiments compare RetNet with Transformer and variants across language modeling, downstream evaluation, training, inference, and architectural ablations. Results show competitive or favorable performance alongside lower training and inference costs.
- Language Modeling: RetNet achieves comparable language-modeling results with Transformer and empirically outperforms it when model size exceeds 2B.The comparison uses 1.3B, 2.7B, and 6.7B models and reports stable training.
- Downstream Evaluation: RetNet achieves comparable performance with Transformer in zero-shot and in-context learning on a wide range of downstream tasks.The evaluation uses 6.7B models with zero-shot and 4-shot settings across HellaSwag, BoolQ, COPA, PIQA, Winograd, Winogrande, and StoryCloze.
- Inference Cost: RetNet inference memory remains nearly constant for long sequences, while Transformer memory grows linearly because of key-value caches.RetNet’s additional memory consumption is about 3%, with model weights occupying 97%.
- Inference Cost: RetNet provides higher, length-invariant decoding throughput and latency that remains almost unchanged across batch sizes and input lengths.The reported comparisons cover memory cost, throughput, and latency for a 6.7B model on an A100-80GB GPU.
- Ablation Studies: Ablations indicate that the swish gate, GroupNorm, multi-scale decay, and decay mechanism improve language-modeling performance or training stability.GroupNorm balances variances across multi-head outputs, while different decay rates across heads improve language-modeling performance.
4 Conclusion
RetNet combines parallel, recurrent, and chunkwise recurrent representations for sequence modeling, targeting efficient inference and competitive performance. The authors position these properties as supporting deployment and future extensions.
- RetNet enables parallel, recurrent, and chunkwise recurrent representations for sequence modeling.
- RetNet provides inference efficiency in memory, speed, and latency, alongside favorable training parallelization and competitive Transformer performance.
- The authors identify O(1) inference complexity as a deployment benefit for large language models.
- Future work includes scaling RetNet, using structured prompting, training multimodal language models, and deploying models on edge devices.
A Hyperparameters
Table 7 lists the hyperparameters used for the models in Section 3.
- Table 7 summarizes the hyperparameters used for the models in Section 3.
B Grouped Results of Different Context Lengths
RetNet’s language-modeling perplexity is evaluated across different context lengths, with results reported as consistently advantageous over Transformer. Longer context can also improve RetNet’s results.
- RetNet outperforms Transformer across different context lengths in language-modeling perplexity.Evaluation uses 2048 text chunks and computes perplexity for the last 128 tokens.
- RetNet can use longer context for better results.