Extended Mind Transformers

TL;DR

Extended Mind Transformers enhance long-input processing in language models by utilizing a memory bank with improved retrieval and positional encoding, outperforming current methods by 6% on a new benchmark.

Contribution

The paper introduces improvements to Memorizing Transformers, enabling better long-range memory retrieval without fine-tuning, and provides a new benchmark for evaluation.

Findings

Outperform state-of-the-art by 6% on average

External memory retrieval is crucial in most decoder layers

Proper positional encoding updates are essential for memory access

Abstract

Pre-trained language models demonstrate general intelligence and common sense, but long inputs quickly become a bottleneck for memorizing information at inference time. We resurface a simple method, Memorizing Transformers (Wu et al., 2022), that gives the model access to a bank of pre-computed memories. We show that it is possible to fix many of the shortcomings of the original method, such as the need for fine-tuning, by critically assessing how positional encodings should be updated for the keys and values retrieved. This intuitive method uses the model's own key/query system to select and attend to the most relevant memories at each generation step, rather than using external embeddings. We demonstrate the importance of external information being retrieved in a majority of decoder layers, contrary to previous work. We open source a new counterfactual long-range retrieval benchmark, and show that Extended Mind Transformers outperform today's state of the art by 6% on average.