Memformer: A Memory-Augmented Transformer for Sequence Modeling

TL;DR

Memformer is a memory-augmented Transformer that uses fixed-size external dynamic memory and recurrence to model long sequences with linear time and constant memory complexity. It also introduces memory replay backpropagation (MRBP) to train recurrent memory with reduced memory cost, achieving comparable performance to Transformer-XL and Compressive Transformer on image generation and language modeling.

Key findings

• Memformer achieved comparable or better performance with much lower cost: on autoregressive MNIST image generation, Memformer (4 encoder + 8 decoder) reached perplexity 1.555 with 5.0B FLOPs, versus Transformer-XL with memory=784 at perplexity 1.611 with 49.1B FLOPs.
• On WikiText-103 language modeling, Memformer (4 encoder + 16 decoder, memory=1,024) reached perplexity 22.74 with 54B FLOPs, versus Transformer-XL memory=1,024 at perplexity 23.67 with 168B FLOPs.
• The paper reports Memformer uses 8.1x less memory space and is 3.2x faster on inference in its efficiency comparisons.
• MRBP reduced training GPU memory to 7,229 MB versus 16,177 MB for BPTT and 9,885 MB for gradient checkpointing, with relative speed x0.90 versus x1.00 for BPTT and x0.48 for gradient checkpointing.
• Memory writer analysis found that around 60% to 80% of memory slots in the middle of a document attended mainly to themselves, while only 5% to 10% attended completely to input tokens later; one slot had forgetting bias magnitude 3.20 versus 1.15 for others.

Methods

• Memformer processes sequences at the segment level with a Transformer encoder-decoder: the encoder reads from and writes to a fixed-size external dynamic memory, and the decoder predicts the next segment conditioned on the current segment and memory.
• Memory reading uses cross-attention from input queries to memory keys/values; memory writing uses slot attention where each memory slot attends only to itself and encoder outputs, plus a Biased Memory Normalization forgetting mechanism.
• Training uses Memory Replay Back-Propagation (MRBP), which replays memory at each timestep to reduce activation storage during backpropagation through time.
• Experiments cover MNIST image generation (28×28 images reshaped to 784 tokens, 256 vocabulary) and WikiText-103 language modeling (28K articles, average 3.6K tokens), with results reported as median from three trials.

Limitations

• The paper states that due to limited computational resources, the authors were unable to experiment on the more recent PG19 dataset.
• The paper notes there is no official implementation of Compressive Transformer, so they re-implemented a version of the same size for comparison.
• Evident from the text: evaluation is limited to MNIST image generation and WikiText-103 language modeling, and the memory writer analysis is qualitative/visual rather than a broad quantitative probe.