CAMELoT: Towards Large Language Models with Training-Free Consolidated Associative Memory

Zexue He,Leonid Karlinsky,Donghyun Kim,Julian McAuley,Dmitry Krotov,Rogerio Feris

2024-02-21

Abstract:Large Language Models (LLMs) struggle to handle long input sequences due to high memory and runtime costs. Memory-augmented models have emerged as a promising solution to this problem, but current methods are hindered by limited memory capacity and require costly re-training to integrate with a new LLM. In this work, we introduce an associative memory module which can be coupled to any pre-trained (frozen) attention-based LLM without re-training, enabling it to handle arbitrarily long input sequences. Unlike previous methods, our associative memory module consolidates representations of individual tokens into a non-parametric distribution model, dynamically managed by properly balancing the novelty and recency of the incoming data. By retrieving information from this consolidated associative memory, the base LLM can achieve significant (up to 29.7% on Arxiv) perplexity reduction in long-context modeling compared to other baselines evaluated on standard benchmarks. This architecture, which we call CAMELoT (Consolidated Associative Memory Enhanced Long Transformer), demonstrates superior performance even with a tiny context window of 128 tokens, and also enables improved in-context learning with a much larger set of demonstrations.

Computation and Language

What problem does this paper attempt to address?

This paper aims to address the high memory and runtime costs faced by large language models (LLMs) when processing long input sequences. The current methods have limitations due to limited memory capacity and the need for retraining to incorporate new LLMs. The paper proposes a new approach called CAMEL OT (Consolidated Associative Memory Enhanced Long Transformer), which is a trainable associative memory module that can be attached to any pre-trained (frozen) attention-based LLM, enabling it to handle input sequences of arbitrary length. This module dynamically integrates the representation of individual tokens based on novelty and recency balance, forming a non-parametric distribution model for modeling long-term context. By retrieving information from and updating the memory bank, LLMs can achieve significant performance improvements in long contexts, such as reducing perplexity in language modeling tasks. This approach does not require additional training, fine-tuning, or adapter learning between the base LLM and the AM module, making it easier to integrate into existing LLM architectures.

CAMELoT: Towards Large Language Models with Training-Free Consolidated Associative Memory

Augmenting Language Models with Long-Term Memory

UniMem: Towards a Unified View of Long-Context Large Language Models

Needle in the Haystack for Memory Based Large Language Models

MemLong: Memory-Augmented Retrieval for Long Text Modeling

Unleashing Infinite-Length Input Capacity for Large-scale Language Models with Self-Controlled Memory System.

CMT: A Memory Compression Method for Continual Knowledge Learning of Large Language Models

InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory

MEMORYLLM: Towards Self-Updatable Large Language Models

HMT: Hierarchical Memory Transformer for Long Context Language Processing

E2LLM: Encoder Elongated Large Language Models for Long-Context Understanding and Reasoning

Enhancing Large Language Model with Self-Controlled Memory Framework

Large Language Models with Controllable Working Memory

Core Context Aware Attention for Long Context Language Modeling

Training-Free Long-Context Scaling of Large Language Models

Efficiently Training 7B LLM with 1 Million Sequence Length on 8 GPUs

LoMA: Lossless Compressed Memory Attention

Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs

LaMemo: Language Modeling with Look-Ahead Memory

Long-context LLMs Struggle with Long In-context Learning