Source-linked AI summary
LongNet: Scaling Transformers to 1,000,000,000 Tokens
Jiayu Ding, Shuming Ma, Li Dong, Xingxing Zhang, Shaohan Huang, Wenhui Wang, Nanning Zheng, Furu Wei
TL;DR
Long-sequence modeling must overcome computational costs without sacrificing Transformer expressivity. LongNet introduces dilated attention and distributed sequence training to scale beyond 1 billion tokens. Experiments report efficient scaling and strong performance on both long and short sequences.
Problem
Existing methods for scaling sequence length are limited by computational complexity, training parallelization, or reduced model expressivity at regular sequence lengths.
Method
LongNet replaces standard attention with dilated attention and uses distributed training across sequence partitions for extremely long inputs.
Results
Nearly constant runtime enables LongNet to scale sequence length to 1 billion tokens, while experiments report strong performance on both long and short sequences.
Takeaways & Limitations
LongNet provides a Transformer-based approach for modeling sequences at billion-token scale while retaining performance on shorter sequences.
Abstract
from arXiv · showhide
Scaling sequence length has become a critical demand in the era of large language models. However, existing methods struggle with either computational complexity or model expressivity, rendering the maximum sequence length restricted. To address this issue, we introduce LongNet, a Transformer variant that can scale sequence length to more than 1 billion tokens, without sacrificing the performance on shorter sequences. Specifically, we propose dilated attention, which expands the attentive field exponentially as the distance grows. LongNet has significant advantages: 1) it has a linear computation complexity and a logarithm dependency between any two tokens in a sequence; 2) it can be served as a distributed trainer for extremely long sequences; 3) its dilated attention is a drop-in replacement for standard attention, which can be seamlessly integrated with the existing Transformer-based optimization. Experiments results demonstrate that LongNet yields strong performance on both long-sequence modeling and general language tasks. Our work opens up new possibilities for modeling very long sequences, e.g., treating a whole corpus or even the entire Internet as a sequence.
1 Introduction
Scaling sequence length offers larger memory and receptive fields, but existing approaches face a trade-off between computational complexity and model expressivity. LongNet addresses this with dilated attention and distributed training, scaling to 1 billion tokens with nearly constant runtime.
- Longer sequences provide larger memory and receptive fields, while exposing more complex causality and reasoning paths for training.
- Existing sequence models struggle to balance efficient long-context computation with sufficient model expressivity.RNN-style models limit training parallelization, while state space models perform less well than Transformers at regular lengths.
- Efficient Transformer variants reduce quadratic self-attention costs, but local or convolutional patterns can forget early tokens and prior methods had not reached 1 billion tokens.
- LongNet replaces vanilla attention with dilated attention, allocating attention exponentially less as token distance grows.The design targets linear computation complexity while retaining access to tokens across long distances.
- Distributed LongNet training partitions the sequence across nodes, enabling 1B-token inputs with nearly constant runtime instead of vanilla Transformer's quadratic growth.
2 LONGNET
LONGNET replaces vanilla self-attention with dilated attention, combining local precision with exponentially expanding long-range coverage. Its design provides approximately linear computation, logarithmic token dependency, and distributed sequence-dimension training.
- Dilated Attention: Dilated attention splits inputs into segments and sparsifies them along the sequence dimension using interval-based row selection.The resulting sparsified segments are processed in parallel and scattered and concatenated into the output.
- Implementation: Dilated attention can be implemented as dense attention between gathering and scattering operations, enabling reuse of vanilla-attention optimizations.The method reduces computation by a factor of N divided by w r^2 over vanilla attention.
- Dilated Attention: Dynamic weights mix dilated attentions in parallel, using attention-softmax denominators rather than learnable fixed weights.For queries attending to keys across dilated attentions, the method gathers the keys and calculates softmax together.
- Dilated Attention: Geometric segment sizes and dilation rates create an exponentially expanding attentive field while preserving multiple attention patterns.The segment size trades globality for efficiency, while dilation reduces computation by approximating the attention matrix.
- Dilated Attention: Different heads shift their sparsified query-key-value selections, and their outputs are concatenated as in vanilla multi-head attention.This gives heads different attention patterns while retaining the standard multi-head output computation.
- Computational Complexity and Token Dependency: Dilated attention has approximately O(Nd) computation complexity and approximately O(log N) token dependency.The construction also supports distributed training by partitioning the sequence dimension and collecting key-value pairs when segments span devices.
3 LONGNET as a Distributed Trainer: Scaling up to 1B Tokens
LONGNET distributes sequence-dimension computation across devices while using dilated attention to keep communication and runtime manageable for extremely long sequences. This design enables scaling to 1 billion tokens, whereas vanilla attention becomes increasingly expensive with sequence length.
- Distributed algorithm: LONGNET is designed to address single-GPU computation and memory limits that remain even after dilated attention reduces complexity to O(Nd).The distributed algorithm specifically targets scaling sequence length beyond what one device can support.
- Distributed algorithm: LONGNET partitions the input sequence across devices to distribute computation along the sequence dimension.The algorithm can scale beyond two devices and operates alongside data, model, and pipeline parallelism.
- Distributed algorithm: For local segments, LONGNET computes attention on each device; for larger segments, it gathers sparsified key-value pairs before cross-attention.Local queries attend to globally gathered key-value pairs, and device outputs are concatenated into the final attention output.
- Distributed algorithm: The distributed backward pass uses reduce-scatter, while gathered key and value sizes remain independent of sequence length, making communication cost constant.This communication property differs from vanilla attention.
- Scaling up to 1B tokens: 1 billion tokens are reached with almost constant latency, while vanilla attention latency rises dramatically as sequence length grows.Both methods use FlashAttention; the reported runtime comparison supports LONGNET’s scaling advantage.
4 Experiments on Language Modeling
The experiments evaluate LONGNET language modeling across sequence lengths, context-window scaling, and model sizes, comparing it with dense and sparse Transformer baselines. LONGNET consistently improves language-modeling efficiency and performance while retaining scaling behavior.
- Experimental setup: LONGNET replaces standard attention with dilated attention in a MAGNETO-based language model trained on The Stack.The base configuration uses hidden dimension 768, 12 attention heads, and 12 decoder layers.
- Experimental setup: The models are compared from 2K to 32K sequence lengths with matched attention-layer differences and constant tokens per batch.LONGNET uses segment lengths from 2048 to 32768 and dilated ratios from 1 to 12.
- Language modeling results: LONGNET consistently outperforms baseline models on Stack perplexity, while longer training sequences generally improve language modeling.Inference extrapolation does not apply when the sequence is much longer than the supported model length.
- Scaling curves of sequence length: LONGNET achieves lower test loss with less computation than vanilla Transformers as training context length increases.Both model types benefit from larger training contexts, but LONGNET scales context length more efficiently.
- Scaling up model size: LONGNET follows a power-law scaling curve across model sizes from 125 million to 2.7 billion parameters.The 2.7B model uses 300B tokens, while the other models use about 40B tokens.
- Long-context prompting: LONGNET test loss gradually decreases as the prompting context window grows from 2K to 32K.The experiment keeps suffixes fixed while increasing prefix lengths.
5 Conclusion and Future Work
The conclusion presents LONGNET as a Transformer variant for billion-token sequences that preserves shorter-sequence performance. It identifies dilated attention and distributed training as central components and points to broader future applications.
- Conclusion: LONGNET scales sequence length to 1 billion tokens and beyond without loss on shorter sequences.The conclusion reports superior performance over strong baselines on both long and short sequences.
- Conclusion: Dilated attention reduces Transformer computation complexity from quadratic to linear.It replaces standard attention as the core component of LONGNET.
- Conclusion: LONGNET parallelizes sequence training across multiple GPU devices as a distributed trainer.The conclusion presents this as part of its support for extremely long sequences.
- Future work: Future work extends LONGNET toward multimodal language modeling, BEiT pretraining, and genomic data modeling.These applications are identified as directions rather than completed experiments.
A Hyperparameters
The appendix includes hyperparameter tables for the models used in Table 2 and for the experiments shown in Figure 7(a).
- Table 3: Table 3 lists hyperparameters for the models evaluated in Table 2.The table is associated with the Table 2 experiments.
- Appendix organization: The appendix organizes hyperparameter information by the corresponding experiment or result display.Table 3 corresponds to Table 2, while Table 4 corresponds to Figure 7(a).
- Table 4: Table 4 lists hyperparameters for the experiments presented in Figure 7(a).The table is associated with the Figure 7(a) experiments.