Alignment-Dependent Inference in Small Language Models via Budgeted Marginalization over Contextual Priors

Alexis Hu, Yilun Du

Keywords: Small Language Models, Alignment-Dependent Inference, Contextual Priors, Budgeted Marginalization

TL;DR: We show that inference-time computation should be allocated between sampling and contextual diversity based on alignment, as diversity improves or degrades performance depending on how it shifts the effective correctness margin.

Abstract: Inference-time reasoning in small language models often relies on aggregating multiple sampled trajectories, yet different ways of diversifying these samples can have opposite effects under the same test-time budget. This raises a basic allocation problem: when should computation be spent on stochastic repetition versus contextual diversity? We formulate inference as budgeted marginalization over latent contextual priors and develop a bias--variance framework in which diversity reshapes the effective correctness margin induced by contextual priors, while repetition reduces within-prior sampling variance. We show that optimal allocation is governed by an alignment drift coefficient describing how the effective margin changes as additional priors are introduced, yielding aligned, misaligned, and saturating regimes. Empirically, we validate these predictions across GSM8K, MATH500, GPQA Diamond, and HumanEval: diversity-first allocation improves performance in aligned settings, repetition becomes more effective under misalignment, and mixed-alignment settings exhibit saturating gains. Margin-based diagnostics further show that allocation effects are mediated by shifts in the dominance of the correct solution mode. Together, these results show that inference-time performance depends not only on the amount of computation, but on how it is allocated across contextual priors.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 56