Learnable Chernoff Baselines for Inference-Time Alignment
Sunil Madhow, Yuchen Liang, Ness Shroff, Yingbin Liang, Yu-Xiang Wang
TL;DR
The paper addresses inference-time alignment for pretrained generative models under KL regularization by introducing Learnable Chernoff Baselines (LCBs). LCBs provide adaptive, model-agnostic rejection sampling envelopes that tilt transition kernels via learned soft-value baselines, enabling efficient sampling with provable total-variation guarantees relative to ideal tilted-rejection schemes. The framework yields end-to-end TV bounds, supports both continuous diffusion and discrete language diffusion, and demonstrates substantial compute savings in Gaussian mixtures and large-language diffusion tasks (LLaDA) without sacrificing alignment quality. This approach offers a scalable, principled path to safer, reward-guided generation at inference time, with practical impact on downstream tasks and potential dual-use considerations.
Abstract
We study inference-time reward-guided alignment for generative models. Existing methods often rely on either architecture-specific adaptations or computationally costly inference procedures. We introduce Learnable Chernoff Baselines (LCBs) as a method for efficiently and approximately sampling from the exponentially tilted kernels that arise from KL-regularized reward alignment. Using only black-box sampling access to the pretrained model, LCBs implement a form of rejection sampling with adaptively selected acceptance probabilities, which allows fine-grained control over inference-compute scaling. We establish total-variation guarantees to the ideal aligned model, and demonstrate in both continuous and discrete diffusion settings that LCB sampling closely matches ideal rejection sampling while using substantially fewer queries to the pretrained model.