Bayesian Test-Time Inference of Task-Aligned Similarity from Weak Interactive Feedback

Weida Liang; Kenji Kawaguchi

cs.LGOct 13, 2026

Bayesian Test-Time Inference of Task-Aligned Similarity from Weak Interactive Feedback

Authors: Weida Liang, Kenji Kawaguchi

Organizations: National University of Singapore

Abstract

Language model systems repeatedly make test-time decisions by comparing similarity: they retrieve exemplars, rank candidates, and revise responses. Yet similarity is not unique: lexical, semantic, structural, and stylistic resemblance can point to different candidates, and the task-aligned notion can change across instances. We introduce BASIS (Bayesian Axis Selection via Interactive Signals), a training-free Bayesian method that uses weak test-time feedback to infer which criterion in a supplied, task-specific library governs the current decision. It maintains a joint posterior over candidate targets and criteria, updates it from scalar scores, pairwise preferences, or verifier outputs, and selects probes by expected information gain. In a finite noisy model, we characterize when the target-criterion pair is recoverable and show why fixed-criterion and fixed-mixture rules can fail under criterion shift. Across concept recovery, reasoning exemplar selection, and response refinement, BASIS improves six-probe concept-recovery Success@1 from 56.0% to 68.1%, reasoning accuracy from 72.3% to 73.1%, and refinement win rate from 61.4% to 63.8% over tuned static mixtures. It also exceeds compute-matched best-of-N selection on reasoning (72.1%) and refinement (61.2%). Gains are largest under criterion mismatch and with larger candidate pools, and matched controls indicate that adaptive criterion weighting explains much of the benefit, with joint target-criterion inference adding further value. BASIS selects only among supplied criteria and does not discover new ones.

Figures & tables

Task	Method	Accuracy / win rate (%)	Time (s)
Reasoning	Dense retrieval	70.0	$4.82\pm 0.11$
	Tuned mixture	72.3	$5.04\pm 0.12$
	Best-of-4 verifier	72.1	$7.91\pm 0.21$
	BASIS	73.1	$7.32\pm 0.18$
Refinement	Tuned mixture	61.4	$6.33\pm 0.15$
	Best-of-6 verifier	61.2	$10.84\pm 0.29$

Table 1: Task outcomes and end-to-end latency at batch size one. Times are means with standard deviations over 100 items. The verifier controls use

N=4

calls for reasoning and

N=6

for refinement. BASIS uses

T=3

feedback calls.

Domain	Both correct	BASIS only	Tuned only	Both wrong
GSM8K	1,801	104	82	513
StrategyQA	1,802	111	88	499
BBH aggregate	2,321	180	150	1,099
Pooled	5,924	395	320	2,111

Table 2: Paired correctness counts over item–seed outcomes: 2,500 each for GSM8K and StrategyQA, 3,750 for BBH, and 8,750 pooled.

Figure 1: Nested

N=64,128,256

candidate pools grow with a heterogeneous support bank, varying both pool size and distractor composition.

$N$	Gather/assembly (ms)	Update/round (ms)	Probe/round (ms)	Marginalization (ms)	Posterior (MiB)
64	4.2	0.07	110	0.03	0.0024
256	14.6	0.18	111	0.06	0.0098
1,024	55.3	0.62	112	0.18	0.0391
4,096	216.8	2.35	114	0.66	0.1563
10,000	527.4	5.68	117	1.55	0.3815

Table 3: Cached-score local inference costs for reasoning. Posterior storage is for the float64 joint distribution over five criteria.

Source $N$	Inference set	Top-1 agreement	Mean top-four overlap	Exact set agreement	Accuracy (%)	Local time (ms)	Posterior (MiB)
1,024	Full pool	100%	4.00/4	100%	68.4	393	0.0391
1,024	Top-128	99.0%	3.96/4	96.1%	68.3	342	0.0049
10,000	Full pool	100%	4.00/4	100%	67.6	897	0.3815
10,000	Top-128	97.8%	3.91/4	91.8%	67.4	356	0.0049
10,000	Top-64	92.8%	3.72/4	81.4%	66.8	348	0.0024

Table 4: Two-stage BASIS with a frozen retriever. Shortlist agreement and overlap are relative to full-pool BASIS. Exact set agreement ignores exemplar order.

Method	Reasoning	Refinement	Distractor accuracy
Tuned mixture	72.3	61.4	65.8
Mean-field $q(z)q(a)$	72.8	63.1	67.4
Criterion marginal, shared probes	72.6	62.7	67.0
Criterion marginal, own probes	72.5	62.4	66.7
History-aware reranker	72.6	62.6	67.0
BASIS	73.1	63.8	68.2

Table 5: Matched-feedback controls. Entries are task accuracy or preference win rate in percent.

Figure 2: Performance and posterior diagnostics under criterion corruption and removal. The removed-criterion condition is an oracle stress test. Maximum posterior reports the mean per-episode maximum mass over criteria.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Task	Candidate space $\mathcal{Z}$	Probe type $q_{t}$	Latent criteria $\mathcal{A}$	Weak feedback
Concept recovery	lexical concepts with relation-based distractors	single concept or pairwise query	semantic, relation, part–whole, morphology, association	scalar score or pairwise preference
Exemplar selection	exemplar bank from training problems	single exemplar or exemplar pair	semantic, lexical, entity, structural, answer-consistency	verifier delta, structural critic, pairwise preference
Response refinement	candidate responses for one instruction	response or response pair	helpfulness, specificity, conciseness, safety, empathy, formality	preference label or rubric score

Appendix

Table 6: Three domains used to instantiate the same latent-criterion inference problem.

Task	Feedback type	Source	Form	Shared with non-oracle baselines?
Concept recovery	scalar similarity / pairwise preference	simulated from the hidden criterion with Gaussian or Bradley–Terry noise	scalar / pairwise	yes
Reasoning exemplar selection	verifier delta	downstream answer verifier applied to candidate exemplars or prompted traces	scalar	yes
Reasoning exemplar selection	structural signal	template overlap or structural critic score computed from reasoning skeletons	scalar	yes
Response refinement	preference feedback	dataset preference label when available, otherwise rubric-based weak supervision	pairwise / scalar	yes

Appendix

Table 7: Feedback sources used in each task. These signals are holistic weak judgments: they do not reveal the hidden criterion label, and they are available to every non-oracle method that uses external feedback.

Criterion	Implementation	Normalization	External model?	Updated after feedback?
Semantic	cosine similarity in the frozen encoder space	pool-wise min–max to $[0,1]$	no	no
Lexical	BM25 or token-overlap relevance	pool-wise min–max to $[0,1]$	no	no
Entity	overlap F1 over entities, slots, or variables	native $[0,1]$ score	no	no
Structural	overlap between extracted reasoning templates, operator sequences, or decomposition graphs	pool-wise min–max to $[0,1]$	optional critic	no
Style	lightweight rubric classifier or rubric score for tone, conciseness, empathy, or formality	affine calibration to $[0,1]$ on a development split	optional	no
Answer-consistency	verifier pass rate for the exemplar’s supplied answer to its own question	binary or calibrated scalar in $[0,1]$	yes	no

Appendix

Table 8: Implementation of the criterion-conditioned scorers. Scores are computed once for the current candidate pool, normalized to a comparable range, and cached across interaction rounds. Feedback updates posterior weights over criteria rather than recomputing base scores.

Method	Same pool	Same criterion scores	Feedback history	Additional interaction calls
Static mixture	yes	yes	no	0
Rubric-guided LLM reranker	yes	rubric or criterion descriptions	no	1
Best-of- $N$ verifier	yes	verifier only	no	$N$
BASIS	yes	yes	yes	$T$
Oracle-criterion	yes	yes + true criterion	no	0

Appendix

Table 9: Information access and additional interaction calls. Call counts begin after the shared criterion-scoring cache has been built. We count verifier, critic, and LLM-judge calls as additional interaction calls. Posterior updates and Monte Carlo probe evaluation are local operations.

Method	S@1@4	S@1@6	Criterion acc.	Q@90 $\downarrow$
Random probing	$17.9\pm 0.8$	$24.1\pm 0.9$	—	$9.8\pm 0.2$
Greedy semantic NN	$35.2\pm 1.1$	$43.1\pm 0.9$	—	$7.7\pm 0.2$
Target-only Bayes	$43.8\pm 0.8$	$51.9\pm 0.8$	—	$6.8\pm 0.1$
Static mix (val-tuned)	$47.6\pm 0.8$	$56.0\pm 0.7$	—	$6.5\pm 0.1$
Active, no criterion	$47.1\pm 0.7$	$55.6\pm 0.7$	—	$6.2\pm 0.1$
BASIS	$59.8\pm 0.7$	$68.1\pm 0.6$	$80.9\pm 0.8$	$4.6\pm 0.1$

Appendix

Table 10: Controlled concept recovery over 500 episodes. Entries are mean

\pm

SD over five seeds. S@1@

k

is Success@1 after

k

probes, Criterion acc. is hidden-criterion identification accuracy, Q@90 is the mean number of probes required to reach 90% of the final Success@1 rate after 12 probes.

Method	GSM8K	StrategyQA	BBH (3 tasks)	Avg.
Random exemplars	65.3	67.1	56.2	62.9
BM25	70.6	71.3	60.4	67.4
Dense retrieval	73.1	73.8	63.1	70.0
ByCS-style Bayesian selection	73.5	74.2	63.7	70.5
RGER-style structural retrieval	74.4	75.1	65.0	71.5
LLM reranker	74.7	74.8	64.6	71.4

Appendix

Table 11: Full reasoning exemplar-selection results. The average is the equal-weight macro-average across the three task columns. Entries are means over five seeds.

Method	HH win	OASST win	UF win	Avg.	Attribute match
Direct preference ranker	57.4	55.9	60.1	57.8	62.9
LLM-as-judge reranking	58.3	57.1	60.5	58.6	63.8
Self-Refine	59.8	58.2	61.4	59.8	64.6
Target-only Bayes	60.7	59.4	62.0	60.7	65.1
Rubric-guided LLM reranker	60.4	58.9	62.3	60.5	65.5
Static mix (uniform)	60.8	59.1	62.7	60.9	66.0

Appendix

Table 12: Full response-refinement results. Entries are means over five seeds. Attribute match measures agreement between the selected response and the hidden preference criterion.

Variant	Concept S@1@6	Reason Avg	Refine Win	Distr. Drop $\downarrow$
BASIS	$\mathbf{68.1}$	$\mathbf{73.1}$	$\mathbf{63.8}$	$\mathbf{4.9}$
No latent criterion	$51.9$	$70.9$	$60.7$	$8.8$
No active probing	$61.7$	$72.0$	$62.4$	$6.6$
Uniform static mixture	$54.8$	$71.7$	$60.9$	$6.9$
Single semantic criterion	$43.1$	$68.9$	$57.9$	$10.8$

Appendix

Table 13: Ablations across tasks. Concept S@1@6 is from concept recovery. Reason Avg and Refine Win are the averages from Tables 11 and 12 . Distr. Drop is the robustness drop under distractors.

Asset / release	License or terms	Upstream source
WordNet 3.0 [ 19 ]	Princeton WordNet License; use, copying, modification, and distribution are permitted with the required copyright notice and disclaimer.	WordNet license
ConceptNet 5.7 [ 26 ]	Data: CC BY-SA 4.0. The repository distinguishes the graph-data license from the Apache 2.0 code license.	License / data credits
SimLex-999, original English release [ 12 ]	The official download page requests citation but does not state a separate dataset license. No SimLex files are reproduced or redistributed here.	SimLex project page
WordSimilarity-353 v1.0 (2002) [ 9 ]	Creative Commons Attribution 4.0 International (CC BY 4.0); cite the original collection.	WordSim-353 source and license
GSM8K, official 2021 release [ 6 ]	MIT License.	OpenAI repository license
StrategyQA, official 2021 release [ 10 ]	MIT License, as stated in the dataset repository.	StrategyQA repository license

Appendix

Table 14: Data assets, release identifiers, and source-stated licenses or terms.

Model identifier	License	Upstream model card
Qwen/Qwen3-Embedding-4B	Apache License 2.0	Qwen model card
Qwen/Qwen3-8B	Apache License 2.0	Qwen model card
mistralai/Ministral-3-14B-Reasoning-2512	Apache License 2.0	Mistral model card

Appendix

Table 15: Model checkpoints and their published license identifiers.

Inference-Time Language Model Alignment

Inference-time language model alignment steers or selects outputs during generation to better meet safety, preference, and task goals without retraining.

See all 78 papers

ARBITER: Reasoning Trajectory Basins and Majority Vote Failures in Test-Time Sampling
May 25, 2026Meng Cai, Lars Kulik, Farhana ChoudhurySelf-Consistency Decoding Test-Time Scaling
School of Computing and Information Systems University of Melbourne
Bayesian Fine-tuning Yields Language Models that are as Bayesian as their Beliefs Allow
Sep 30, 2026Polina Tsvilodub, Andreas Waldis, Linlu Qiu +2Fine-Tuning LLM Fine-Tuning
University of Tübingen · Massachusetts Institute of Technology · New York University
On the Rejection Criterion for Proxy-based Test-time Alignment
Apr 17, 2026Ayoub Hammal, Pierre Zweigenbaum, Caio CorroLLM Alignment Inference-Time Language Model Alignment
Université Paris-Saclay, CNRS, LISN · INSA Rennes, IRISA, CNRS, Université de Rennes

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

ARBITER: Reasoning Trajectory Basins and Majority Vote Failures in Test-Time Sampling

Meng Cai, Lars Kulik, Farhana Choudhury

School of Computing and Information Systems University of Melbourne

Self-Consistency Decoding Test-Time Scaling LLM Reliability LLM Reasoning Majority Voting

Open paper AskPDF arXiv

When language models use test-time sampling, they generate multiple reasoning trajectories and select an answer by majority vote. We show that these trajectories are not independent: for a given question, they concentrate into a small number of clusters, or reasoning basins, each defined by a normalized final answer and the solutions that reach it. A majority vote therefore selects the most stable basin rather than the most accurate one, which creates wrong-majority failures where the correct answer is present but outvoted. We introduce ARBITER, a model-agnostic approach that models interactions between basins using only the base model's own sampled outputs, hidden states, and derived evidence. Most direct correction strategies fail; ARBITER instead uses conservative additive evidence on top of consensus. In its simplest parameter-free form, ARBITER-Δ adds same-model evidence to the majority prior, while ARBITER-Enc augments this with bounded residual signals from hidden states over complete solutions. On GSM8K with Qwen3-4B, consensus over K=24 samples achieves around the mid-94% range, while a same-pool top-2 oracle reaches around the mid-96% range. ARBITER recovers a subset of these cases using zero external information. Across three model families and three math benchmarks, it yields consistent gains with no net-negative cases; for example, on Llama-3.1-8B MMLU-HS-Math, it improves accuracy from the mid-78% range to the mid-82% range, recovering about 22% of the available oracle headroom, indicating that this headroom can be partially recovered from the sample pool itself.

Abstract

Inference-Time Language Model Alignment

Explore similar work

ARBITER: Reasoning Trajectory Basins and Majority Vote Failures in Test-Time Sampling

Bayesian Fine-tuning Yields Language Models that are as Bayesian as their Beliefs Allow

On the Rejection Criterion for Proxy-based Test-time Alignment