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

Lang Zhou; Shuxuan Li; Zhuohao Li; Shi Liu; Wei-Shi Zheng; Zhilin Zhao

cs.CLMar 19, 2026

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

Authors: Lang Zhou, Shuxuan Li, Zhuohao Li, Shi Liu, Wei-Shi Zheng, Zhilin Zhao

Organizations: School of Computer Science and Engineering, Sun Yat-sen University · Shenzhen Loop Area Institute · Southern University of Science and Technology

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. Code and dataset are available at https://github.com/Tommy307/UT-ACA.

Figures & tables

Figure 1: Comparison of context management strategies. The figure contrasts full-context decoding, fixed-size context selection, and the proposed adaptive context allocation strategy. Full-context decoding attends to all available KV entries, fixed selection uses the same retrieved budget at every step, and UT-ACA expands the window only when uncertainty indicates that additional evidence is needed.

Figure 2: Workflow of our framework. The figure illustrates how UT-ACA connects long-context input, token-level uncertainty estimation, and adaptive context allocation during decoding. (a) The prompt provides the long-context input and target instruction. (b) The uncertainty detector converts output logits and semantic embeddings into a Generation Difficulty Metric (GDM). (c) The adaptive context window accepts low-uncertainty tokens, but rolls back, expands the retrieved context, and regenerates the token when the GDM indicates insufficient support.

Figure 3: Token-generation scenarios for the Generation Difficulty Metric. The figure categorizes grounded answers, abstentions, and hallucinations by whether the required information is supported by the context or by the intrinsic knowledge of the model.

Figure 4: Example comparison between F1 score and Conceptual Accuracy Score (CAS). F1 score gives partial credit for surface overlap, while CAS assigns zero because the prediction is semantically incorrect.

Methods	mAcc $({\%})$	Recall N $({\%})$	Recall P $({\%})$
Top-10 Std	57.43	35.66	73.19
Logit Margin	55.54	48.09	60.93
Evidential Mass	58.29	55.87	60.04
Evidential	70.85	74.45	68.25
UT-ACA	83.43	89.71	81.27

Table 1: Token-level uncertainty detection performance. Unknown and Hallucination are grouped as uncertain, while Correct is treated as certain. mAcc denotes classification accuracy; Recall N and Recall P denote recall on uncertain and certain tokens, respectively.

Methods	Settings	mTokens $\downarrow$	mAcc ${}_{\text{conc}}({\%})$ $\uparrow$
InfLLM	$K=2$	32	91.37
InfLLM	$K=3$	48	99.83
	Update: Set.1	25	96.59
UT-ACA	Update: Sub.1	29	99.08

Table 2: Generation performance on the validation set. UT-ACA is compared with fixed-budget InfLLM in terms of average context tokens (mTokens) and conceptual accuracy (mAcc

{}_{\text{conc}}

), showing that adaptive allocation preserves high accuracy with fewer tokens.

Models	Methods	Settings	mTokens $\downarrow$	mTime ${}_{\text{tok}}$ (s) $\downarrow$	mAcc ${}_{\text{conc}}$ (%) $\uparrow$
Llama-3.1-8B-it	† RetrievalAttn 33	$K_{\text{tok}}=2048$	2k	0.021	15.25
	† LLMLingua-2 36	$K_{\text{com}}=25\%$	32k	-	18.08
	† SnapKV 28	$K_{\text{tok}}=1024$	1k	0.210	23.43
	TokenSelect 45	$K_{\text{tok}}=1024$	1k	-	39.43
	OmniKV 11	$K_{\text{tok}}=6.7\%$	8k	-	40.47
	† InfLLM-V2 55	$K_{\text{tok}}=1024$	1k	-	40.29

Table 3: Context-management performance on the long-context test set, where Update: Sub.

\emph{N}

denotes the policy of decreasing the block budget by

\emph{N}

after each accepted token, with a minimum of 1. The table compares token selection, compression, fixed-budget retrieval, and UT-ACA across two backbones using context usage (mTokens), per-token latency (mTime

{}_{\text{tok}}

), and conceptual accuracy (mAcc

{}_{\text{conc}}

). Methods marked with † exclude samples with out-of-memory errors or incoherent outputs.

Method	mAcc ${}_{\text{conc}}({\%})$ $\uparrow$	mTok. $\downarrow$	E2E T. (s) $\downarrow$	GPU (GiB)
InfLLM	25.81	128	95.80	26.42
	51.23	256	97.46	26.42
	72.26	512	99.56	26.44
UT-ACA	45.23	123	98.42	26.44
	62.17	218	99.56	26.47
	70.48	344	100.70	26.51

Table 4: Quality–latency–memory trade-off on the long-context test set. InfLLM rows correspond to

K=\{8,16,32\}

and UT-ACA rows to

K_{\max}=\{32,48,64,96\}

, ordered from top to bottom. mTok., E2E T., and GPU denote selected context tokens, end-to-end latency, and peak GPU memory, respectively.

LogM	SE	LSTM	mAcc $(\%)$	mAcc ${}_{\text{conc}}(\%)$
✔	✘	✘	26.82	-
✘	✔	✘	78.95	69.58
✔	✘	✔	53.71	-
✘	✔	✔	80.06	68.58
✔	✔	✘	80.19	69.42
✔	✔	✔	83.43	70.48

Table 5: Component ablation of the uncertainty detector. The table evaluates the contribution of the Logit Margin branch (LogM), the Semantic Embedding branch (SE), and the LSTM module by removing them in different combinations. mAcc measures detector accuracy, and mAcc

{}_{\text{conc}}

measures downstream conceptual generation accuracy.

Method	$\infty$ Bench				LongBench				RULER
	Sum Task		QA Task		narrativeqa		qmsum		QA-16k		QA-8k
	mTok $\downarrow$	Score $\uparrow$	mTok $\downarrow$	Score $\uparrow$	mTok $\downarrow$	Score $\uparrow$	mTok $\downarrow$	Score $\uparrow$	mTok $\downarrow$	Score $\uparrow$	mTok $\downarrow$	Score $\uparrow$
	1024	11.25	1024	1.14	64	5.35	64	13.84	128	12.00	128	14.30
	2048	11.52	2048	0.93	128	4.68	128	14.25	256	11.00	256	16.60
H 2 O	4096	10.59	4096	1.27	256	4.44	256	14.81	512	10.70	512	21.00
	1024	27.24	1024	16.27	64	14.95	64	18.40	128	14.05	128	14.90

Table 6: Results on

\infty

Bench, LongBench, and RULER using Llama-3.1-8B-Instruct. The table compares H 2 O, InfLLM, and UT-ACA under the official evaluation protocols of the corresponding benchmarks, without retraining the UT-ACA uncertainty detector. mTok denotes the average number of selected tokens, Score follows the corresponding benchmark metric, and bold values indicate the best result in each column.

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Attribute Names	Example Data
Training Set	• Marriage Date • Job Title • Current City • Email Address • Phone Number • Favorite Color • User Agent • Credit Card Provider • Currency Used • Catch Phrase • Street Address • Vehicle License Plate • Favorite File Extension • Domain Name • Cryptocurrency • Timezone • Isbn Code • Lucky Number • Hobby • Marital Status • Personality Type	• The marriage date of Laetitia Guyot de Paris is 2010-03-29. • The job title of Franck Jourdan is Magazine features editor. • The current city of Marcial Escalona Yuste is Augerdan. • The email address of Jeremy Wright is [email protected]. • The phone number of Clarence Byrd is 07703087971. • The favorite color of Letizia Fagotto is Black. • The user agent of Inés Novoa Sarabia is Mozilla/5.0 • The credit card provider of Douglas Baker is Maestro. • The currency used of Susann Jacobi Jäckel is Bahraini dinar. • The catch phrase of Olga Torre Rodriguez is Interfaccia estesa … • The street address of Gabriela Jesus is Bährgasse 62. • The vehicle license plate of Erdal Mälzer is GE 3041 CB. • The favorite file extension of Ornella Tomasini is tiff. • The domain name of Théo Macedo is chretien.com. • The cryptocurrency of Anaïs Gilbert du Rousset is Vertcoin. • The timezone of Eutimio Tejada Ibañez is Europe/Paris. • The isbn code of Jeanne-Sabine Brunet is 978-0-456-54959-9. • The lucky number of Zoé Chauvet is 61. • The hobby of María Manuela Guardia Moles is Sleeping. • The marital status of Gustavo Vallés is Widowed. • The personality type of Salvi Toscanini-Ferrabosco is Extrovert.
Validation Set and Test Set	• Birth Date • Birth Place • Company • Major • University • Work Place	• The birth date of Olivia Garcia is March 22, 1985. • The birth place of Daniel Lee is Seoul, South Korea. • The company of Olivia Chen is Microsoft. • The major of Ethan Clark is Computer Science. • The university of Aiden Gomez is Stanford University. • The work place of Elijah Robinson is Singapore, Singapore.

Appendix

Table 7: Attribute sets and example facts in the synthetic dataset. The table lists representative attribute names and example facts for the training, validation, and test splits. Training uses diverse single-attribute facts, while validation and test focus on six held-out biography attributes.

Wikipedia Source Articles
Architecture	Mathematics	Pacific Ocean	Thermodynamics	Blockchain
Art	Number	Sahara	Virus	Cryptography
Baroque	Sport	South America	Communism	Engineering
Cinema	Vaccine	Volcano	Culture	Internal combustion engine
Classical music	French Revolution	Age of Enlightenment	Democracy	Internet
Impressionism	Maya civilization	American Civil War	Epistemology	Nanotechnology

Appendix

Table 8: Wikipedia source articles used for test-set context extension. These passages are inserted around synthetic biographies to create long distractor contexts, testing whether methods can retrieve relevant biography facts from long, heterogeneous inputs.

Figure 5: Standardized pattern of synthetic summaries. The figure illustrates the controlled person–attribute–value structure used to define ground-truth facts. This standardized structure keeps the ground truth unambiguous while allowing the surrounding biography text to vary in wording and length.

Figure 6: Prompt template for generating naturalistic biographies. The figure shows the LLM instruction used to convert structured ground-truth facts into fluent biography passages. The prompt asks the LLM to preserve the provided summary facts while embedding them in natural narrative text.

Figure 7: Prompt for reference-constrained token-level uncertainty labeling. GPT-OSS-120B assigns each prediction token as Correct (1), Hallucination (0), or Unknown (2) by comparing it with the explicit ground-truth subject, attribute, and answer under prioritized labeling rules.

Dataset	Example Summarization	Example Biography
Training Set	The email address of Jeremy Wright is [email protected].	The sky is really blue. … The email address of Jeremy Wright is [email protected]. … The sky is really blue.
Validation Set	The birth date of Olivia Garcia is March 22, 1985.	The sky is really blue. … The birth date of Olivia Garcia is March 22, 1985. … The sky is really blue.
Test Set	• The birth date of Emma Thompson is 12/29/1880. • The birth place of Emma Thompson is Tokyo, Japan. • The university of Emma Thompson is Harvard University. • The major of Emma Thompson is Computer Science. • The company of Emma Thompson is Google. • The work place of Emma Thompson is San Francisco Bay Area.	… Emma Thompson was born on December 29, 1880, in the bustling metropolis of Tokyo, Japan. … She applied to Harvard University and was accepted into the Computer Science program. … After careful consideration, she chose to join Google, a company known for its innovative culture and cutting-edge technology. … Emma Thompson’s journey from Tokyo to the San Francisco Bay Area is a testament to her unwavering determination and exceptional talent. …

Appendix

Table 9: Representative examples from the training, validation, and test splits. The table shows how the three splits differ in attribute complexity and context length. Training and validation examples contain one explicit fact in a short template-like passage, whereas test examples contain six biography attributes expressed in natural language and surrounded by long-context distractors.

Component	Parameter	Value
Input Dimensions	Vector Dim ( $D_{in}$ )	4096
	Score Dim	1
	Hidden Model Dim ( $d_{model}$ )	64
Encoders	Vec Dropout	0.5
Encoders	Score MLP Inner Dim	32
Residual Head	Number of Blocks	3

Appendix

Table 10: Hyperparameter configuration of the Llama-3.1-8B uncertainty detector. The table reports the input dimensions, encoder settings, residual-head configuration, and number of output classes. The detector projects high-dimensional LLM features into a compact 64-dimensional hidden space before classification.

Figure 8: Extraction points for semantic embeddings in the final transformer layer. The figure marks the attention output, MLP output, and residual stream in the final layer (e.g., Layer 32 of Llama-3.1-8B). These positions are evaluated as candidate internal states for token-level semantic features.

Figure 9: Prompt template for evaluating Conceptual Accuracy Score (CAS). The figure shows the LLM-judge instruction used to score semantic correctness rather than surface overlap. For each subtask, the prompt provides the question, the model prediction, and the ground-truth answer, then requests a floating-point score between 0 and 1 that reflects conceptual accuracy.

ID	Name	Str	CAS	Hum	Cor	ID	Name	Str	CAS	Hum	Cor
1	Emma Thompson	.000	.000	.000	0/6	16	Lucas Scott	.262	1.00	1.00	6/6
2	Liam Carter	.127	.167	.167	1/6	17	Harper Coleman	.114	.167	.167	1/6
3	Olivia Bennett	.230	1.00	1.00	6/6	18	Henry Green	.260	.500	.667	3/6
4	Noah Wright	.278	.830	.833	5/6	19	Evelyn Perry	.308	1.00	1.00	6/6
5	Ava Sullivan	.299	.667	.667	4/6	20	Oliver Hughes	.358	1.00	1.00	6/6
6	Ethan Brooks	.364	.833	.833	5/6	21	Abigail Kelly	.333	.833	.833	5/6

Appendix

Table 11: Preliminary human comparison of Conceptual Accuracy Score on biography extraction. The table compares string matching ( Str ), conceptual accuracy score ( CAS ), and human annotation ( Hum ) on 30 test samples. Cor reports the number of correctly extracted attributes out of six according to human annotators.

Case #3: Olivia Bennett				Case #17: Harper Coleman
(Str: 0.230, LLM: 1.000, Hum: 1.000)				(Str: 0.114, LLM: 0.167, Hum: 0.167)
Attr	GT	Pred.	Hum	Attr	GT	Pred.	Hum
Birth Date	June 3, 1888	1888-06-03	✔	Birth Date	10 Nov, 1907	1976	✘
Birth Place	Sydney, NSW	Sydney, NSW	✔	Birth Place	Mexico City	Madrid, Spain	✘
University	MIT	Mass. Inst. Tech.	✔	University	NUS	Complutense Madrid	✘
Major	Medicine	Medicine	✔	Major	Intl. Rel.	Chem. Eng.	✘

Appendix

Table 12: Case study of CAS under format variation and semantic errors. Case #3 shows that date and entity-format differences can still be semantically correct, whereas Case #17 contains wrong attribute values that should receive low scores. ✔ and ✘ indicate human judgments for each attribute.

Setting	mAcc (%) $\uparrow$	Recall ${}_{\text{N}}$ (%) $\uparrow$	Recall ${}_{\text{P}}$ (%) $\uparrow$	mTokens $\downarrow$	mAcc ${}_{\text{conc}}$ (%) $\uparrow$
Two-Way Metric	82.74 ( $\downarrow$ 0.69)	88.04 ( $\downarrow$ 1.67)	80.60 ( $\downarrow$ 0.67)	348	67.99 ( $\downarrow$ 2.49)
Three-Way Metric	83.43	89.71	81.27	344	70.48

Appendix

Table 13: Ablation study of two-way versus three-way difficulty metrics. The two-way metric merges Unknown and Hallucination into one negative class, while the three-way metric keeps them distinct. mAcc denotes detector accuracy, Recall

{}_{\text{N}}

and Recall

{}_{\text{P}}

denote recall on negative and positive samples, mTokens denotes average context token consumption, and mAcc

{}_{\text{conc}}

denotes conceptual generation accuracy.

		QA-16k		QA-8k
Method	Setting	mTok $\downarrow$	Score $\uparrow$	mTok $\downarrow$	Score $\uparrow$
Original	topk8/sub4	105.00	14.65	108.00	14.40
Calibration	topk8/sub4	101.77	19.30	102.68	23.10
Original	topk32/sub16	385.00	21.65	352.00	26.85
Calibration	topk32/sub16	364.15	26.10	367.24	26.90

Appendix

Table 14: Effect of limited downstream calibration on RULER. Original denotes the uncertainty detector trained only on the synthetic biography data, while Calibration denotes the same detector after calibration on 900 RULER trajectories. mTok denotes the average number of selected context tokens, and Score follows the official RULER string-match accuracy.

Setting	Top-1	Rank of ‘‘29’’	Log $p$ ( ‘‘29’’ ) $\uparrow$
Before Expansion	‘‘9’’	3rd	$-5.13$
After Expansion	‘‘29’’	1st	$-4.56$

Appendix

Table 15: Token ranking and log-probability before and after context expansion. The table reports the candidate ranking for the day token in a birth-date extraction example. Before expansion, the correct token ‘‘29’’ ranks third; after expansion, it becomes top-ranked with log-probability improving from

-5.13

-4.56

Case	Intervention	Quality (%) $\uparrow$	mTokens $\downarrow$	Rollbacks $\downarrow$
FP-1	Suppress expansion	100.0 $\rightarrow$ 100.0	379.36 $\rightarrow$ 367.83 ( $-3.0\%$ )	12 $\rightarrow$ 10
FP-2	Suppress expansion	100.0 $\rightarrow$ 100.0	370.14 $\rightarrow$ 359.75 ( $-2.8\%$ )	9 $\rightarrow$ 9
FN-1	Force expansion	16.7 $\rightarrow$ 66.7	173.47 $\rightarrow$ 238.00 ( $+37.2\%$ )	1 $\rightarrow$ 6
FN-2	Force expansion	8.3 $\rightarrow$ 75.0	316.80 $\rightarrow$ 385.47 ( $+21.7\%$ )	1 $\rightarrow$ 9

Appendix

Table 16: Representative false-positive and false-negative failures under single-decision counterfactual replay. We suppress one unnecessary expansion for false positives and force one missed expansion for false negatives. Quality measures per-instance generation quality, mTokens the average attended context tokens, and Rollbacks the total rollback operations.

Long-Context Inference

Long-context inference studies how language models retrieve, retain, and reason over information spanning very large inputs.

See all 52 papers

Feb 25, 2026

DySCO: Dynamic Attention-Scaling Decoding for Long-Context Language Models

Understanding and reasoning over long contexts is a crucial capability for language models (LMs). Although recent models support increasingly long context windows, their accuracy often deteriorates as input length grows. In practice, models often struggle to keep attention aligned with the most relevant context throughout decoding. In this work, we propose DYSCO, a novel decoding algorithm for improving long-context reasoning. DYSCO leverages retrieval heads--a subset of attention heads specialized for longcontext retrieval--to identify task-relevant tokens at each decoding step and explicitly up-weight them. By doing so, DYSCO dynamically adjusts attention during generation to better utilize relevant context. The method is training-free and can be applied directly to any off-the-shelf LMs. Across multiple instruction-tuned and reasoning models, DYSCO consistently improves performance on challenging long-context reasoning benchmarks, yielding relative gains of up to 25% on MRCR and LongBenchV2 at 128K context length with modest additional compute. Further analysis highlights the importance of both dynamic attention rescaling and retrievalhead guided selection for the effectiveness of the method, while providing interpretability insights into decoding-time attention behavior. Our code is available at https://github.com/princeton-pli/DySCO.

Xi Ye, Wuwei Zhang, Fangcong Yin +2

Princeton Language and Intelligence, Princeton University · Department of Computer Science, New York University

Long-Context Reasoning in Language Models Language Model Decoding Long-Context Language Model Inference

Ask PDF

May 12, 2026

Training-Inference Consistent Segmented Execution for Long-Context LLMs

Transformer-based large language models face severe scalability challenges in long-context generation due to the computational and memory costs of full-context attention. Under practical computation and memory constraints, many inference-efficient long-context methods improve efficiency by adopting bounded-context or segment-level execution only during inference, while continuing to train models under full-context attention, resulting in a mismatch between training and inference execution and state-transition semantics. Based on this insight, we propose a training-inference consistent segment-level generation framework, in which training and inference follow the same segment-level forward execution semantics. During training, consistency with inference is enforced by restricting gradient propagation to KV states carried over from the immediately preceding segment, while permitting head-specific access to past KV states during the forward pass without involving them in gradient propagation. Across long-context benchmarks, our approach achieves performance comparable to full-context attention, while achieving competitive latency-memory trade-offs against strong inference-efficient baselines, and substantially improving scalability at very long context lengths (e.g., approximately 6x lower peak prefill memory at 128K compared to full-context attention with FlashAttention).

Xianpeng Shang, Jiang Li, Zehua Duo +2

College of Computer Science, Inner Mongolia University, Hohhot 010021, China · National & Local Joint Engineering Research Center of Intelligent Information Processing Technology for Mongolian, Hohhot 010021, China · Inner Mongolia Key Laboratory of Multilingual Artificial Intelligence Technology, Hohhot 010021, China +1

Self-Attention Long-Context Language Modeling Long-Context Inference Efficient Language Model Inference

Ask PDF

May 14, 2026

EndPrompt: Efficient Long-Context Extension via Terminal Anchoring

Extending the context window of large language models typically requires training on sequences at the target length, incurring quadratic memory and computational costs that make long-context adaptation expensive and difficult to reproduce. We propose EndPrompt, a method that achieves effective context extension using only short training sequences. The core insight is that exposing a model to long-range relative positional distances does not require constructing full-length inputs: we preserve the original short context as an intact first segment and append a brief terminal prompt as a second segment, assigning it positional indices near the target context length. This two-segment construction introduces both local and long-range relative distances within a short physical sequence while maintaining the semantic continuity of the training text--a property absent in chunk-based simulation approaches that split contiguous context. We provide a theoretical analysis grounded in Rotary Position Embedding and the Bernstein inequality, showing that position interpolation induces a rigorous smoothness constraint over the attention function, with shared Transformer parameters further suppressing unstable extrapolation to unobserved intermediate distances. Applied to LLaMA-family models extending the context window from 8K to 64K, EndPrompt achieves an average RULER score of 76.03 and the highest average on LongBench, surpassing LCEG (72.24), LongLoRA (72.95), and full-length fine-tuning (69.23) while requiring substantially less computation. These results demonstrate that long-context generalization can be induced from sparse positional supervision, challenging the prevailing assumption that dense long-sequence training is necessary for reliable context-window extension. The code is available at https://github.com/clx1415926/EndPrompt.

Han Tian, Luxuan Chen, Xinran Chen +10

1Nankai University · 2Baidu Inc. · 3Shanghai Jiao Tong University +1

Long-Context Language Modeling LLM Fine-Tuning LLM Context Management

Ask PDF

Abstract

Long-Context Inference

Explore similar work

DySCO: Dynamic Attention-Scaling Decoding for Long-Context Language Models

Training-Inference Consistent Segmented Execution for Long-Context LLMs

EndPrompt: Efficient Long-Context Extension via Terminal Anchoring