Source-linked AI summary
Augmenting Language Models with Long-Term Memory
Weizhi Wang, Li Dong, Hao Cheng, Xiaodong Liu, Xifeng Yan, Jianfeng Gao, Furu Wei
TL;DR
Existing LLMs are limited by fixed input lengths and therefore cannot readily exploit rich information from long past contexts. LONGMEM uses a frozen backbone LLM to encode cached memory and a residual SideNet to retrieve and read it, separating memory encoding from adaptation. It improves long-context modeling and memory-augmented in-context learning, including 40.5% accuracy on ChapterBreak and gains with thousands of demonstrations.
Problem
Fixed input-length limits restrict LLMs from utilizing rich long-context information and sufficient demonstration examples from past inputs.
Method
LONGMEM freezes the backbone LLM as a memory encoder and uses an adaptive residual SideNet as a decoupled memory retriever and reader.
Results
40.5% ChapterBreak identification accuracy is reported as state of the art, alongside -1.38∼-1.62 perplexity improvements across Gutenberg-2022 length splits and broad in-context learning gains.
Takeaways & Limitations
LONGMEM enables language models to memorize and use long-form context, including thousands of auxiliary demonstration examples for memory-augmented in-context learning.
Takeaways & Limitations
Memorizing Transformer’s coupled design is vulnerable to memory staleness, while LONGMEM’s efficient SideNet setup considers fewer SideNet layers than backbone layers.
Abstract
from arXiv · showhide
Existing large language models (LLMs) can only afford fix-sized inputs due to the input length limit, preventing them from utilizing rich long-context information from past inputs. To address this, we propose a framework, Language Models Augmented with Long-Term Memory (LongMem), which enables LLMs to memorize long history. We design a novel decoupled network architecture with the original backbone LLM frozen as a memory encoder and an adaptive residual side-network as a memory retriever and reader. Such a decoupled memory design can easily cache and update long-term past contexts for memory retrieval without suffering from memory staleness. Enhanced with memory-augmented adaptation training, LongMem can thus memorize long past context and use long-term memory for language modeling. The proposed memory retrieval module can handle unlimited-length context in its memory bank to benefit various downstream tasks. Typically, LongMem can enlarge the long-form memory to 65k tokens and thus cache many-shot extra demonstration examples as long-form memory for in-context learning. Experiments show that our method outperforms strong long-context models on ChapterBreak, a challenging long-context modeling benchmark, and achieves remarkable improvements on memory-augmented in-context learning over LLMs. The results demonstrate that the proposed method is effective in helping language models to memorize and utilize long-form contents. Our code is open-sourced at https://aka.ms/LongMem.
1 Introduction
LONGMEM addresses fixed input-length limits by separating long-context memory encoding from memory retrieval and fusion. It improves long-context modeling and memory-augmented in-context learning while avoiding the memory staleness associated with coupled designs.
- Motivation: Fixed input-length limits prevent existing LLMs from using rich information beyond a single context window.Scaling context length increases training cost, while dense Transformer attention remains quadratically expensive.
- Motivation: Memorizing Transformer handles up to 65k tokens but suffers memory staleness because one updated model both encodes and uses cached representations.Parameter updates can create distributional shifts between older cached representations and representations from the latest model.
- Method: LONGMEM freezes the backbone LLM as a memory encoder and uses a residual SideNet as the memory retriever and reader.The decoupled architecture caches previous-context key-value pairs and fuses retrieved memory through the SideNet.
- Method: The decoupled design resolves memory staleness and avoids catastrophic forgetting during efficient memory-augmented adaptation.The backbone retains pretrained knowledge while only the SideNet is adapted.
- Results: LONGMEM consistently outperforms strong baselines on long-text language modeling and memory-augmented in-context learning.It supports full-length book contexts and thousands of task-relevant demonstration examples in memory.
- Results: 40.5% identification accuracy on ChapterBreak is reported as state of the art, while Gutenberg-2022 perplexity improves by -1.38∼-1.62 across length splits.With 2k demonstration examples in memory, LONGMEM also shows pronounced in-context learning improvements over MemTRM and non-memory-augmented baselines.
2 Methods
LONGMEM combines a frozen backbone LLM with a trainable residual SideNet and cached memory bank to retrieve and fuse long-form past context. Its memory module stores previous key-value pairs, retrieves relevant chunks, and supports efficient memory-augmented language modeling.
- Architecture: LONGMEM uses three components: a frozen backbone LLM, a residual SideNet, and a Cache Memory Bank.The backbone processes current and previous inputs, while SideNet performs memory-augmented processing.
- Memory Encoding and Caching: The frozen backbone encodes current and previous inputs without gradient calculation, while storing self-attention key-value pairs from previous inputs in the Cached Memory Bank.These cached pairs are later recalled as memory augmentations for future inputs.
- Memory Encoding and Caching: The Cached Memory Bank is a head-wise vector queue that retains key-value pairs from the latest M previous inputs and removes the oldest sequences when updating.This update preserves sequence-level language-modeling causality and keeps nearby previous context.
- Residual SideNet: SideNet is a smaller Transformer with normal decoder layers and one memory-augmented layer that jointly processes current hidden states and retrieved memory.For efficiency, SideNet can use fewer layers than the backbone, and its parameters are adapted while the backbone and output head remain frozen.
- Residual SideNet: Cross-network residual connections fuse backbone representations into SideNet by adding hidden-state differences from corresponding backbone layers.This design transfers pretrained backbone knowledge while adapting the lightweight SideNet.
- Memory Retrieval and Fusion: Token-to-chunk retrieval mean-pools contiguous memory chunks, selects the top-(K/csz) chunks by query-key dot product, and flattens them into K token-level key-value pairs.Chunk retrieval reduces the retrieval index size and accelerates retrieval; chunk size controls context granularity.
3 Experiments
Experiments evaluate LONGMEM across long-context language modeling, memory-augmented in-context learning, and ablations of memory configuration. The studies examine cached long-context usage, training organization, and retrieval granularity and capacity.
- Experimental Scope: LONGMEM is evaluated on long-text language modeling and infinite-length in-context learning by loading past contexts or demonstrations into cached memory.The evaluation covers long-context modeling and memory-augmented learning settings.
- Training Setup: Training batch construction preserves document-level causality by distributing consecutive segments from each document across consecutive batches.Documents are grouped, shuffled within groups, concatenated, and truncated into ordered segments before batch construction.
- Implementation: The implementation uses a GPT-2 407M-parameter backbone with ALiBi position embeddings and GPU-based FAISS retrieval over mean-pooled chunk keys.Each GPU constructs and updates its own retrieval module; the cached memory holds 65,536 token key-value pairs per GPU.
- Long-Context Language Modeling: On long-context benchmarks, LONGMEM significantly outperforms baselines, including -1.38 to -1.62 perplexity on PG-22 and -1.0 PPL on ArXiv.The evaluation includes PG-22, ArXiv, and ChapterBreak, with token-level perplexity reported for language modeling.
- Ablation Studies: A chunk size of 2 performs best on five NLU datasets, while a 16k memory size yields the best perplexity for books averaging 8k–50k tokens.The ablations vary chunk size over {2, 4, 8} and inference memory size over {8k, 16k, 32k, 65k}.
4 Related Work
Prior work extends transformer context through larger inputs, sparse attention, and side networks, while LONGMEM introduces a distinct residual fusion design for decoupled memory augmentation.
- Large Language Models: Large language models exhibit emergent abilities including few-shot in-context learning and multi-step reasoning.
- x-formers: Transformer variants reduce self-attention complexity or reuse past attention states to support longer contexts.
- Side-Tuning: Side-Tuning trains a lightweight side-network fused with a fixed pre-trained network, whereas LONGMEM uses a different objective and cross-network fusion.
5 Conclusion
LONGMEM augments frozen language models with decoupled long-term memory, enabling retrieval of past context for long-context modeling and memory-augmented in-context learning.
- LONGMEM caches attention key-value pairs from past inputs and retrieves them through a residual SideNet while the backbone LLM remains frozen.
- The method improves long-context language modeling and enables thousands of auxiliary demonstration examples to be cached for in-context learning.
A Inference Efficiency and GPU-Memory Efficiency
LONGMEM loads out-of-boundary inputs into cached memory for long-sequence comprehension, with inference speed and GPU-memory utilization compared against dense self-attention.
- LONGMEM loads out-of-boundary inputs into cached memory as previous context when comprehending long sequences.
- Table 6 compares LONGMEM with fully dense self-attention using inference speed and GPU-memory utilization.
B Training Details
The training setup reproduces a GPT-2 backbone and adapts memory augmentation on long text using fixed hardware and specified optimization settings.
- The reproduced GPT-2* pre-training uses 117B tokens, batch size 512, and 1024-token fixed segments.
- Memory-augmented adaptation training uses Adam on 16 32GB-Tesla-V100 GPUs.
- Table 7 reports the memory-augmented adaptation and architectural hyperparameters.
C Prompting Templates
This section presents the hand-crafted prompting templates and labels used for zero-shot evaluation across five NLU datasets and the Squad QA dataset.
- Table 8 presents hand-crafted in-context learning prompting templates and labels for five NLU datasets and the Squad QA dataset.
- The prompts are used to query model predictions during zero-shot evaluation.