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

Abstract: Recent works proposed test-time alignment methods that rely on a small aligned model as a proxy that guides the generation of a larger base (unaligned) model. The implicit reward approach skews the large model distribution, whereas the nudging approach defers the generation of the next token to the small aligned model when the large base one is unconfident about its outcome. In this work, we first show that both approaches can be reduced to sampling from similar graphical models, where they differ only in the definition of a rejection criterion (or distribution). Moreover, we argue that the confidence criterion is ill-motivated due to linguistic phenomena like ambiguous phrasing. We propose a novel rejection criterion based on a conservative confidence bet. Experimentally, our novel approach outperforms previous work on several datasets.

• The paper formalizes proxy-based alignment within a unified probabilistic graphical model to capture token-level deferral decisions.
• It critiques traditional confidence-based rejection, introducing a conservative confidence bet that defers only when the proxy clearly outperforms the base model.
• Empirical evaluations on reasoning and commonsense tasks demonstrate improved accuracy and efficiency, notably in mathematical reasoning scenarios.

Proxy-based Test-time Alignment: Rejection Criterion Characterization and Innovation

Background and Motivation

Test-time alignment of LLMs strives to circumvent the computationally expensive and poorly scaling supervised fine-tuning and preference learning phases by biasing the generation process directly. Contemporary alternatives largely center on reward-based or proxy-based mechanisms:

• Reward-based test-time alignment: Utilizes partial or full rewards to guide generation at inference, facing issues with expressivity, compute efficiency, and training complexity.
• Proxy-based test-time alignment: Employs a small aligned model as a proxy to guide or filter the outputs of a larger, unaligned base model. Notable methods include implicit reward extraction, nudging, and Knapsack Approximate Deferral (KAD).

The paper "On the Rejection Criterion for Proxy-based Test-time Alignment" (2604.16146) advances understanding of proxy-based approaches by formalizing them within a unified probabilistic graphical model. It interrogates the theoretical underpinnings of proxy selection criteria, reveals interconnections among popular methods, critiques standard confidence-based rejection decisions, and introduces a new conservative confidence bet criterion.

Formalization Within a Probabilistic Graphical Model

The paper proposes a canonical probabilistic graphical model (PGM) for proxy-based test-time alignment. The generative process at each step samples a draft token from the base model, applies a binary rejection rule, and—depending on acceptance or rejection—either emits the sampled token or defers to an aligned proxy model. Central to this construction is the parameterization of the rejection criterion, which can subsume existing approaches:

• Nudging: Deferral is based on min-entropy or maximum confidence thresholding over the base model's output distribution; the decision is distribution-level and deterministic.
• KAD (Dual): Decision is token-specific, deferring when the sampled token's probability falls below a threshold.
• Implicit Reward: The aligned small model is used to compute a reward as $q^*/q$, shifting the base model distribution via token-wise weighting. The paper proves that, under certain sufficient conditions, such mixtures can also be cast within the PGM framework as a rejection rule, although not universally.

Through this formalism, the paper exposes structural similarities and reducibilities between seemingly disparate alignment procedures, facilitating systematic analysis and extension.

Critical Analysis of Confidence-based Rejection Decisions

The prevalent reliance on base model confidence as a rejection proxy is empirically and theoretically challenged, citing:

• Linguistic ambiguity: Probability mass may be distributed across several equally valid tokens (e.g., "like" vs "such"), which artificially depresses confidence even when the model is accurate.
• Model underconfidence: Lower token probability does not reliably signal poor quality; it often reflects multiple correct completions rather than model uncertainty. This phenomenon is highlighted in recent findings on generative model calibration [dinh2025generativeunderconfident].

Consequently, confidence-based thresholds risk misclassifying ambiguous but satisfactory generations, degrading practical alignment performance.

Proposed Conservative Confidence Bet Criterion

To remediate the inadequacies of absolute and deterministic confidence thresholds, the paper introduces the conservative confidence bet rejection rule:

$$\pi(r = 1 \mid \bar = v) = \llbracket p_v < \max_w q^*_w - \lambda \rrbracket$$

Here, the prediction from the base model ($p_v$) is compared to the maximal token probability under the aligned proxy ($\max_w q^*_w$), optionally with a margin $\lambda$. Deferral occurs if the base model's confidence is strictly lower than what the proxy can achieve. This criterion ensures that deferral is only triggered when the proxy is demonstrably superior for a given token, enabling a more rational collaboration.

Empirical Evaluation

The evaluation spans reasoning tasks (GSM8K, MATH500, SVAMP) and commonsense QA datasets (ARC-Challenge, CommonsenseQA), using OLMo~2 and Qwen~3 model families at two contrasting scales per family. Results demonstrate:

• The proposed conservative confidence bet achieves higher or competitive accuracy relative to previous proxy-based methods and their building blocks.
• For OLMo~2 models, the improvement is particularly marked in mathematical reasoning tasks (e.g., MATH500; $26.4\%$ accuracy, highest among mixtures).
• For Qwen~3, performance gains are more apparent in math tasks but less pronounced in commonsense reasoning, reflecting a smaller gap between base and aligned models.

The approach consistently delivers mixture outcomes superior to both base and proxy models individually, and approaches the performance ceiling set by the fully aligned large model.

Implications and Future Directions

The paper establishes a robust framework for analyzing proxy-based test-time alignment, enabling principled design and evaluation of rejection criteria. The conservative confidence bet leverages both base and proxy information, minimizing unnecessary deferral while maintaining alignment efficacy.

Practical implications include improved sample efficiency and throughput in inference-aligned LLMs, especially in scenarios where computational budget or alignment training costs are prohibitive. Theoretically, the unified PGM offers extensibility to richer proxy collaboration mechanisms, potentially incorporating uncertainty quantification or multi-proxy selection.

Future developments may explore:

• Generalization to multiple proxies: Joint decision criteria across diverse aligned models or modalities.
• Learned rejection distributions: Adapting rejection parameters dynamically based on downstream task or interaction feedback.
• Integration with reward-based alignment: Bridging domain-specific reward architectures with proxy-guided deferral logic for hybrid test-time alignment schemes.

Conclusion

"On the Rejection Criterion for Proxy-based Test-time Alignment" (2604.16146) delivers a comprehensive formal characterization of proxy-based alignment methods, highlights the limitations of confidence-based rejection criteria, and introduces a superior, rational rejection rule grounded in conservative confidence comparison. This work advances both theoretical understanding and empirical performance in efficient, modular test-time alignment for LLMs, with meaningful practical and foundational repercussions for instruction-following, safety, and controlled generation.