PaperlayerView exact revision on arXiv

Source-linked AI summary

Memformer: A Memory-Augmented Transformer for Sequence Modeling

Qingyang Wu, Zhenzhong Lan, Kun Qian, Jing Gu, Alborz Geramifard, Zhou Yu

arXiv:2010.06891v2cs.CL

TL;DR

Transformers and other sequence models face efficiency and memory constraints when processing long histories. Memformer uses recurrent interaction with fixed-size external dynamic memory, alongside memory replay back-propagation for training. The paper reports comparable performance with improved efficiency and information retention, while noting limits of linear sparse attention for autoregressive access to the past.

  • Problem

    Transformers compute quadratic self-attention and most efficient attention approaches still require memory space that grows linearly with sequence length.

    Source: §1 · p. 1

  • Method

    Memformer combines a fixed-size external dynamic memory with recurrence and uses memory replay back-propagation to reduce training memory cost.

    Source: §1 · p. 2

  • Results

    Memformer achieved comparable performance with great efficiency and preserved information from the distant past.

    Source: §5 · p. 8

  • Takeaways & Limitations

    Memformer processes long sequences with linear time complexity and constant memory complexity, and its memory slots retain information over extended periods.

    Source: §1 · p. 2, §5 · p. 8

    Linear sparse attention cannot guarantee that an autoregressive decoder can access every past token; the passage identifies Sparse Transformer as a theoretical exception.

    Source: §2.1 · p. 2

Abstract

from arXiv · show

Transformers have reached remarkable success in sequence modeling. However, these models have efficiency issues as they need to store all the history token-level representations as memory. We present Memformer, an efficient neural network for sequence modeling, that utilizes an external dynamic memory to encode and retrieve past information. Our model achieves linear time complexity and constant memory space complexity when processing long sequences. We also propose a new optimization scheme, memory replay back-propagation (MRBP), which promotes long-range back-propagation through time with a significantly reduced memory requirement. Experimental results show that Memformer has achieved comparable performance compared to the baselines by using 8.1x less memory space and 3.2x faster on inference. Analysis of the attention pattern shows that our external memory slots can encode and retain important information through timesteps.

1 Introduction

Transformer efficiency methods reduce self-attention computation, but most still require memory that grows linearly with sequence length. Memformer combines a fixed-size external dynamic memory with recurrence, and the paper reports comparable task performance with improved efficiency.

1 Introduction

  • Most efficient self-attention approaches still require linear memory space, despite reducing self-attention complexity to process longer sequences.The passage describes Transformer attention as having quadratic computation complexity.Source: §1 · p. 1
  • Transformer-XL and Compressive Transformer discard distant-past information, imposing a theoretical maximum temporal range with fixed-size memory.Transformer-XL stores raw hidden states, while Compressive Transformer further compresses them.Source: §1 · pp. 1–2
  • Memformer combines a fixed-size external dynamic memory with Transformer recurrence, yielding theoretically infinite temporal memorization range, linear computation, and constant memory space complexity.The paper also introduces memory replay back-propagation to reduce the memory cost of training with large memory representations.Source: §1 · p. 2
  • Across autoregressive image generation and language modeling, Memformer performs on par with Transformer and Transformer XL with large memory while using computation and memory more efficiently.The authors also report that Memformer retains information for an extended period.Source: §1 · p. 2

2 Related Work

The related work reviews efficient attention and recurrent-memory approaches, noting limits in autoregressive access to the past and in memory use. It positions Memformer’s dynamic memorization as an alternative to token-level history storage and complex external-memory mechanisms.

2.1 Sparse Attention

  • Sparse Transformer, Longformer, and Big Bird use sparse attention patterns and global tokens to reduce attention cost, with Longformer and Big Bird described as reaching O(N) complexity.Their global tokens encode sequence-wide information while local attention covers nearby tokens.Source: §2.1 · p. 2
  • Linear sparse attention cannot guarantee that an autoregressive decoder can access every past token; the passage identifies Sparse Transformer as the exception with theoretical full-history coverage.Sparse attention’s theoretical soundness is proven for bidirectional encoders, not decoders.Source: §2.1 · p. 2

2.2 Linear Attention

  • Linear Transformer’s autoregressive cumulative summation can grow very large over long, variable-length inputs, causing overflow and gradient instability.The passage also notes that Linformer had not been applied to autoregressive decoding.Source: §2.2 · p. 2

2.3 Recurrence and Memory

  • Transformer-XL and Compressive Transformer use recurrence over cached hidden states, but their context has a theoretical maximum range and good performance requires huge memory.Compressive Transformer further compresses previous cached hidden states.Source: §2.3 · p. 2

2.3.1 Dynamic Memorization

  • Dynamic memorization stores information beyond token-level history representations and has no theoretical upper bound on temporal range, while prior external-memory systems are slow and unstable to train.The passage names Neural Turing Machine and Differential Neural Computer as examples of those prior systems.Source: §2.3.1 · p. 3

3 Methods

Memformer processes sequences in segments, using an encoder to interact recurrently with fixed-size external memory and a decoder to predict the next segment. Memory reading uses cross-attention from input queries to memory slots.

3.1 Segment-level Sequence Modeling

  • Long sequences are divided into segments because interacting with a large external memory for every token is unaffordable.Each segment contains L tokens.Source: §3.1 · p. 3
  • Memformer’s encoder encodes each segment into memory and retrieves the previous memory, while the decoder predicts the next segment; generated segments are fed back to model longer sequences autoregressively.The method overview includes segment-level modeling, memory reading and writing, and MRBP training.Source: §3.1 · p. 3, §3 · p. 3

3.2 External Dynamic Memory Slots

  • External dynamic memory stores high-level representations of past inputs and changes as the model recurrently encodes and retrieves information.The design allocates a constant number of vectors as memory slots, so inference memory consumption remains constant as input length grows.Source: §3.2 · pp. 3–4

3.3 Memory Reading

  • Memory reading uses cross-attention: input-sequence queries attend to key-value pairs projected from all memory slots to produce hidden states.Reading occurs in every encoder layer, increasing the chance of retrieving needed information from a large memory.Source: §3.3 · p. 4

3.4 Memory Writing

Memory writing updates fixed slots with segment information, while a forgetting mechanism clears less important information. Slot attention restricts each slot’s inputs, and Biased Memory Normalization (BMN) controls forgetting while stabilizing memory representations.

3.4 Memory Writing

  • Memory writing updates information with a slot-attention module and removes unimportant information with a forgetting method.Writing occurs only at the encoder’s last layer to store high-level contextual representations.Source: §3.4 · p. 4

3.4.1 Update via Memory Slot Attention

  • Each memory slot attends only to itself and encoder outputs, preventing direct writes between slots; temperature τ < 1 sharpens attention toward fewer slots or token outputs.Attention lets a slot preserve its old information or update it with new information.Source: §3.4.1 · p. 4

3.4.2 Implementation of Memory Writer

  • A slot attending only to itself remains unchanged at the next timestep; otherwise, its position relative to the terminal state influences how readily its information is forgotten.The terminal state is also the learnable initial state; ma is easier to forget than mb because it is closer to that state.Source: §3.4.2 · p. 5, §3.4.3 · p. 5

3.4.3 Forgetting Mechanism

  • BMN normalizes memory slots at every step to prevent unbounded growth and maintain gradient stability, while a learnable bias helps forget previous information.Normalized slots lie on a sphere; the bias controls forgetting speed and direction.Source: §3.4.3 · p. 5

3.5 Memory Replay Back-Propagation

Training Memformer with back-propagation through time (BPTT) would require storing many intermediate activations, making memory consumption impractical. Memory Replay Back-Propagation (MRBP) replays memory at each timestep to backpropagate across long unrolls more efficiently.

3.5 Memory Replay Back-Propagation

  • BPTT requires storing intermediate activations from the unrolled computational graph, leading to impractically high memory consumption for Memformer.The memory writer needs BPTT during training to learn to retain long-term information.Source: §3.5 · pp. 5–6
  • MRBP replays memory at each timestep to backpropagate over long unrolls, offering a more efficient alternative to gradient checkpointing.It traverses only the critical path during the forward pass and recomputes the local timestep’s partial graph during the backward pass.Source: §3.5 · p. 6

4 Experiments

Across image generation and language modeling, Memformer is reported to match or outperform Transformer-XL with lower computation cost. Its memory slots exhibit distinct update patterns, including slots that retain information from earlier timesteps.

4.1 Computation and Memory Cost

  • Vanilla Transformer computation grows as O(N^2), while Transformer-XL and Memformer have O(N) computation; memory costs are O(K × L) for Transformer-XL and O(K) for Memformer.In the tested large-memory setting, Memformer used 8.1x less GPU memory than Transformer-XL.Source: §4.1 · p. 6, Figure 5 · p. 7

4.2 Autoregressive Image Generation

  • For autoregressive image generation, Memformer with 4 encoder and 8 decoder layers achieved 1.555, versus 1.611 for the best Transformer-XL baseline (memory size 784), using nearly 10% of its FLOPs.It also outperformed the Transformer Decoder with the entire input sequence.Source: Table 1 · p. 6, §4.2 · p. 7
  • Reducing Memformer to 4 encoder and 4 decoder layers lowered image-generation performance to 1.628, indicating the importance of decoder-layer count.In the language-modeling ablations, reducing memory size to 512 lowered performance; removing the memory module raised perplexity to 25.57, similar to Transformer-XL with memory size 128.Source: §4.2 · p. 7, §4.3 · p. 8

4.3 Language Modeling

  • Memformer with 4 encoders, 16 decoders, and memory size 1,024 achieved the best performance, requiring computation cost 54 and performing much better than Transformer-XL with memory size 1,024.Table 2 reports experimental results on language modeling and states that the method is 3.2 times faster here.Source: Table 2 · p. 7, §4.3 · p. 8

4.3.1 Memory Writer Analysis

  • During document processing, around 60% to 80% of memory slots focused on themselves and did not update at the current timestep; another type partly attended to itself and aggregated information from other tokens.The self-focused slots were suggested to carry information from the distant past.Source: §4.3.1 · p. 8
  • Input-focused slots accounted for nearly all memory slots at the beginning but only 5% to 10% later; their forgetting-bias magnitude was 3.20 versus 1.15 for some other slots.For one such slot, attention to named entities and verbs formed a compressed sentence representation.Source: §4.3.1 · p. 8

5 Conclusion

Memformer combines external dynamic memory with autoregressive modeling, with linear time and constant memory complexity. The authors report efficient performance and retention of information from the distant past, and suggest potential relevance to dialog and interactive systems.

5 Conclusion

  • Memformer uses external dynamic memory to process long sequences with linear time and constant memory complexity.The conclusion also describes it as an autoregressive model.Source: §5 · p. 8
  • The authors report comparable performance with efficiency and preservation of information from the distant past.Source: §5 · p. 8
  • The authors suggest that enhanced memory capacity may motivate work on recurrence and autoregressive modeling for dialog and interactive systems.Source: §5 · pp. 8–9

A MRBP Efficiency Test

MRBP is compared with BPTT and gradient checkpointing under matched Memformer hyperparameters. It saves more GPU memory than the alternatives with only slight speed degeneration.

A MRBP Efficiency Test

  • The comparison uses the Memformer model with all hyperparameters held the same across methods.The methods compared are BPTT, gradient checkpointing, and MRBP.Source: A MRBP Efficiency Test · p. 11
  • Table 3 compares Memory Replay Back-Propagation performance, with evaluation speed measured in seconds per sample.The table caption defines BPTT as back-propagation through time and GC as gradient checkpointing.Source: Table 3 · p. 11
  • MRBP saves more GPU memory with only slight speed degeneration, while gradient checkpointing is much slower (x0.48) and MRBP has speed x0.90.BPTT is fastest but costs the most memory because it unrolls the entire computational graph.Source: A MRBP Efficiency Test · p. 11

B Training Details

The training details identify the hardware and approximate training durations for image generation and language modeling.

B Training Details

  • Table 4 is identified as the paper's training-details table.Source: Table 4 · p. 11
  • Image generation training took about one day on one GPU, while language modeling took approximately four days on four GPUs.The hardware used was NVIDIA V100 16GB and 2080Ti 11GB.Source: B Training Details · p. 11
  • Training used NVIDIA V100 16GB and 2080Ti 11GB GPUs.Source: B Training Details · p. 11

C Effects of Time Horizon and Memory Size

The section examines how back-propagation horizon and memory size affect Memformer performance. Longer horizons improve perplexity until gains nearly disappear, while memory-size gains diminish beyond eight slots.

C Effects of Time Horizon and Memory Size

D Implementation of Memory Writer

The Memory Slot Attention module updates memory from the previous memory state and encoder outputs, using sparse attention to limit interference between slots.

D Implementation of Memory Writer

Loading 2010.06891v2…