On the Rejection Criterion for Proxy-based Test-time Alignment
| Authors | Ayoub Hammal et al. |
| Year | 2026 |
| Field | NLP |
| arXiv | 2604.16146 |
| Download | |
| Categories | cs.CL |
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.
Engineering Breakdown
Plain English
This paper addresses a practical problem in test-time alignment: how to use a small, well-aligned model to guide generation from a larger unaligned base model. The authors show that two existing approaches—implicit reward and nudging—can be unified under a single graphical model framework that differs only in their rejection criteria. They identify that the standard confidence-based rejection criterion is flawed because linguistic ambiguity (like ambiguous phrasing) makes confidence scores unreliable signals. The core contribution is a new rejection criterion based on conservative confidence betting that better decides when to defer token generation to the smaller proxy model.
Core Technical Contribution
The paper's key insight is that implicit reward and nudging approaches, despite appearing different, optimize the same underlying graphical model—they only differ in how they decide when to switch from the large model to the small aligned proxy. The authors challenge the assumption that model confidence is a good rejection criterion, showing that linguistic phenomena create false confidence signals that don't correlate with actual alignment quality. They propose a conservative confidence betting criterion that replaces raw confidence thresholds with a principled decision rule that accounts for the cost of incorrect rejections. This unification framework and alternative rejection criterion represent both a theoretical clarification and a practical improvement over prior test-time alignment methods.
How It Works
The system operates on two models: a large base model (unaligned) and a smaller proxy model (aligned). At generation time, for each token position, the rejection criterion evaluates whether to sample from the large model or defer to the small aligned model. The unified graphical model formulation shows both methods are sampling from p(tokens | base_model_output, proxy_guidance, rejection_criterion), where the rejection criterion determines the conditional probability of deferring. Instead of using the base model's confidence score directly—which peaks even on ambiguous tokens—the conservative confidence betting criterion creates a threshold that accounts for the empirical cost of rejection errors. When the rejection criterion fires (typically when confidence is low and the difference between top choices is small), generation switches to the small aligned model, which has been fine-tuned to produce safe, aligned outputs. The output is a sequence of tokens drawn from a mixture distribution that intelligently selects between the two models based on the rejection decision at each step.
Production Impact
For teams running large language models in production, this approach enables alignment without full model fine-tuning—a significant cost savings when you have a large base model and smaller aligned model available. The practical impact is reducing hallucinations and unsafe outputs by deferring uncertain generations to a trustworthy proxy, while maintaining the base model's speed and capability for confident, well-aligned predictions. Integration is straightforward: replace the standard confidence threshold with the conservative confidence betting criterion during inference, adding minimal computational overhead (just one additional scoring pass per token). The trade-off is that rejections increase latency slightly (the small model is slower per-token), but the paper's framework lets you tune this trade-off by adjusting the bet threshold—more conservative bets mean more deferrals to safety but slower throughput. For systems already running multiple models (e.g., a main LLM plus a safety classifier), this becomes a natural fit into existing multi-model inference pipelines.
Limitations and When Not to Use This
The paper assumes you have access to a well-aligned small model, which may not exist for all use cases or domains—creating such a proxy requires substantial labeled data and fine-tuning effort. The conservative confidence betting criterion is motivated against ambiguous linguistic cases, but the paper's abstract doesn't reveal how well this criterion generalizes across different domains, model sizes, or types of misalignment (factuality vs. safety vs. refusals). The approach only addresses test-time behavior and doesn't improve the base model itself, meaning you're paying an inference-time cost indefinitely rather than solving alignment at training time. Edge cases around code generation, multilingual outputs, and highly specialized domains (medicine, law) aren't explored, so applicability to vertical-specific models remains unclear. Additionally, the paper appears incomplete (abstract cuts off mid-sentence), so evaluation metrics, actual performance improvements, and computational overhead numbers are not available for assessment.
Research Context
This work builds on a growing line of test-time alignment research that uses smaller models or reward models to steer larger base models without modifying their weights—an attractive alternative to full fine-tuning given model scale. It directly extends recent work on implicit reward methods (which adjust sampling probabilities using a proxy's preferences) and nudging/deferral methods (which hand off generation to the proxy under certain conditions). The paper advances beyond prior work by providing a unifying theoretical framework that exposes what these approaches have in common and clarifying why confidence alone is a poor rejection signal. The research opens a direction toward better rejection criteria for mixture-of-models inference, with potential applications to speculative decoding, ensemble generation, and runtime model selection in resource-constrained environments.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::