APCE: Adaptive Progressive Context Expansion for Long Context Processing

TL;DR

The paper tackles the memory and performance degradation of long-context transformers (ContextRot) by introducing Adaptive Progressive Context Expansion (APCE), an input-chunk sparsification approach that selects the most semantically relevant chunks via a query-chunk similarity. APCE computes low-dimensional embeddings for input chunks and the current query, uses cosine similarity to pick the top-k chunks for attention, and supports reprioritization and asynchronous generation to adapt as decoding progresses, decoupling from hardware-specific kernels. This yields a drastic reduction in attention complexity from $O(N^2)$ to $O((k m)^2)$ and substantial memory savings (e.g., up to 55.6% in prefill and 32.8% in KV-cache) while matching or surpassing full-dense baselines on long-context summarization tasks like BookSum with context lengths up to 30k. The method is hardware-agnostic, complementary to existing sparsification techniques, and opens avenues for applying context-aware chunk selection to a broader set of long-context tasks, with tunable trade-offs via reprioritization and asynchronous generation.

Abstract

Deploying useful Long-Context Transformer Models (LCTMs) requires addressing two key challenges: (1) A growing memory footprint due to quadratic self-attention and linear KV-cache scaling in memory as sequence length increases; (2) the ContextRot phenomena where empirical evidence suggests that transformer architecture's performance degrades with increasing context length. Given the shared dependency on the input, a natural question arises: Can we surgically select the most important input chunks for processing to synergistically (a) reduce the memory footprint, and (b) mitigate the ContextRot effects? In this paper, we answer this question in the affirmative for long-context summarization tasks. We propose APCE as a context-aware solution to select the most important input chunks through low-dimensional semantic similarity matching with the current query. By directly operating on the input, APCE decouples from strict dependency on underlying hardware or CUDA environments, promising a compatible solution scalable to different deployment systems. Our empirical evaluations have demonstrated superior or on-par summarization performance for APCE compared to the full dense baseline using a fraction (50%-70%) of the input sequence resulting in KV-cache and self-attention memory efficiency improvements. We hope our findings inspire further research on context-aware efficiency solutions for LCTMs geared towards other relevant long-context tasks.

Figures (4)

• Figure 1: Overview of APCE for an example with 6 chunks. (left) The full dense baseline attention suffers from quadratic complexity and long time-to-first-token (TTFT). (middle) APCE performs prefill chunk selection using semantic similarity, resulting in sparse attention reducing complexity and TTFT. (right) Chunk selection is updated as generation progresses to maintain performance.
• Figure 2: Summarization Performance scores (left) and inference timings (right) for APCE compared to the baseline when varying the number of loaded chunks on Llama-3.2-3B-Instruct dubey2024llama the 30K context length group from BookSum kryscinski2021booksum. Chunk size is fixed to 800 tokens.
• Figure 3: Comparing APCE and the Full Dense baseline across the 3 groups of long context when varying the number of chunks at fixed chunk size of 800. The top row shows the performance results on BookSum kryscinski2021booksum, whereas the bottom row shows the timing evaluations.
• Figure 4: Comparing APCE and the Full Dense baseline across the 3 groups of long context when varying the chunk size at 70% of available chunks selected. The top row shows the performance results on BookSum kryscinski2021booksum, whereas the bottom row shows the timing evaluations.