CALIBUDGET: Calibration-Guided Source Allocation for Fixed-Budget Mixed-Reasoning Adaptation

Yupeng Chang; Yuan Wu

cs.LGSep 29, 2026

CALIBUDGET: Calibration-Guided Source Allocation for Fixed-Budget Mixed-Reasoning Adaptation

Authors: Yupeng Chang, Yuan Wu

Organizations: School of Artificial Intelligence, Jilin University · Key Laboratory of Symbolic Computation and Knowledge Engineering, Jilin University

Abstract

Fixed-budget adaptation from heterogeneous data sources requires deciding not only how much data to use, but how much exposure each source receives. Size-proportional rules can crowd out small sources, whereas difficulty-only rules can chase noisy estimates or allocate residual budget to nearly saturated pools. We introduce CALIBUDGET, a floor-protected, reliability-aware integer allocator that treats source exposure as an explicit adaptation variable. From small train-internal calibration splits, it combines model need, post-floor availability, and bootstrap stability, then produces exact capacity-respecting quotas without changing the model, objective, or total budget. In a controlled setting combining mathematical and commonsense data, CALIBUDGET improves CommonAvg, FragileAvg, and MacroAvg over validation-error-with-floor, the strongest matched comparator, in all three paired LLaMA-2-7B LoRA+ runs. The respective mean gains are 0.56, 0.46, and 0.41 percentage points (pp). Overall increases by 0.18 pp, whereas MathAvg decreases by 0.20 pp, exposing a coverage-retention boundary rather than a uniform gain. CALIBUDGET changes only 1.14-1.42% of the source budget but improves performance in 15 of 24 comparisons across commonsense tasks and seeds. These results suggest that small changes in source quotas can matter; example-level selection can then determine which examples fill each quota.

Figures & tables

Figure 1: Overview of CALIBUDGET. A heterogeneous multi-source pool is split into train-internal calibration data and candidate pools. Floor protection guarantees minimum exposure; post-floor availability, source need, and estimate reliability define the utility used for residual top-up, producing capacity-respecting integer quotas under budget $B$ . The allocated source subset is combined with a fixed companion training set and adapted under an otherwise unchanged training protocol.

		Retention	Source-balanced		Aggregate
Method	Allocation signal	Math (%)	Common (%)	Fragile (%)	Macro (%)	Overall (%)
Pooled uniform	pooled random	30.43	69.98	67.57	62.07	50.21
Proportional	source size	30.72	69.48	66.84	61.72	50.10
Equal	equal source quota	30.51	71.09	71.86	62.97	50.80
Floor-Sqrt	floor + $\sqrt{\text{residual size}}$	30.47	71.03	71.19	62.92	50.75
Val-error+floor	floor + calibration error	30.91	71.31	72.15	63.23	51.11

Table 1: Three-seed mean performance (%) on LLaMA-2-7B under matched data budgets, adaptation backend, and evaluation. Bold denotes the best descriptive mean; underlining marks the strongest matched comparator when it is runner-up.

Seed	Math	Common	Fragile	Macro	Overall
42	$-1.12$	$+0.49$	$+0.59$	$+0.16$	$-0.32$
43	$+0.27$	$+0.38$	$+0.32$	$+0.36$	$+0.33$
44	$+0.26$	$+0.81$	$+0.48$	$+0.70$	$+0.53$
Mean	$-0.20$	$+0.56$	$+0.46$	$+0.41$	$+0.18$
SD	$0.80$	$0.22$	$0.14$	$0.27$	$0.44$

Table 2: Paired performance differences (CALIBUDGET minus Val-error+floor, pp). Differences are computed before display rounding; SD is across seeds.

Frozen references, seed 42
Diagnostic	Math	Common	Fragile	Macro	Overall
CALIBUDGET	30.20	71.29	71.79	63.07	50.74
Val-error+floor	31.32	70.80	71.20	62.91	51.06
Exploratory normalization variants, seed 42
Group-norm A	30.85	71.46	72.06	63.34	51.16
Group-norm B	30.48	71.84	72.86	63.57	51.16

Table 3: Exploratory seed-42 normalization variants and a fixed-index seed-43 check (%). Internal variants A–C alter source normalization; recovery adds a short retention-oriented phase. None revises the frozen CALIBUDGET configuration.

Figure 2: Diagnostics under the frozen comparison: (a) paired seed effects relative to validation-error-with-floor; (b) the observed MathAvg–FragileAvg boundary, including exploratory variants; (c) task-level accuracy differences (pp); and (d) mean source-quota shifts under the 10,000-example budget. No diagnostic panel selects a method or hyperparameter.

Model	Method	MathAvg (%)	CommonAvg (%)	FragileAvg (%)	MacroAvg (%)	Overall (%)
Qwen2.5-7B	Val-error+floor	$65.77_{\pm 0.30}$	$84.95_{\pm 0.29}$	$88.42_{\pm 0.37}$	$81.12_{\pm 0.29}$	$75.36_{\pm 0.30}$
	CALIBUDGET	$66.13_{\pm 0.31}$	$85.23_{\pm 0.21}$	$88.77_{\pm 0.32}$	$81.41_{\pm 0.18}$	$75.68_{\pm 0.19}$
LLaMA-3.1-8B	Floor-Sqrt	$60.77_{\pm 0.37}$	$82.99_{\pm 0.17}$	$84.60_{\pm 0.42}$	$78.54_{\pm 0.07}$	$71.88_{\pm 0.11}$
	CALIBUDGET	$60.72_{\pm 0.55}$	$82.97_{\pm 0.19}$	$84.57_{\pm 0.26}$	$78.52_{\pm 0.26}$	$71.84_{\pm 0.37}$

Table 4: Cross-backbone audits. Qwen2.5-7B ( Qwen et al., 2025 ) averages three independently constructed adaptive allocations; LLaMA-3.1-8B ( Grattafiori et al., 2024 ) averages three training seeds while holding each method’s seed-42 allocation fixed. Comparators are matched within each audit. Main values are mean accuracy (%); subscripts report SD (pp).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Package component	Purpose	Verification boundary
src/calibudget/allocation.py	Deterministic floor–residual allocation, capacity-aware largest-remainder rounding, and selected-index hashing.	Model- and dataset-independent synthetic inputs.
src/calibudget/proxy.py	Canonical subgroup-score record validation and stable four-shard aggregation.	Validates structure and ordering without exposing experiment-specific score values.
src/calibudget/data.py	Record normalization, normalized prompt hashing, and pairwise-disjoint role checks.	Dataset-agnostic integrity contracts.
src/calibudget/modeling.py	Response-only tokenization-length contract with caller-supplied tokenizer.	No model download or model identity is required.
src/calibudget/sketch.py	Deterministic CountSketch for finite one-dimensional vectors.	Synthetic numeric arrays only.
src/calibudget/audit.py	Strict JSON-object loading, SHA-256 calculation, and regular-file binding checks.	Local file integrity and provenance checks.

Appendix

Table 5: Map of the supplementary package. The compact modules expose the contracts exercised by the default tests, while the sanitized snapshots preserve additional implementation context.

Evidence layer	What the evidence supports	What it does not establish
Main paired study	Consistent gains on CommonAvg, FragileAvg, and MacroAvg under matched LLaMA-2 budgets and LoRA+ controls.	Uniform improvement, large-sample significance, or absence of a MathAvg trade-off.
Reliability ablation	The complete utility is consistently stronger than removing $\rho_{k}$ on source-balanced metrics in the three observed seeds.	Independent causal effects among interacting utility terms.
Quota geometry	Small, structured source reallocations can accompany measurable aggregate changes.	A monotonic quota–accuracy response for individual tasks.
Qwen2.5 audit	Directionally similar descriptive gains under an independently recalibrated second backbone.	Backbone-invariant superiority; all reported bootstrap intervals cross zero.
LLaMA-3.1 and modern panel	Transferred LLaMA-3.1 allocations are near parity, and the broader modern-task panel shows no reliable gain.	Reliable gains under transferred quotas, on code tasks, or with full-parameter tuning.
Floor and split audits	Quota rankings are stable across calibration splits, while the floor operating point remains consequential.	A universal floor value or robustness to arbitrary source partitions.

Appendix

Table 6: Claim boundary. Each evidence layer is paired with the conclusion supported by the current results and with a stronger conclusion that those results do not establish.

Method	Math	Common	Fragile	Macro	Overall
Val-error+floor	65.38	84.55	88.04	80.71	74.96
CALIBUDGET	65.11	84.93	88.54	80.97	75.02
$\Delta$ (pp)	$-0.26$	$+0.39$	$+0.50$	$+0.26$	$+0.06$

Appendix

Table 7: Separately audited Qwen2.5-7B seed-42 reconstruction (accuracy, %). Differences are CALIBUDGET minus Val-error+floor and are computed before display rounding.

Data Mixture Optimization

Data mixture optimization studies how to allocate training examples across domains and tasks to improve model performance under data and compute constraints.

See all 40 papers

Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning
May 26, 2026Manh Nguyen, Sunil Gupta, Hung LeEfficient Language Model Reasoning Adaptive Sampling
Applied Artificial Intelligence Initiative Deakin University, Australia
Rethinking Calibration for High-Sparsity LLM Pruning: From Single Sources to Capability Coverage
Jun 2, 2026Hu Xu, Zhaolong Xing, Congcong Liu +5LLM Pruning Model Calibration
Shanghai Jiao Tong University
Instance-Optimal Estimation with Multiple LLM Judges on a Budget
May 22, 2026Junghyun Lee, Sanghwa Kim, Yassir Jedra +2Parameter Estimation LLM Evaluation
KAIST AI · ICL EEE · KTH EECS

May 26, 2026J/K move · Enter open · S save

Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning

Manh Nguyen, Sunil Gupta, Hung Le

Applied Artificial Intelligence Initiative Deakin University, Australia

Efficient Language Model Reasoning Adaptive Sampling Uncertainty Sampling Inference-Time Compute Allocation LLM Reasoning

Open paper AskPDF arXiv

Sampling multiple responses improves language model reasoning, but uniform compute allocation is inefficient because easy questions are over-sampled while hard questions remain under-explored. We propose Uncertainty-Aware Budget Allocation (UAB), a concave integer optimization framework that reallocates a fixed sampling budget using uncertainty estimated from the initial samples themselves. In Phase-1, every question receives a small fixed number of generations. Their answer disagreement, measured by vote entropy, provides a difficulty signal while these generations contribute to the final vote. In Phase-2, the remaining budget is allocated by a marginal-greedy algorithm that optimally solves a concave coverage-maximization surrogate, concentrating samples on questions whose initial answers disagree. Across five open-weight models (1.5B--27B parameters) and five reasoning benchmarks of varying difficulty, UAB improves average accuracy by

+2.3\%

over uniform allocation, and achieves gains of up to

+5.5\%

on individual benchmarks, with the largest gains in low-resource settings. Moreover, UAB meets the target budget exactly and requires no auxiliary model or additional LLM calls. Code is publicly available at https://github.com/manhitv/UAB.

Abstract

Data Mixture Optimization

Explore similar work

Uncertainty-Aware Budget Allocation for Adaptive Test-Time Reasoning

Rethinking Calibration for High-Sparsity LLM Pruning: From Single Sources to Capability Coverage

Instance-Optimal Estimation with Multiple LLM Judges on a Budget