UT-ACA: Uncertainty-Triggered Adaptive Context Allocation for Long-Context Inference

TL;DR

UT-ACA introduces a dynamic context allocation method for long-context inference in language models, reducing context usage while maintaining quality by adjusting based on token uncertainty during decoding.

Contribution

It proposes a novel uncertainty-triggered adaptive framework that dynamically adjusts context size during inference, addressing fixed context limitations in existing methods.

Findings

Reduces average context usage significantly.

Maintains generation quality in long-context scenarios.

Demonstrates effectiveness across multiple benchmarks.

Abstract

Long-context inference remains challenging for large language models due to attention dilution and out-of-distribution degradation. Context selection mitigates this limitation by attending to a subset of key-value cache entries, yet most methods allocate a fixed context budget throughout decoding despite highly non-uniform token-level contextual demands. To address this issue, we propose Uncertainty-Triggered Adaptive Context Allocation (UT-ACA), an inference-time framework that dynamically adjusts the context window based on token-wise uncertainty. UT-ACA learns an uncertainty detector that combines semantic embeddings with logit-based confidence while accounting for uncertainty accumulation across decoding steps. When insufficient evidence is indicated, UT-ACA selectively rolls back, expands the context window, and regenerates the token with additional support. Experiments show that UT-ACA substantially reduces average context usage while preserving generation quality in long-context settings.