[Paper Review] In-context Autoencoder for Context Compression in a Large Language Model

Tao Ge, Jing Hu+4 more|arXiv (Cornell University)|Jul 13, 2023

Speech Recognition and Synthesis7 citations

TL;DR

This paper proposes In-context Autoencoder (ICAE), a lightweight method that uses a LoRA-fine-tuned LLM to compress long contexts into compact memory slots, enabling efficient inference with up to 4× context compression. The approach achieves significant latency and memory savings while preserving performance, offering new insights into LLM memorization and context management.

ABSTRACT

We propose the In-context Autoencoder (ICAE), leveraging the power of a large language model (LLM) to compress a long context into short compact memory slots that can be directly conditioned on by the LLM for various purposes. ICAE is first pretrained using both autoencoding and language modeling objectives on massive text data, enabling it to generate memory slots that accurately and comprehensively represent the original context. Then, it is fine-tuned on instruction data for producing desirable responses to various prompts. Experiments demonstrate that our lightweight ICAE, introducing about 1% additional parameters, effectively achieves $4 imes$ context compression based on Llama, offering advantages in both improved latency and GPU memory cost during inference, and showing an interesting insight in memorization as well as potential for scalability. These promising results imply a novel perspective on the connection between working memory in cognitive science and representation learning in LLMs, revealing ICAE's significant implications in addressing the long context problem and suggesting further research in LLM context management. Our data, code and models are available at https://github.com/getao/icae.

Motivation & Objective

• To address the long context problem in LLMs by compressing input sequences into shorter, informative memory representations.
• To enable efficient inference with reduced latency and GPU memory cost by replacing long contexts with compact memory slots.
• To explore how self-supervised pretraining enhances LLM memorization capabilities, drawing parallels to human working memory.
• To develop a parameter-efficient, plug-and-play solution compatible with existing LLMs without architectural changes.
• To investigate the potential of memory slots for downstream tasks like chain-of-thought reasoning and task-specific compression.

Proposed method

• ICAE uses a LoRA-adapted LLM as a learnable encoder to compress long input contexts into a small number of memory slots.
• The decoder is the target LLM itself, which conditions on the generated memory slots to generate responses or reconstruct the original context.
• The model is pretrained jointly on autoencoding and language modeling objectives to ensure memory slots are both reconstructive and semantically meaningful.
• Fine-tuning on instruction data improves the quality of memory slots for diverse prompting scenarios and enhances interaction with prompts.
• Memory slots are learned via a special token format and learnable embeddings, enabling the model to generate compact, high-fidelity representations.
• The method is parameter-efficient, adding fewer than 1% of the original model parameters via LoRA.

Experimental results

Research questions

• RQ1 Can a large language model be effectively used to compress long input contexts into a small set of memory slots without architectural changes?
• RQ2 To what extent can such memory slots preserve the semantic and structural information of the original context?
• RQ3 How does self-supervised pretraining affect the model’s ability to memorize and reconstruct long contexts?
• RQ4 What is the trade-off between context compression ratio and inference performance in real-world scenarios?
• RQ5 Can memory slots be used effectively across diverse downstream tasks, including chain-of-thought reasoning and task-specific prompting?

Key findings

• ICAE achieves up to 4× context compression on Llama models, significantly reducing GPU memory usage and inference latency.
• Using 2048 memory slots instead of 4096 context tokens results in minimal performance degradation, saving approximately 20GB of GPU memory.
• The pretraining phase with autoencoding and language modeling objectives greatly enhances the model’s ability to memorize and reconstruct long contexts.
• The method demonstrates strong compatibility with existing LLMs, requiring only 1% additional parameters via LoRA fine-tuning.
• Memory slots are effective for chain-of-thought reasoning and can be further fine-tuned for task-specific compression.
• The study reveals a strong analogy between LLM memorization and human working memory, suggesting shared mechanisms in representation learning.