On the Rejection Criterion for Proxy-based Test-time Alignment
For practitioners of large language model alignment, this work provides a theoretically motivated improvement to proxy-based test-time alignment methods.
The paper unifies two test-time alignment methods (implicit reward and nudging) under a common graphical model, differing only in rejection criterion. It proposes a novel conservative confidence bet criterion that outperforms prior work on multiple datasets.
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.