Source-linked AI summary
Augmenting Self-attention with Persistent Memory
Sainbayar Sukhbaatar, Edouard Grave, Guillaume Lample, Herve Jegou, Armand Joulin
TL;DR
The paper asks whether transformer feedforward sublayers can be removed while preserving the contextual modeling of self-attention. It replaces them with persistent key-value memory vectors inside a unified all-attention layer, which achieves competitive language-modeling results and matches or exceeds selected benchmarks.
Problem
Transformer feedforward layers contain most model parameters, yet their role alongside self-attention motivates revisiting whether they are necessary as a separate component.
Method
The paper merges self-attention and feedforward sublayers into an all-attention layer augmented with persistent key-value vectors that store context-independent information.
Results
The architecture is competitive with transformers across standard character- and word-level language-modeling benchmarks, including 3.4 ppl better than the previous best small model on WikiText-103.
Takeaways & Limitations
Persistent vectors can replace transformer feedforward layers without loss of performance, while simplifying the layer into a unified attention mechanism.
Abstract
from arXiv · showhide
Transformer networks have lead to important progress in language modeling and machine translation. These models include two consecutive modules, a feed-forward layer and a self-attention layer. The latter allows the network to capture long term dependencies and are often regarded as the key ingredient in the success of Transformers. Building upon this intuition, we propose a new model that solely consists of attention layers. More precisely, we augment the self-attention layers with persistent memory vectors that play a similar role as the feed-forward layer. Thanks to these vectors, we can remove the feed-forward layer without degrading the performance of a transformer. Our evaluation shows the benefits brought by our model on standard character and word level language modeling benchmarks.
1 Introduction
Transformers combine self-attention with feedforward layers, whose substantial parameter count suggests both components contribute to representation building. The paper replaces this two-stage structure with an all-attention layer using persistent memory and reports competitive language-modeling performance.
- Transformer layers gather contextual information with self-attention and transform it into rich representations with a subsequent feedforward layer.
- The proposed all-attention layer merges self-attention and feedforward sublayers into one unified attention layer.Persistent key-value vectors replace the feedforward sublayer and provide context-independent information.
- The all-attention layer directly builds representations from context and persistent memory without a feedforward transformation.
- The architecture is evaluated on standard word-level and character-level language-modeling benchmarks, with performance competitive with transformers.
2 Related work
Related work places the paper among research on neural language modeling and attention-based sequence models. Prior approaches span recurrent, convolutional, transformer, caching, self-attention, vocabulary, and regularization techniques.
- Language-modeling research has explored feed-forward, recurrent, gated convolutional, and transformer architectures.
- Prior transformer language models introduced caching and learnable self-attention spans to improve efficiency or handle longer sequences.
- Word-level language models address large-vocabulary computation with approximate or sampled softmaxes, subword units, and hierarchical softmax.
- Regularization work studies dropout, tied embedding and classifier weights, and adaptive softmax for improving generalization or reducing memory footprint.
- Related attention research includes self-attention in autoregressive models and applications across natural language processing and computer vision.
3 Transformer layer
A transformer layer stacks multi-head self-attention and position-wise feedforward sublayers, each followed by add-norm processing. Self-attention computes context-weighted representations, while the feedforward component independently transforms each position.
- Each transformer layer contains multi-head self-attention followed by a feedforward sublayer, with add-norm applied after each sublayer.The layer is part of a stack of identical transformer layers.
- Multi-head self-attention applies multiple attention heads in parallel to sequence vectors and combines their outputs.The concatenated head outputs are multiplied by an output matrix to produce d-dimensional vectors.
- Attention uses learned query, key, and value transformations with position encoding, then softmax-weighted context representations produce each output.Relative position encoding is defined as p(t, c) = u_t-c in this paper.
- The feedforward sublayer independently applies two affine transformations with a pointwise nonlinearity at every sequence position.Its ReLU activation is σ(x) = max(0, x), and the hidden dimension is typically 4 times d.
- Add-norm combines a sublayer output with a residual connection and then applies layer normalization.Layer normalization normalizes the output activations using their average and standard deviation.
4 Our approach
The approach replaces each transformer layer’s separate self-attention and feedforward sublayers with one all-attention layer. It adds persistent key-value vectors to self-attention so the layer can represent information independent of the immediate context while retaining language-modeling components such as relative positions and adaptive attention.
- Feedforward reinterpretation: The feedforward sublayer can be recast as an attention layer by replacing ReLU with Softmax and removing biases.Under this interpretation, the feedforward parameters act as key and value vectors, enabling a common implementation with self-attention.
- All-attention layer: The all-attention layer replaces both the self-attention and feedforward sublayers with a single attention layer.It applies attention simultaneously to input-conditioned sequence vectors and unconditioned persistent vectors.
- Persistent memory: Persistent vectors are shared key-value pairs that store information independent of the input context.They are stacked in matrices M_k and M_v and can represent general task knowledge or feedforward-like information.
- Attention computation: Persistent keys and values are concatenated with the input-conditioned key and value pools before attention scores and outputs are computed.The extended context includes the ordinary context plus persistent-vector indices, whose position encodings are zero.
- Attention computation: Multi-head all-attention concatenates head outputs, applies an output projection, and follows the sublayer with AddNorm.Persistent vectors are not shared between heads.
- Language modeling: For language modeling, the model uses relative position embeddings, caching, adaptive attention spans, and adaptive input-output mechanisms.These components address unbounded contexts, learned per-head context sizes, and large-vocabulary computation.
5 Experiments
The all-attention architecture is evaluated on character- and word-level language modeling benchmarks, where it matches or improves on strong transformer baselines. Ablations show that persistent vectors and their joint integration with context attention are important for performance.
- Main results: The experiments compare the all-attention network with state-of-the-art models on character-level enwik8 and text8 and word-level WikiText-103 benchmarks.The reported metrics are bits per character for character modeling and perplexity for word modeling.
- Main results: Our small enwik8 model outperforms similarly sized models, while the large model matches state-of-the-art performance with significantly fewer parameters.The cited result is reported for character-level modeling on enwik8.
- Main results: Our large text8 model is 0.01 bpc below the state-of-the-art with half the number of parameters, while the small model matches the best-performing model.These comparisons are reported on the text8 character-level benchmark.
- Main results: On WikiText-103, the network is 3.4 ppl better than the previous best comparable-size Transformer-XL, while larger models remain about 2 perplexity points better.The comparison is for word-level language modeling among small models, with a separate larger-model state of the art reported for context.
- Ablation study: Persistent vectors are crucial for Text8 performance, with good performance already at N = 1024; removing them performs poorly.The N = 0 model is equivalent to a transformer without feedforward sublayers, although it retains 38M parameters in its projection matrices.
- Ablation study: Joint attention over persistent and context vectors outperforms separated attention, while single-head, head-split, and FF-attn variants perform worse than the default all-attn model.The variants have the same number of parameters except head-split, which uses fewer because half of its persistent vectors are unused.
6 Conclusion
The paper proposes a unified attention layer that aggregates contextual and general information using persistent vectors. These vectors can replace transformer feedforward layers without performance loss and may clarify how transformer-like models process and store information.
- Persistent vectors extend self-attention with information complementary to short-term context.
- The unified attention layer replaces transformer feedforward layers without loss of performance.
- The simplified layer may improve understanding of how transformer-like sequence models process and store information.