Learnable Chernoff Baselines for Provable Inference-Time Alignment

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, which reveal the dimension-independent quantities governing the tradeoff between accurate sampling and inference compute. We empirically demonstrate in both continuous and discrete diffusion settings that LCB sampling closely matches ideal rejection sampling, but uses substantially fewer queries to the pretrained model. Our experiments also include real-world data experiments on a diffusion large language model.