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

Category: Machine Learning Systems
InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory

What is the Problem?

Large language models (LLMs) are typically pre-trained on sequences with limited maximum lengths (a few thousand tokens), which restricts their ability to process much longer sequences required in real-world applications such as LLM-driven agents and streaming inputs. Existing solutions often require expensive continual pre-training on longer sequences, which is computationally intensive and can degrade performance on shorter contexts. The challenge is to enable LLMs to efficiently and effectively process extremely long sequences—well beyond their training context window—without any additional training or architectural changes.

Summary

The paper introduces InfLLM, a training-free, memory-based method that enables LLMs to process extremely long sequences by augmenting the standard sliding window attention mechanism with an efficient external context memory. InfLLM stores distant context information in memory units and dynamically retrieves only the most relevant units for each token during attention computation. This approach allows LLMs to capture long-distance dependencies and avoid the distraction caused by irrelevant or noisy contexts, all without any further training. The method is evaluated on challenging long-context benchmarks, demonstrating that LLMs pre-trained on short sequences can achieve performance comparable to or better than models that have undergone costly continual training on long sequences.

Key Insights

  • LLMs possess an intrinsic, underutilized capacity to process and reason over extremely long sequences if provided with a mechanism to efficiently retrieve relevant distant context, even without further training.
  • By combining sliding window attention with a block-level, dynamically managed context memory, InfLLM enables efficient long-context extrapolation, maintaining both computational efficiency and high performance on long-sequence tasks.

Notable Design Details/Strengths

  • Training-Free and Model-Agnostic: InfLLM requires no additional training or model architecture changes, making it applicable to any existing LLM.
  • Block-Level Context Memory: Past key-value vectors are organized into blocks, and only the most semantically significant tokens within each block are used for relevance computation, reducing computational and memory overhead.
  • Dynamic Memory Management: An offloading mechanism stores most memory units in CPU memory, with frequently accessed units cached on GPU using an LRU strategy, enabling processing of sequences up to 1M tokens with modest GPU resources.
  • Effective Long-Range Dependency Modeling: InfLLM consistently outperforms or matches both training-free and continually trained baselines on long-context benchmarks, demonstrating robust long-range reasoning.

Limitations/Weaknesses

  • Increased CPU Memory Usage: Storing large amounts of past key-value cache in CPU memory can be demanding, especially for extremely long sequences.
  • Inference Speed: While GPU memory usage is optimized, there is still room to further accelerate inference, particularly for very long sequences.
  • Dependence on Base Model Quality: The effectiveness of InfLLM can be limited by the base model’s ability to filter noise and represent context, as seen in some tasks with weaker base models.

Summary of Key Results

  • Comparable or Superior Performance: InfLLM enables LLMs pre-trained on short sequences (e.g., 8K or 32K tokens) to match or exceed the performance of models continually trained on long sequences, across diverse tasks such as question answering, summarization, and retrieval.
  • Scalability: InfLLM successfully processes sequences up to 1,024K tokens, maintaining high accuracy in tasks requiring retrieval of information from distant context.
  • Efficiency: Achieves significant reductions in GPU memory usage and inference time compared to full-attention or continually trained long-context models, making long-context inference practical on a single GPU.
  • Generalization: Outperforms retrieval-augmented generation (RAG) approaches on context retrieval tasks, without requiring additional data or retriever training.

Open Questions

  • How can the memory unit segmentation and representative token selection be further optimized, possibly with lightweight training, to improve relevance and efficiency?
  • Can InfLLM be effectively combined with other context compression or retrieval techniques to further reduce memory and computational requirements, especially for deployment in resource-constrained environments?: inf-llm.pdf: https://github.com/thunlp/InfLLM
Tags: llm long-context memory research