Source-linked AI summary
Compressive Transformers for Long-Range Sequence Modelling
Jack W. Rae, Anna Potapenko, Siddhant M. Jayakumar, Timothy P. Lillicrap
TL;DR
Long-range sequence models need to retain useful information from the past without storing every activation. The paper compresses old Transformer memories into additional coarse memories and reports state-of-the-art language-modelling results, including 0.97 bpc on Enwik8 and 17.1 perplexity on WikiText-103.
Problem
Long-range sequence models face a storage problem when preserving useful information from distant past inputs.
Method
The Compressive Transformer compresses old hidden-state memories into smaller representations while attending jointly to recent granular and longer-term coarse memories.
Results
The Compressive Transformer outperforms existing architectures at long-range language modelling, with 0.97 bpc on Enwik8, and is also applied to speech and reinforcement learning.
Takeaways & Limitations
PG-19 provides a larger book-level benchmark for comparing and developing long-context language models.
Takeaways & Limitations
The method adds complexity and is unlikely to provide additional benefit when the task lacks long-range reasoning.
Abstract
from arXiv · showhide
We present the Compressive Transformer, an attentive sequence model which compresses past memories for long-range sequence learning. We find the Compressive Transformer obtains state-of-the-art language modelling results in the WikiText-103 and Enwik8 benchmarks, achieving 17.1 ppl and 0.97 bpc respectively. We also find it can model high-frequency speech effectively and can be used as a memory mechanism for RL, demonstrated on an object matching task. To promote the domain of long-range sequence learning, we propose a new open-vocabulary language modelling benchmark derived from books, PG-19.
1 INTRODUCTION
The Compressive Transformer addresses the Transformer’s storage and attention costs by compressing past activations into longer-term coarse memories while retaining granular short-term memories. The paper applies this mechanism to language, speech, and reinforcement learning, and introduces PG-19 for long-context language modelling.
- Motivation: Transformers preserve every time-step’s hidden activation and attend over this external memory, avoiding the compressed-state bottleneck of recurrent models.This representation retains more historical information but increases memory and attention costs.
- Motivation: Storing all activations increases attention computation and memory requirements, while sparse attention leaves storage costs unresolved and can require custom sparse kernels.The paper therefore returns to compact representations built with dense components such as convolutions.
- Compressive Transformer: The Compressive Transformer maps past hidden activations to fewer compressed memories and attends jointly to short-term granular and longer-term coarse memories.It is presented as a simple extension of the Transformer using the same attention mechanism over both memory sets.
- Applications: 400,000 steps: the Compressive Transformer models high-frequency speech with a trend of lower likelihood than TransformerXL and Wavenet.The model also compresses and uses past observations as a memory component within the IMPALA reinforcement-learning agent.
- Benchmark: Over double the size of existing language-modelling benchmarks: PG-19 contains Project Gutenberg book text with much longer contexts.The benchmark is introduced to promote long-context sequence modelling.
2 RELATED WORK
Related work extends Transformer context through linear-time operators, persistent memories, sparse attention, and dynamic attention spans. The Compressive Transformer incorporates Transformer-XL’s memory and relative positional embeddings while addressing limitations of existing approaches.
- Alternative attention mechanisms: Wu et al. (2019) show a linear-time convolution-like operator can exceed quadratic-time self-attention on sentence-to-sentence translation and sentence-level language modelling.
- Memory-based Transformers: Transformer-XL retains past activations and introduces relative positional embeddings, both of which the Compressive Transformer incorporates.
- Sparse attention: The Sparse Transformer attends to roughly √n memory locations but still requires retaining all memories during training.
- Sparse attention: On Enwik8, an 8,000-location attention window improves Sparse Transformer performance but does not significantly outperform simpler Transformer-XL with a much smaller window.
- Dynamic attention spans: Dynamic attention spans achieve state-of-the-art character-based language modelling, but dynamic sparse computation is inefficient on current dense-linear-algebra accelerators.The approach could be combined with a compressive memory.
3 MODEL
The Compressive Transformer extends Transformer-XL by retaining old activations in a compressed memory rather than discarding them. Its compression-based design can increase temporal range without increasing attention cost, while attention-reconstruction provides the most effective auxiliary objective among those considered.
- Architecture: The Compressive Transformer compacts past activations into a compressed memory, extending Transformer-XL’s layerwise memory instead of discarding sufficiently old activations.It is a variant of the Transformer and uses attention to propagate information over time.
- Memory update: At each layer, fixed-size input windows push hidden activations into fine-grained memory, while the oldest memories are compressed and stored in a secondary FIFO memory.The model defines nm and ncm as the memory and compressed-memory slots per layer; c controls compression coarseness.
- Compression functions: The model considers pooling, 1D and dilated convolutions, and most-used selection as compression functions, with pooling serving as a fast, simple baseline.The most-used scheme preserves memories ranked by average attention usage.
- Auxiliary losses: Attention-reconstruction works best because it preserves attended information while allowing information no longer attended to be discarded.This lossy objective reconstructs content-based attention over original and compressed memories and stops compression-loss gradients from entering the main network.
- Efficiency: 2 times greater temporal range is obtained with ncm = nm = n/2 and c = 3 at identical attention cost to TransformerXL.The Compressive Transformer’s maximum temporal range is l × (nm + c ∗ncm), with attention cost O(n2 + ns(nm + ncm)).
4 PG-19 BENCHMARK
PG-19 is proposed as an open-vocabulary language-modelling benchmark built from books published before 1919, providing long-form text for evaluating long-range sequence models. It uses word-level perplexity for model comparison and includes qualitative topic analysis.
- Benchmark motivation and construction: PG-19 uses Project Gutenberg books published before 1919 to provide training and benchmark data containing long-range dependencies.The benchmark is motivated by the need to evaluate models with longer-range memories on larger contexts.
- Benchmark motivation and construction: 28,752 books are included in PG-19, compared with 1,527 books in NarrativeQA.NarrativeQA requires corresponding Wikipedia summaries, resulting in a smaller book selection.
- Evaluation: PG-19 retains an open vocabulary without replacing rare words with an unknown token.The benchmark supports character-, byte-, or subword-based tokenization while retaining word-level perplexity as the comparison metric.
- Evaluation: Word-level perplexity is computed as e^(L/n_words), using the exact word count from Table 2 for normalization.The loss is accumulated over the validation or test subset under the chosen tokenization scheme.
- Qualitative analysis: An LDA topic model qualitatively identifies topics spanning art, education, exploration, war, ancient civilisations, and human-condition themes.The topic variety contrasts with the more objective domains of Wikipedia and news corpora.
5 EXPERIMENTS
The experiments evaluate the Compressive Transformer on language modelling, speech, and reinforcement-learning tasks, showing strong long-range modelling and effective compressed-memory use. Results include 0.97 bits-per-character on Enwik8, larger gains on infrequent WikiText-103 words, and human-level performance on an object-matching task.
- Enwik8: The Compressive Transformer achieves 0.97 bits-per-character on Enwik8, establishing a new state-of-the-art result on the benchmark.Evaluation used a memory size of 4096 for TransformerXL and compressed memory of 3072 for the Compressive Transformer.
- WikiText-103: Naive dynamic evaluation lowers WikiText-103 test perplexity to 16.1, compared with the published 16.4 from Krause et al. (2019).The dynamic-evaluation result comes from loading a checkpoint and continuing training for one epoch on the test set; such methods are excluded from Table 6.
- WikiText-103: On WikiText-103, modelling improves by 2.6% over TransformerXL for frequent words but by approximately 20% for infrequent words.The Compressive Transformer also shows a 10X improvement on rare words over the prior state-of-the-art LSTM language model.
- Memory analysis: The model’s average attention is evaluated across compressed memory, memory, and sequence buckets to determine whether it uses compressed memory.The analysis averages attention over 20,000 Enwik8 sequences with all three regions set to size 768.
- Reinforcement learning: Agents equipped with the Compressive Transformer solve the reinforcement-learning task to human-level, whereas compression rate 1 does not reach the same proficiency.Memory and compressed memory are fixed at 64, and results are averaged over 3 seeds across compression rates.
6 CONCLUSION
The paper presents compression as a way to extend Transformers’ temporal receptive field, showing benefits across text, speech, and reinforcement-learning vision tasks while introducing PG-19. It identifies added complexity as the main limitation and proposes combining recent granular memories with coarser compressed memories.
- 6 CONCLUSION: Compression extends the temporal receptive field of Transformer-based sequence models and improves long-range language modelling, while introducing the PG-19 book-level language-modelling benchmark.PG-19 is proposed for comparing long-range language models or pre-training long-range reasoning language models.
- 6 CONCLUSION: Compressive memories apply beyond text to speech waveform modelling and vision in a reinforcement-learning maze-like memory task.The comparisons use strong baselines: Wavenet for speech and IMPALA for reinforcement learning.
- 6 CONCLUSION: The main limitation is additional complexity, and the method is unlikely to help tasks without long-range reasoning.Compression is presented as a simpler alternative to dynamic or sparse attention, using components such as convolutions.
- 6 CONCLUSION: Future models may combine granular recent memories with coarser compressed memories rather than compressing the past immediately.Suggested directions include adaptive layerwise compression rates, shallow long-range memory with deep short-range memory, and RNNs as compressors.
A. Baevski and M. Auli. Adaptive input representations for neural language modeling. arXiv preprint arXiv:1809.10853, 2019.
This related-work section lists prior research on neural language modeling, sequence modeling, reinforcement learning, and representation learning. The cited works include approaches using attention, convolution, recurrence, sparse computation, caching, and pretraining.
- Related work: The bibliography also covers neural machine translation, latent Dirichlet allocation, BERT pretraining, DeepMind Lab, and distributed deep reinforcement learning.These entries respectively address jointly learned alignment and translation, latent Dirichlet allocation, bidirectional Transformer pretraining, a reinforcement-learning environment, and scalable distributed deep RL.
- Related work: Sequence-modeling research cited here spans convolutional, trellis, quasi-recurrent, hierarchical multiscale recurrent, and sparse Transformer architectures.The references explicitly identify these model families as convolutional sequence modeling, trellis networks, quasi-recurrent neural networks, hierarchical multiscale recurrent neural networks, and sparse Transformers.
A. Graves. Generating sequences with recurrent neural networks. arXiv preprint arXiv:1308.0850, 2013. … SUPPLEMENTARY MATERIALS
The supplementary materials comprise bibliographic references spanning recurrent sequence generation, long short-term memory, compression, attention, memory mechanisms, optimization, and related sequence-learning applications. The references also cover language modeling, speech synthesis, reading comprehension, visual explanation, and large-scale neural architectures.
- A. Graves. Generating sequences with recurrent neural networks. arXiv preprint arXiv:1308.0850, 2013.: The Graves subsection cites Neural Turing Machines, hybrid computing with dynamic external memory, Hypernetworks, and explicit memory representations.
- S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997.: The Hochreiter–Schmidhuber subsection is accompanied by references on neural text degeneration and linear-time neural machine translation.
- M. Hutter. The human knowledge compression contest. URL http://prize. hutter1. net, 6, 2012.: The Hutter subsection includes references on optimization, reading comprehension, multiplicative LSTMs, dynamic evaluation, product-key memory, and pointer-sentinel mixtures.
- M. Hutter. The human knowledge compression contest. URL http://prize. hutter1. net, 6, 2012.: Additional Hutter-subsection references address recurrent language modeling, WaveNet speech generation, Parallel WaveNet, broad-context word prediction, and sparse memory reads and writes.
- M. Hutter. The human knowledge compression contest. URL http://prize. hutter1. net, 6, 2012.: The bibliography further covers activation memorization, persistence and transience of memory, back-propagation, relational recurrent neural networks, and model-parallel language-model training.
- M. Hutter. The human knowledge compression contest. URL http://prize. hutter1. net, 6, 2012.: Optimization and transformer-related citations include increasing batch size without learning-rate decay, adaptive attention spans, the Transformer architecture, lightweight dynamic convolutions, and XLNet.
- M. Hutter. The human knowledge compression contest. URL http://prize. hutter1. net, 6, 2012.: Related application references include end-to-end dense video captioning with masked transformers and aligning books and movies for story-like visual explanations.
- M. Hutter. The human knowledge compression contest. URL http://prize. hutter1. net, 6, 2012.: The references conclude with recurrent highway networks, extending the supplementary bibliography’s coverage of recurrent architectures.
A COMPRESSION ACROSS LAYERS
The section analyzes attention-reconstruction compression loss across network layers using trained models on Enwik8 and WikiText-103. Character-based language modelling has roughly an order-of-magnitude lower compression loss than word-level modelling.
- A COMPRESSION ACROSS LAYERS: Compression loss is examined by layer index to investigate how representation compressibility varies with network depth.Here, compression loss denotes the attention-reconstruction attention loss.
- A COMPRESSION ACROSS LAYERS: The analysis uses a 24-layer model trained on Enwik8 and an 18-layer model trained on WikiText-103.
- A COMPRESSION ACROSS LAYERS: About one order of magnitude separates character-based from word-level language modelling compression loss, with the character-based loss lower.
B COMPARISON OF COMPRESSED MEMORY SIZES · C PG-19 PREPROCESSING
The section identifies 3, 072 as the best compressed memory size in the reported sweep and describes the preprocessing applied to PG-19’s Gutenberg texts. It also specifies the compression setup and Enwik8 model configuration used for comparison.
- B COMPARISON OF COMPRESSED MEMORY SIZES: 3, 072 was the best compressed memory size identified in the sweep for the compared models.The comparison used the best test perplexity obtained on WikiText-103 and Enwik8 across compressed memory sizes.
- B COMPARISON OF COMPRESSED MEMORY SIZES: Both models used a 1D convolution compression network with a compression rate of 3.
- B COMPARISON OF COMPRESSED MEMORY SIZES: The Enwik8 model used embedding size 1024, 8 attention heads, 24 layers, and MLP hidden size 3072.
- B COMPARISON OF COMPRESSED MEMORY SIZES: The Enwik8 model used a sequence window size of 768 and a memory size of 768.
- B COMPARISON OF COMPRESSED MEMORY SIZES: Table 8 reports compressed memory size versus test performance for Enwik8.
- B COMPARISON OF COMPRESSED MEMORY SIZES: Table 9 reports compressed memory size versus test performance for WikiText-103.
- C PG-19 PREPROCESSING: The Gutenberg texts were minimally pre-processed by removing boilerplate license text.
- C PG-19 PREPROCESSING: Discriminatory words were replaced with a unique ⟨DWx⟩token using the Ofcom list of discriminatory words.
D PG-19 TOPICS
This section presents examples of topics from the PG-19 corpus, with top words generated using an LDA topic model.
- D PG-19 TOPICS: Top words are presented for selected topics in the PG-19 corpus.The section introduces these top-word lists as examples from the corpus.
- D PG-19 TOPICS: The topic representations were generated with an LDA topic model.The method is attributed to Blei et al. (2003).
- D PG-19 TOPICS: Table 10 provides examples of top topics on the PG-19 corpus.It summarizes the topic examples presented in this section.
E PG-19 SAMPLES
PG-19 samples show that the Compressive Transformer can continue book excerpts in their style while sustaining character information over hundreds of words. Because training uses unreset state, short-context prompts begin out of sample, but samples generally improve toward their ends.
- Sample quality: The model continues PG-19 test-set book excerpts in their style, generating artificial dialogue or descriptive text and remembering character names over hundreds of words.Samples use Nucleus Sampling with p = 0.98.
- Sample quality: Unreset state makes relatively short prompts slightly out of sample because memory and compressed memory may initially be empty, although samples usually improve toward the end.During training, both memory components are always full.