In plain words: A new sequence-processing design called retention runs three ways: all at once for fast training, step by step for cheap generation, or in chunks for long texts. It matches the usual transformer's language quality while keeping each new word's cost constant, so decoding uses less memory and runs faster.
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 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.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuqing Xia, Jilong Xue, Jianyong Wang, Furu Wei
arXiv:2307.08621 · cs.CL, cs.LG · submitted Jul 17, 2023 · updated Aug 9, 2023
abstract · pdf · html
- Retention replaces the softmax in attention with an exponential decay along the sequence dimension. This allows formulating retention in a recurrent form for efficient O(1) inference.
- Retention heads use different decay rates (gamma values) for multi-scale modeling. Attention heads use the same softmax.
- Retention outputs are normalized per-head with GroupNorm before concatenation. Attention uses LayerNorm on the concatenated output.
- Retention can be computed in parallel, recurrent, or chunkwise recurrent modes. Attention is only parallel.
- The recurrent form enables RetNets to summarize long previous context into a fixed-size state during inference. Attention recomputes on the full context each step.
- So in summary, retention adapts attention to enable recurrent modeling and multi-scale decays. This provides efficiency benefits and competitive performance to Transformers.