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

AI summary

Overview

• Research area: Test-time inference and selection for language model systems, combining Bayesian posterior inference, active probe selection, and retrieval/reranking.
• Technical level: Intermediate (the math is presented formally, but the core idea and experiments are described accessibly).
• Scope: A training-free Bayesian method, BASIS, that uses weak interactive feedback to infer which criterion in a supplied task-specific library governs the current selection decision, evaluated on concept recovery, reasoning exemplar selection, and response refinement.

What This Paper Is About

Language model systems repeatedly make test-time decisions by comparing similarity: they retrieve exemplars, rank candidates, and revise responses. The problem is that similarity is not unique, since lexical, semantic, structural, and stylistic resemblance can point to different candidates, and the task-aligned notion can change from instance to instance, so a fixed embedding or fixed mixture score may select a plausible but task-misaligned candidate. The paper's goal is to infer, from weak scalar, pairwise, or verifier feedback at test time, which criterion in a supplied library governs the current decision, and to use that posterior to select better targets without any training.

Key Contributions

1. The authors formulate feedback-driven selection as joint inference over a candidate and a finite library of relevance criteria, and introduce BASIS (Bayesian Axis Selection via Interactive Signals) for this task, with a joint posterior over candidate targets and criteria updated from scalar scores, pairwise preferences, or verifier outputs.
2. They characterize when the target–criterion pair is recoverable, under separability and probe coverage assumptions, and identify settings in which a fixed criterion or a fixed mixture cannot track instance-dependent relevance (Proposition 1 on fixed-criterion non-identifiability and Proposition 2 on static mixtures).
3. They evaluate concept recovery, reasoning exemplar selection, and response refinement, using matched controls (mean-field, criterion-marginal, history-aware reranker) to separate adaptive criterion weighting from joint target–criterion inference.
4. They show that probes selected by expected information gain correspond exactly to the mutual information between the latent target–criterion pair and the feedback (Proposition 3), making the greedy probe the natural one-step choice.

Main Findings

• Concept recovery gains: Across 500 episodes with 50 candidates and up to 12 probes, BASIS improves six-probe concept-recovery Success@1 from 56.0% (tuned static mixture) to 68.1%, against an oracle-criterion upper bound of 79.5%. Criterion identification accuracy is 80.9%.
• Probe efficiency: BASIS reaches 90% of its final 12-probe recovery rate after a mean of 4.6 probes, compared with 6.5 probes for the tuned mixture.
• Downstream reasoning and refinement: Reasoning accuracy rises from 72.3% (tuned mixture) to 73.1%, and refinement win rate from 61.4% to 63.8%. BASIS also exceeds compute-matched best-of-N selection on reasoning (72.1%) and refinement (61.2%).
• Statistical strength of the gains: The paired item-clustered interval for the 0.8-point reasoning difference is [-0.9, 2.5], while the 2.4-point refinement gain has interval [1.3, 3.5]. BASIS closes 29.6% of the reasoning gap and 54.5% of the refinement gap between the tuned mixture and oracle-criterion results.
• Paired item–seed outcomes: Pooled over 8,750 item–seed outcomes, 5,924 were correct for both methods, 395 were BASIS-only, 320 were tuned-mixture-only, and 2,111 were wrong for both. Reasoning results are equal-weight macro-averages over GSM8K, StrategyQA, and aggregate BBH.
• Robustness under distractors: On reasoning, BASIS drops 4.9 points with distractors (73.1% to 68.2%) versus 6.5 points for the tuned mixture (72.3% to 65.8%); BASIS at T = 5 feedback rounds reaches 71.4% on the distractor split. On refinement, clean/mismatch win rates are 63.8/61.5 for BASIS, 61.4/57.2 for the tuned mixture, and 65.8/64.4 for the oracle-criterion upper bound, corresponding to drops of 2.3, 4.2, and 1.4 points.
• Larger candidate pools help: With a heterogeneous support bank of 12,416 unique examples (at least 10,218 eligible per query), BASIS improves over the tuned mixture by 0.8, 2.4, and 2.8 points at N = 64, 128, and 256 candidates.
• Latency cost: End-to-end latency increases by 2.28 seconds (45%) for reasoning and 2.75 seconds (43%) for refinement relative to the tuned mixture, but BASIS remains 0.59 and 1.76 seconds faster than the corresponding verifier controls. Most added time is spent acquiring feedback.
• What joint inference adds: Relative to mean-field inference, the joint posterior has point-estimate gains of 0.3 points on reasoning, 0.7 on refinement, and 0.8 on distractor accuracy. Against the criterion-marginal controls, the joint model adds 0.5–0.6 points on reasoning, 1.1–1.4 on refinement, and 1.2–1.5 on distractors. Removing the latent criterion lowers reasoning and refinement to 70.9 and 60.7; disabling active probing yields 72.0 and 62.4.
• Probe approximation and feedback-model sensitivity: With K = 16 and M = 32, the joint-information policy selects the same probe as a K = 64, M = 64 reference in 95.2% of rounds, taking 110 ms per round versus 856 ms for the reference. A criterion-entropy surrogate takes 68 ms and gives 73.0/63.5 on reasoning/refinement. Under Laplace, Student-t3, or heteroscedastic Gaussian feedback while inference retains a Gaussian likelihood, six-probe concept-recovery Success@1 is 67.5, 67.0, and 66.6, versus 68.1 under matched Gaussian feedback.
• Scaling of local inference: At N = 10,000, BASIS retains exact posterior marginalization over the full pool; posterior operations grow approximately linearly with pool size while probe selection remains nearly flat (117 ms per round at N = 10,000 versus 110 ms at N = 64). A Top-128 shortlist at N = 10,000 retains 3.91 of four exemplars on average and reaches 67.4% accuracy versus 67.6% for full-pool inference, while reducing local inference time from 897 to 356 ms.
• Library coverage and fallback: A fallback based on criterion entropy and posterior-predictive likelihood achieves AUROC 0.86 and triggers on 72% of missing-criterion episodes versus 9% of full-library episodes, raising reasoning from 70.7 to 72.4 and refinement from 61.0 to 61.9 under criterion removal, and reducing high-confidence failures from 9.2% to 4.8%. Expanding the libraries from 5/6 to 24 irrelevant, correlated, or redundant criteria changes accuracy by at most 0.2 points, while probe-selection time rises from 110 to 424 ms per round.
• Qualitative traces: In a GSM8K trace, BASIS selects an exemplar aligned with the updated remainder and returns $81, while the tuned mixture returns $63. In a BBH ordering failure, the tuned mixture answers hummingbird (D) correctly but BASIS selects a direct-position exemplar and answers cardinal (B); the final pairwise feedback favors the direct-position example with posterior-predictive probability 0.684, and semantic, structural, and answer-consistency posterior masses change from (0.309, 0.202, 0.198) to (0.276, 0.172, 0.266). In an HH refinement trace, the empathy posterior rises from 0.17 to 0.55, and when humor is absent BASIS favors formality.

Methodology in Plain English

BASIS treats selection as inference over two unknowns at once: which candidate is useful, and which criterion from a supplied menu explains why. The method is training-free and reuses a frozen representation backbone. Each criterion defines a projection or scoring rule, producing a criterion-conditioned similarity between a probe and a candidate. The system starts with a prior over the joint candidate–criterion state and updates it after every weak feedback observation, using a Gaussian model for scalar scores, a Bradley–Terry model for pairwise preferences, and a logistic model for pass/fail verifier signals. At each round, it picks the probe expected to reduce the most uncertainty, which the authors show equals maximizing one-step mutual information; it uses top-K = 16 candidate pruning with M = 32 Monte Carlo feedback samples, and at most three downstream feedback rounds. For target recovery it returns the maximum-a-posteriori candidate; for retrieval or refinement it ranks candidates by the posterior-weighted utility. Feedback is holistic and weak: no non-oracle method sees the hidden criterion label. All methods share candidate pools, backbone, prompts, and feedback or verifier access. Experiments use 500 concept-recovery episodes with 50 candidates and up to 12 probes; GSM8K, StrategyQA, and three BIG-Bench Hard tasks with 64 support examples, an initial dense shortlist of 32, at most three feedback rounds, and four exemplars selected; and 1,000 prompts each from Anthropic HH, OpenAssistant, and UltraFeedback with eight candidate responses. The frozen retriever is Qwen/Qwen3-Embedding-4B, the reasoning generator is Qwen/Qwen3-8B with thinking disabled, and the shared verifier is mistralai/Ministral-3-14B-Reasoning-2512, with greedy decoding for final answers. Results average five seeds, with confidence intervals from 10,000 paired bootstrap samples.

Why This Matters

The paper reframes a common failure mode of similarity-based selection: systems can be confident and wrong when a fixed scoring rule does not match the instance-specific notion of relevance the feedback actually rewards. It shows a training-free way to make that notion an explicit, inspectable latent variable, and it identifies concrete settings (criterion mismatch, larger candidate pools, refinement) where adaptation pays off, while also reporting the latency and coverage limits of the approach.

Real-world applications:

• Retrieval-augmented and few-shot systems that must pick demonstrations or passages matching an operation sequence, entity structure, or style rather than surface topic.
• Response refinement and preference alignment, where the winning revision may match tone, empathy, or formality rather than subject matter.
• Reasoning assistants that choose among candidate solutions using verifier or answer-check signals when topical similarity is misleading.
• Systems with a low-confidence routing need, where a criterion-entropy fallback can send uncertain cases to a fixed selector instead of acting on a possibly wrong criterion.

Industry relevance: the method requires no training and plugs into existing frozen encoders, verifiers, and rerankers, so the practical cost is added test-time feedback calls rather than model updates. The reported timing (2.28 seconds and 2.75 seconds of added end-to-end latency, and second-stage inference as low as 356 ms with a shortlist) gives a concrete sense of the trade-off, and the authors note that for routine reasoning instances with a small accuracy difference a static selector remains a lower-latency option.

Future Directions

• Criterion discovery: BASIS selects only among supplied criteria and cannot discover new ones; the authors state this explicitly, and an open question is how to expand or learn the criterion library rather than assume a useful criterion is already represented.
• Live feedback and preference drift: most feedback comes from offline verifiers or rubric scores, and live human feedback and preference drift remain untested.
• Cross-model and serving generalization: the experiments use one retriever, one generator, and one verifier stack, and the reasoning results describe a matched, thinking-disabled setup that is not directly comparable to benchmark scores reported with default thinking enabled.
• Reducing the feedback and scoring cost: full-pool scoring can remain costly for criteria requiring external models, and the authors note that timing excludes candidate generation, first-stage retrieval, structural extraction, and one-time external scoring.

Target Audience

Researchers and practitioners working on test-time compute, retrieval and reranking, Bayesian inference for in-context behavior, and active learning, along with engineers building selection or refinement pipelines who want a training-free method with inspectable criterion posteriors and clear latency numbers. Readers should be comfortable with basic probability and preference models, but the paper also supplies the assumptions, theorems, proofs, and experimental protocol in sufficient detail for those who want to verify the claims.

Authors’ 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.