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

TL;DR

This paper analyzes test-time alignment methods using proxy models, identifies issues with confidence-based rejection criteria, and proposes a new conservative confidence bet that improves performance.

Contribution

It introduces a novel rejection criterion based on a conservative confidence bet, addressing limitations of previous confidence-based methods.

Findings

The new rejection criterion outperforms previous methods on multiple datasets.

Both implicit reward and nudging approaches can be reduced to similar graphical models.

Confidence criteria are ill-motivated due to linguistic ambiguities.

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.