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

2026-04-17

AI Summary

This work addresses the limitations of existing test-time alignment methods, whose reliance on ad hoc rejection criteria—such as confidence thresholds—lacks theoretical grounding in scenarios involving linguistic ambiguity, thereby constraining alignment performance. To overcome this, the paper unifies implicit reward and nudging approaches within a single graphical model framework and introduces a novel rejection criterion based on conservative confidence games, which effectively mitigates misjudgments caused by language ambiguity. By integrating proxy-model guidance, graph-based sampling, and the proposed rejection mechanism, the method achieves significant improvements over current approaches across multiple benchmark datasets, demonstrating both the effectiveness and robustness of the new criterion.

Technology Category

• Machine Learning: Calibration & Uncertainty Quantification
• Reasoning under Uncertainty: Graphical Models
• Natural Language Processing: Language Grounding & Multi-modal NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

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.

Problem

• test-time alignment
• proxy-based
• rejection criterion
• confidence
• language ambiguity

Innovation

• test-time alignment
• proxy-based alignment
• rejection criterion
• conservative confidence bet
• large language models