Source-linked AI summary
InfLLM: Training-Free Long-Context Extrapolation for LLMs with an Efficient Context Memory
Chaojun Xiao, Pengle Zhang, Xu Han, Guangxuan Xiao, Yankai Lin, Zhengyan Zhang, Zhiyuan Liu, Maosong Sun
TL;DR
LLMs trained on short sequences face out-of-domain and distraction problems on longer inputs, while continual long-sequence training is costly and may alter capabilities. InfLLM adds a training-free memory that retrieves relevant distant contexts alongside local context, enabling comparable performance to continually trained baselines and long-distance dependency capture at 1,024K tokens. The method remains limited by CPU memory usage, inference speed, and heuristic context segmentation.
Problem
LLMs pretrained on restricted sequence lengths struggle with longer inputs because of out-of-domain and distraction issues, while continual long-sequence training is costly and may change model capabilities.
Method
InfLLM combines sliding-window attention with an external block-level context memory that retrieves relevant distant contexts for each computation step without additional training.
Results
InfLLM achieves comparable performance to continually trained baselines without additional training and effectively captures long-distance dependencies in sequences up to 1,024K tokens.
Takeaways & Limitations
InfLLM extends the use of short-sequence-pretrained LLMs to extremely long streaming inputs while using a finite-size context window.
Takeaways & Limitations
InfLLM stores substantial past KV cache in CPU memory, still has room for inference speed-up, and uses heuristic context segmentation whose dynamic improvement remains future work.
Abstract
from arXiv · showhide
Large language models (LLMs) have emerged as a cornerstone in real-world applications with lengthy streaming inputs (e.g., LLM-driven agents). However, existing LLMs, pre-trained on sequences with a restricted maximum length, cannot process longer sequences due to the out-of-domain and distraction issues. Common solutions often involve continual pre-training on longer sequences, which will introduce expensive computational overhead and uncontrollable change in model capabilities. In this paper, we unveil the intrinsic capacity of LLMs for understanding extremely long sequences without any fine-tuning. To this end, we introduce a training-free memory-based method, InfLLM. Specifically, InfLLM stores distant contexts into additional memory units and employs an efficient mechanism to lookup token-relevant units for attention computation. Thereby, InfLLM allows LLMs to efficiently process long sequences with a limited context window and well capture long-distance dependencies. Without any training, InfLLM enables LLMs that are pre-trained on sequences consisting of a few thousand tokens to achieve comparable performance with competitive baselines that continually train these LLMs on long sequences. Even when the sequence length is scaled to $1,024$K, InfLLM still effectively captures long-distance dependencies. Our code can be found in \url{https://github.com/thunlp/InfLLM}.
1 Introduction
LLMs trained on short sequences struggle with longer inputs because of unseen lengths and distracting noisy contexts, while continual long-sequence training is costly and may weaken short-context performance. InfLLM addresses this with a training-free external memory that retrieves relevant distant contexts, achieving comparable performance to continually trained baselines and retaining long-distance dependencies at 1,024K tokens.
- LLMs pretrained on only a few thousand tokens cannot reliably process longer sequences because of unseen lengths and distracting noisy contexts.
- Continual training on longer sequences incurs substantial costs, requires high-quality long-sequence datasets, and may weaken short-context performance.
- InfLLM combines sliding-window attention with an external memory that retrieves relevant distant contexts while ignoring irrelevant noise.
- InfLLM organizes past key-value vectors into blocks and selects semantically significant tokens as memory-unit representations for relevance lookup.
- InfLLM achieves comparable performance to continually trained baselines without additional training and captures long-distance dependencies at 1,024K tokens.
2 Related Work
Prior approaches address long-context processing through positional extrapolation or more efficient attention computation, but each leaves important constraints. Positional methods do not resolve distraction from noisy contexts, while efficient-attention methods generally require architectural modification or retraining.
- Long-context research commonly separates into context-length extrapolation and efficient context computation.
- Context Length Extrapolation: Context-length extrapolation methods modify positional encoding or reuse position indices to let short-sequence models process longer inputs.
- Context Length Extrapolation: Position downscaling and reuse methods alleviate unseen-length issues but do not address distraction from noisy contexts.
- Efficient Context Computation: Efficient-attention approaches reduce computation through sparse or approximated attention and state-space models, but require architectural changes and retraining.
- Efficient Context Computation: Key-value eviction improves inference efficiency but cannot extrapolate LLM context windows without further training because unseen positions remain problematic.
3 Methodology
InfLLM extends limited-context LLMs to long streams without training by combining sliding-window attention with a memory that retrieves relevant distant contexts. It organizes evicted key-value vectors into block-level units, selects representative tokens, and loads only the most relevant units for each computation step.
- 3.1 Overall Framework: InfLLM combines sliding-window attention with an additional context memory to capture long-distance dependencies without further training.Each token attends to local contexts and relevant contexts retrieved from memory.
- 3.1 Overall Framework: The long input is encoded chunk-by-chunk and generated token-by-token because the full sequence cannot fit in GPU memory.For encoding, the current-token length equals the chunk size; for decoding, it equals one.
- 3.1 Overall Framework: Past key-value vectors are divided into initial, evicted, and local tokens, with evicted tokens stored in context memory.The current cache concatenates initial tokens, retrieved memory units, and local tokens; empty retrieval reduces InfLLM to methods that discard distant contexts.
- 3.2 Context Memory: InfLLM partitions evicted tokens into contiguous blocks and represents each block with the tokens having the highest representative scores.Representative scores measure a token’s influence on other tokens in its local window and require no additional parameters.
- 3.2 Context Memory: Only the km memory units with the highest relevance scores are loaded for the current attention computation.This block-level lookup avoids constructing massive token-level memory units and ignores irrelevant contexts to reduce computation.
- 3.2 Context Memory: The representative-token memory representation is training-free, while learning an additional encoder for more expressive representations remains future work.The authors explicitly identify the additional encoder as an open direction.
- 3.2 Context Memory: InfLLM assigns identical positional encodings to tokens beyond the local window to address both length extrapolation and discontinuous-cache mismatch issues.Its memory system also offloads most units and, for extremely long sequences, representative tokens to CPU memory with a k-nearest-neighbor index.
4 Experiments
InfLLM is evaluated as a training-free approach for extending short-context LLMs to long sequences, with experiments covering benchmarks, efficiency, scaling, memory design, and ablations. Across these studies, it retrieves distant context while limiting computation, but heuristic memory segmentation remains a stated limitation.
- 4.1 Settings: InfLLM is evaluated on ∞-Bench using Mistral-7B-Instruct-v0.2 and Llama-3-8B-Instruct, whose maximum pre-trained sequence lengths are 32K and 8K.The benchmark’s average sequence length is 145.1K, and its 95% quantile is 214K.
- 4.3 Main Results: InfLLM significantly outperforms sliding-window methods and can generalize Llama-3 from 8K to more than 16 times its length on ∞-Bench.The paper attributes this to context memory supplying relevant distant contextual information, whereas position downscaling methods remain affected by noisy contexts.
- 4.3 Main Results: InfLLM achieves comparable performance to the continually fine-tuned Llama-1M while using less computation and memory.Llama-1M extends its context window to 1048K through further fine-tuning on long-text and chat data.
- 4.1 Settings: InfLLM supplements a 4K local window with relevant memory units, using block-level memory and offloading to process long sequences within limited resources.The reported configurations use memory unit size 128, four representative tokens, and model-specific numbers of loaded memory units.
- 4.6 Extra Studies: Increasing representative tokens or selected memory units generally improves performance, while the optimal memory unit size varies across tasks.Larger unit counts improve recall but increase memory scheduling time; heuristic segmentation can be suboptimal because semantic unit sizes differ by task.
- 4.7 Ablation Study: Ablations associate reduced memory lookup with significant performance decline, while mean key-vector representations remain competitive with InfLLM’s representative-token approach.The results support memory lookup during both input encoding and output decoding, and motivate more efficient unit representations.
- 4.8 Scaling to 1,024K Context: InfLLM reaches 100% accuracy when sequence length scales to 1,024K, whereas LM-Infinite’s performance rapidly declines as sequence length increases.The result is presented as evidence that InfLLM captures long-distance dependencies despite length noise.
5 Conclusion
InfLLM is a training-free method that extends LLM length generalizability by adding context memory to sliding window attention. Experiments show it improves processing of extremely long sequences and captures long-distance dependencies.
- InfLLM improves LLM length generalizability without additional training.It is based on sliding window attention and an additional context memory module.
- The context memory selects relevant information from massive contexts to capture long-distance dependencies.
- Experiments on two widely-used long-text benchmarks show that InfLLM helps models trained on few-thousand-token sequences process extremely long sequences.
Limitations
InfLLM has memory and speed limitations despite extending LLM context windows without additional training.
- InfLLM stores substantial past KV cache in CPU memory, increasing CPU memory usage.KV-cache quantization is proposed as a future way to reduce this requirement.
- InfLLM reduces long-text computational overhead but still has room for inference-speed improvements.The authors suggest integration with llama.cpp and vllm as future directions.
A Cache Management Strategy
InfLLM manages offloaded memory units through a fixed GPU cache, relevance-based loading, and LRU eviction. The LRU strategy has a lower cache missing rate than Random and FIFO on GovReport data.
- InfLLM offloads most memory units to CPU memory while retaining frequently used and currently needed units in a fixed-size GPU cache.This reduces GPU memory use and communication between CPU and GPU.
- At each computation step, InfLLM computes memory-unit relevance and transfers needed units from CPU memory to GPU memory when they miss the cache.
- After attention computation, LRU assigns frequency scores and offloads the lowest-scoring GPU-cache units back to CPU memory.The frequency score incorporates current attention scores and decayed influence from previous lookups.
- On a GovReport sample, LRU has a lower cache missing rate than Random and FIFO.The reported lower missing rate is intended to ensure that offloading does not introduce significant time overhead.
B Positional Encoding
InfLLM assigns identical positional encodings to tokens beyond the local window, relying on decoder-only causality to preserve relative positional information. A multi-passkey evaluation tests whether this information remains usable.
- InfLLM gives all tokens beyond the local window the same positional encoding rather than explicitly encoding their positions.
- Decoder-only unidirectionality lets later hidden states incorporate earlier context, preserving relative positional information between distant tokens.A later segment can use an earlier segment, whereas the earlier segment cannot access subsequent tokens.
- The Retrieve.Passkey evaluation uses two randomly positioned passkeys in 50 sequences of length 64K and requires outputting them in order.This setup is used to verify the model’s ability to capture relative positional information.
C.1 Implementation Details
LongBench results show InfLLM outperforming other streaming-input models across diverse tasks, with context memory improving model performance. Experiments used specified cache, precision, hardware, and scaling settings.
- Implementation Details: The context memory is constructed for all LLM layers.
- Implementation Details: The GPU cache size is 32, twice the number of loaded units per step.
- Implementation Details: Experiments use half-float precision and NVIDIA A100 or A800 GPUs.
- Implementation Details: For 1,024K-context experiments, InfLLM uses 2048-token encoding chunks and one representative token.
- Evaluation Results: InfLLM outperforms other models capable of processing streaming inputs across diverse LongBench tasks.The reported improvement is attributed to context information supplied by the context memory.
C.3 Experiments on Vicuna
InfLLM extends Vicuna's 4K context length to 128K, improving some retrieval tasks but not more complex retrieval and mathematical tasks. The results are presented for Vicuna-based models in Table 6.
- Context Extension: InfLLM extends Vicuna's 4K context length to 128K.Vicuna has a maximum length of only 4K, while the extended setting reaches 128K.
- Results: InfLLM significantly improves Vicuna performance on Retrieve.Passkey and Retrieve.Number.
- Results: InfLLM shows no performance gains on Retrieve.KV and Math.Find with Vicuna.The passage attributes this to Vicuna hidden vectors having limited ability to filter noise in extremely long texts.
- Results: Table 6 reports results for Vicuna-based models.