PaperlayerView exact revision on arXiv

Source-linked AI summary

In-context Autoencoder for Context Compression in a Large Language Model

Tao Ge, Jing Hu, Lei Wang, Xun Wang, Si-Qing Chen, Furu Wei

arXiv:2307.06945v4cs.CLcs.AIcs.LG

TL;DR

Long-context modeling remains challenging, motivating a more compact representation that preserves context information. ICAE uses a LoRA-adapted LLM encoder, the target LLM as decoder, and autoencoding, continuation, and instruction fine-tuning to produce usable memory slots. On Llama, it achieves 4× context compression with about 1% additional parameters, while experiments were limited to models up to 13 billion parameters.

  • Problem

    Long-context modeling is challenging for Transformer LLMs, and context compression asks whether the same information can be represented with a shorter context.

  • Method

    ICAE encodes contexts into compact memory slots with a LoRA-adapted LLM, then conditions the target LLM on them after autoencoding, continuation, and instruction training.

  • Results

    4× context compression is achieved on Llama, with the resulting ICAE adding about 1% learnable parameters to the target LLM.

  • Takeaways & Limitations

    ICAE represents the same content with a shorter context or more information at the same context length, reducing inference computation and memory overheads in practical scenarios.

  • Takeaways & Limitations

    Experiments were conducted only on Llama models up to 13 billion parameters because of computational limitations.

Abstract

from arXiv · show

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\times$ 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.

1 INTRODUCTION

Long-context modeling is difficult for Transformer LLMs, motivating context compression as a way to represent the same information more compactly. ICAE uses an LLM to encode contexts into memory slots that support responses to prompts, with reported 4× compression and potential inference savings.

  • Transformer self-attention makes long-context modeling a fundamental challenge, while prior architectural approaches can reduce computation and memory but still suffer performance declines on long contexts.
  • Context compression seeks a shorter representation that preserves the information needed for the same response, moving from 2,572 characters to 512 words and potentially fewer memory slots.
  • ICAE encodes a long context into a small number of memory slots and lets the target LLM condition on them to accomplish varied prompted tasks.
  • 4× context compression is achieved by Llama-based ICAE after autoencoding and language-modeling pretraining followed by instruction fine-tuning.
  • ICAE is presented as a context-management approach that can improve inference latency and memory cost while also enabling investigation of LLM memorization.

2 IN-CONTEXT AUTOENCODER

ICAE combines a lightweight LoRA-adapted encoder with the unchanged target LLM as decoder, producing memory slots from contexts. It trains these slots through autoencoding, continuation, and instruction-response objectives.

  • 2.1 MODEL ARCHITECTURE: ICAE uses an encoder-decoder structure in which a LoRA-adapted LLM encodes contexts and the target LLM decodes from the resulting memory slots.
  • 2.1 MODEL ARCHITECTURE: The encoder appends k memory tokens to a context of length L and outputs k memory slots, with k much smaller than L.
  • 2.2 PRETRAINING: Autoencoding trains the decoder to restore the original context from the produced memory slots, using the special [AE] token to identify the task.
  • 2.2.2 TEXT CONTINUATION: Text continuation adds a self-supervised objective that predicts the continuation of the context and improves generalization beyond exclusive reliance on autoencoding.
  • 2.3 INSTRUCTION FINE-TUNING: Instruction fine-tuning uses context-prompt-response samples so the decoder can produce desirable responses from memory slots conditioned on prompts.

3 EXPERIMENTS

The experiments pretrain ICAE on the Pile and instruction-tune it with PWC, evaluating Llama targets under a default 512-token training context. The model uses 128 memory slots and adds about 1% learnable parameters.

  • Data: ICAE is pretrained on the Pile and instruction-tuned on PWC, which contains 240k training and 18k testing context-prompt-response samples.
  • Data: The default maximal training context length is 512 tokens for both the encoder and decoder.
  • Model Configuration: Llama serves as the target LLM, with LoRA applied to query and value projections in its multi-head attention.
  • Model Configuration: The default memory-slot length is 128 and the LoRA rank is 128.
  • Model Configuration: The resulting ICAE adds about 1% learnable parameters on top of the target LLM.

3.2 RESULTS

The pretrained ICAE reconstructs long contexts accurately with 128 memory slots, but performance degrades as contexts lengthen or memory slots shrink. Fine-tuned models retain useful response quality under compression, especially with stronger base models and pretraining.

  • 3.2.1 PRETRAINED ICAE: Below 0.05 loss and nearly 100% BLEU and EM occur for contexts within 300 tokens using 128 memory slots.At 500 tokens, median BLEU remains over 0.98 and median EM approaches 0.6.
  • 3.2.1 PRETRAINED ICAE: As context length exceeds 400 tokens, BLEU and EM decline, indicating insufficient capacity in the 128-length memory slots.
  • 3.2.1 PRETRAINED ICAE: Over 4× compression is challenging: at context length 500, k = 128 exceeds 95% BLEU, whereas k = 64 and 32 are much less satisfactory.The loss curve shows the same degradation at higher compression ratios.
  • 3.2.1 PRETRAINED ICAE: Restoration favors familiar text: normal passages are well restored, while random text produces much worse loss and BLEU scores.Observed restoration errors also resemble selective human memorization, with some information preserved and other details altered.
  • 3.2.2 FINE-TUNED ICAE: On PWC, Llama-7b ICAE with 128 memory slots beats Alpaca and StableLM using original contexts, with win rates of 56.7% and 74.1%, respectively.Its win+tie rate is 73%–81%, but it underperforms GPT-4 in around 70% of cases.
  • 3.2.2 FINE-TUNED ICAE: Larger base models and pretraining improve compressed response quality, while 128-length memory slots significantly outperform 128-token natural-language summaries.A pretrained ICAE with 8× compression can match a non-pretrained ICAE with 4× compression, and memory slots achieve an approximately 2× win/lose ratio over summaries.

3.3 ANALYSIS

The analysis examines how target-model strength, compression ratio, and multiple memory-slot spans affect ICAE quality and inference efficiency. Stronger target LLMs improve compression, while 4× compression can substantially reduce latency and GPU memory with minimal quality degradation.

  • Target LLMs: More powerful target LLMs achieve better context compression ratios.
  • Inference efficiency: Over 2× speedup is achieved with ICAE’s 4× context compression during inference.Acceleration reaches around 3.5× in compute-intensive scenarios and may exceed 7× when compressed memory slots are cached.
  • Multiple spans: Multiple context chunks can be compressed individually and concatenated, but this pattern initially fails because it was absent during training.A small amount of multiple-span concatenation training data can be incorporated to address this issue.
  • Multiple spans: 2048 memory slots perform on par with 4096-token contexts, saving about 20GB of GPU memory with minimal quality degradation.

4 RELATED WORK

Related work addresses prompt compression, compressed representations, and general compression, but the cited approaches differ from ICAE in computational cost, context length, model compatibility, or training requirements.

  • Prompt compression: Prompt-compression methods learn compact soft prompts but require back-propagation for each incoming prompt, limiting their application.
  • Prompt compression: GIST produces gist tokens similarly to ICAE memory slots but is limited to compressing short prompts.Its prompts are usually short because they refer to task instructions before input texts.
  • Prompt compression: GIST requires fine-tuning the LLM, and its gist tokens appear incompatible with an untouched LLM.
  • Context distillation: AutoCompressors recursively compress long text into summary vectors, requiring a tuned LLM and sophisticated recursive-compression training.
  • Vector-token compression: Other related studies compress indescribable concepts into vector tokens, including vision objects and text styles.
  • General compression: Broader compression research studies kNN prediction with general-purpose compressors and the compression abilities of LLMs as versatile predictors.

5 CONCLUSION AND FUTURE WORK

The conclusion presents ICAE as a compact memory-slot approach for reducing context overhead while preserving information, with implications for long-context handling and LLM memory research. Future work is constrained by experiments on Llama models up to 13 billion parameters and extends toward larger, multimodal, and discrete-memory systems.

  • Conclusion: ICAE generates compact, informative memory slots that can represent the same content with a shorter context.The approach is described as reducing computation and memory overhead for inference in scenarios such as retrieval-augmented generation and advanced prompting.
  • Conclusion: ICAE offers a perspective on LLM memorization and the connection between LLM and human memory.
  • Future work: Experiments were limited by computational resources to Llama models with at most 13 billion parameters.The authors plan to validate ICAE on larger and stronger LLMs.
  • Future work: Future directions include multimodal LLM applications and discrete memory slots for compact cross-modal representations.

A MODEL TRAINING CONFIGURATION

The training configuration uses text-continuation pretraining and instruction fine-tuning, with implementation details including hardware and numerical precision. Figures depict the continuation and instruction-training pipelines.

  • Training objectives: Pretraining uses a text-continuation objective to predict next tokens.
  • Implementation: ICAE is trained on 8 Nvidia A100 GPUs with 80GB each and uses bf16 by default.
  • Training pipeline: Figure 7 illustrates pretraining with text continuation, while Figure 8 illustrates instruction fine-tuning.Instruction fine-tuning makes generated memory slots interact with prompts and responses.

B PROFILING SETUP

Latency is tested on one Nvidia A100 GPU under a specified CPU, memory, and software configuration.

  • Latency testing uses 1 Nvidia A100 GPU with 80GB memory.The test machine also has an AMD EPYC 7413 CPU with 24 cores and 216GB RAM.

C PROMPT-WITH-CONTEXT DATASET

The PWC dataset pairs contexts with GPT-4-generated prompts and answers for instruction fine-tuning and evaluation. It includes 240k training examples and 18k test examples.

  • The PWC dataset contains text, prompt, and answer triples.It is constructed by sampling texts from the Pile and generating prompts and answers with GPT-4.
  • GPT-4 generates 15 prompts per text, including 10 specific and 5 general prompts.The prompts cover diverse aspects such as topic, genre, structure, style, polarity, key information, and details.
  • The dataset uses GPT-4 to generate prompt-answer pairs from sampled texts.The construction process samples 20k texts from the Pile before generating the pairs.
  • 240k examples are used for training and 18k examples for testing.The test set includes contexts with varied length distributions.

D GPT-4 EVALUATION

GPT-4 evaluates two assistant responses by judging instruction following and correctness, using examples and a structured choice among the alternatives or a tie. The supplied examples illustrate this procedure across factual and instruction-following tasks.

  • GPT-4 compares two model-generated answers and selects the better response or a tie.The evaluation prompt provides a task description and three examples covering choices for Assistant A, Assistant B, and ties.
  • The evaluation prioritizes whether responses follow the instruction and are correct.If both answers correctly respond to the prompt, GPT-4 should judge them as a tie.
  • The GPT-4 rater tends to prefer longer responses, potentially making ICAE’s reported evaluation performance lower than its actual performance.The paper attributes this concern to ICAE’s generally short responses after instruction fine-tuning with the PWC dataset.
  • The prompt requires GPT-4 to return a structured JSON object containing a reason and a choice.The examples specify the expected keys and instruct the evaluator to keep its reasoning brief.
  • GPT-4 is shown comparing a pretrained and non-pretrained ICAE in example evaluations.Table 9 presents outputs from the target LLM conditioned on 128 memory slots produced by both variants.
  • GPT-4 identifies the answer that correctly states a maximum prison sentence when comparing two responses.This example demonstrates factual correctness as an evaluation criterion.
  • GPT-4 favors the response that correctly identifies Ranga Rangachari as Red Hat’s vice president of Storage and Hyperconverged Infrastructure.The alternative incorrectly identifies him as Red Hat’s CEO.
Loading 2307.06945v4…