Learnable Chernoff Baselines for Inference-Time Alignment

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.

• The paper introduces Learnable Chernoff Baselines (LCBs) that construct state-dependent acceptance envelopes for efficient, reward-guided inference-time alignment.
• It leverages moment-generating function objectives and Chernoff bounds to provide rigorous control over sampling error and computational cost in both continuous and discrete settings.
• Empirical results show that LCBs drastically reduce proposal complexity while achieving reward distributions comparable to traditional methods like rejection sampling and Best-of-N.

Learnable Chernoff Baselines for Inference-Time Alignment

Problem Formulation and Motivation

The paper addresses the critical challenge in generative modeling of efficiently aligning outputs at inference time, specifically when a pretrained model must be adapted to optimize a downstream reward function under a KL-regularized constraint. The canonical setting involves sampling from a distribution

$$p^*(x_0) \propto \exp(r(x_0)/\alpha) \, p^{\mathrm{pre}}(x_0)$$

where $r$ is a reward, $\alpha$ is a temperature, and $p^{\mathrm{pre}}$ is the base model's output law. While reinforcement learning (RL) finetuning classically targets this distribution, it necessitates open-weight access and substantial computational cost. Inference-time alignment, by contrast, aims to sample directly from $p^*$ given only black-box access, enabling adaptation without retraining.

Existing inference-time approaches (e.g., SMC/PF, Best-of-$N$ selection, classifier guidance) suffer from compute inefficiency, architectural dependence, or lack robust theoretical guarantees. This work proposes a new method—Learnable Chernoff Baselines (LCBs)—that generalizes adaptive rejection sampling and offers rigorous bounds on both accuracy and computational efficiency, applicable to both continuous and discrete generative models.

LCB Framework: Theory and Algorithms

The paper introduces LCBs as adaptive baseline functions that serve as state-dependent acceptance envelopes for rejection sampling. The paradigm starts with estimation of soft-value functions $v_t$ (log-expected exponentiated reward conditioned on intermediate state). The aligned target model is realized as a Markov tilt at each timestep, and the sampling problem reduces to rejection sampling the transition kernels with proportional acceptance probabilities.

However, naive rejection sampling with a global uniform baseline is inefficient—requiring up to $e^{2B}$ proposals per accepted transition in the worst case. LCBs accelerate this process by learning tight, local upper bounds, $B_{t+1}(x_{t+1})$, on $v_t$ over successor states, with a user-tunable tail parameter $\delta$ that controls both the rate of approximation and proposal complexity.

The method is formalized by leveraging Chernoff bounds to guarantee that the baseline covers the value function with probability at most $\delta$. LCBs are optimized by minimizing a moment-generating function-based objective, with explicit uniform convergence bounds derived for empirical risk minimization. The result is a rejection sampling algorithm that is model-agnostic, efficient, and end-to-end theoretically controlled in total variation distance.

Figure 1: Guided language diffusion trajectories, with LCB (solid) tracking the soft-value function (dotted), providing more efficient state-dependent baselines than rejection sampling's global envelope.

Theoretical Guarantees

A key contribution is rigorous control of the sampling error relative to the ideal value-tilted model:

$$d_{TV}(\hat q(x_0), \hat p(x_0)) \lesssim \delta \sum_t e^{J^*_t + 2\epsilon_0}$$

where $J^*_t$ is determined by the LCB objective and $\epsilon_0$ reflects statistical risk from baseline training.

For sub-Gaussian value estimators, the bound tightens further:

$$d_{TV}(\hat q(x_0), \hat p(x_0)) \lesssim \delta \sum_t e^{\sigma_t \sqrt{\log(1/\delta) + 2\epsilon_0}}$$

and, in DDPM-style Gaussian mixture targets, $\sigma_t = O(\beta_t + \tilde \beta_t)$, linking error to the local variance structure of the generative process.

The paper also proves that composing baseline-induced errors with value estimation errors yields an end-to-end bound:

$$d_{TV}(\hat q(x_0), p^*(x_0)) \leq d_{TV}(\hat q(x_0), \hat p(x_0)) + o_n(1)$$

where $n$ is the sample count for value function pretraining, reflecting the amortized nature of baseline and value estimation costs.

Empirical Results

Experiments span both continuous (Gaussian mixtures under DDPM) and discrete (large language diffusion—LLaDA) regimes.

Gaussian Mixture DDPM

LCBs match the alignment strength of Rejection Sampling (RS) and Best-of-$N$ (BoN) selection using orders-of-magnitude fewer proposals—$8\%$ of RS and $14\%$ of BoN's compute. Effective alignment is achieved even in ill-conditioned, low-temperature regimes, demonstrating the practical efficiency of adaptive baselines.

Figure 2: LCB achieves comparable alignment to RS and BoN while requiring significantly fewer proposals; chaining LCB with BoN further boosts reward with minimal compute overhead.

Figure 3: Proposal complexity comparison as a function of method and temperature, showing substantial gains for LCB versus RS and BoN.

LLaDA Language Diffusion

In high-temperature language generation, LCBs deliver statistically indistinguishable reward distributions compared to RS, while reducing model query complexity by $25$–$87\%$ depending on alignment temperature. The difference in text quality metrics between LCB and RS is not statistically significant across key attributes (sentence structure, vocabulary), confirming that reward amplification does not degrade output diversity.

As $\delta$ is tightened, LCB's proposal complexity remains strictly sublinear in $1/\delta$, exhibiting diminishing returns that are easily tunable in deployment.

Figure 4: Reward distributions induced by RS and LCB are nearly identical across alignment temperatures, with LCB requiring fewer proposals.

Figure 5: Diminishing returns in reward mass as $\delta$ decreases, but proposal complexity increases only modestly for smaller $\delta$.

Analysis of Baseline Learning and Coverage

LCBs are learned sequentially—each is fitted via ERM using samples from the guided trajectory distributions, ensuring mutual compatibility and avoiding covariate shift. High-probability coverage of state space is explicitly validated, with empirical verification that nearly all states satisfy the baseline's upper bound property, even at challenging, noisier timesteps.

Figure 6: Empirical coverage statistics for LCBs in LLaDA; most states robustly satisfy baseline coverage for all their successors.

Trajectory and Value Function Evolution

Analysis of guided trajectories in language diffusion reveals near-optimal alignment between learned baselines and the value function throughout the trajectory, supporting the theoretical connection between adaptive acceptance probabilities and accurate reward-guided sampling.

Figure 7: Unrolled trajectories in LLaDA, with LCB baselines tracking value function evolution across rewards and timesteps.

Proposal Complexity Scaling and Distributional Fidelity

LCB-based sampling rigorously controls both the proposal complexity and the total variation error to the target, enabling fine-grained tuning of the alignment vs compute budget. Distributional comparisons (scatter plots, histograms) show no observable discrepancy between LCB and RS at matched reward levels.

Figure 8: Scatter plots and histograms for Gaussian mixture models, showing close match between LCB, RS, and BoN when tuned for alignment.

Practical and Theoretical Implications

Practically, LCBs present a scalable, flexible inference-time alignment method suited for deployment in settings with closed-weight models or costly queries (API-based LLMs, large diffusion models). Theoretically, the LCB formulation unifies adaptive envelope rejection sampling with moment generating function analysis for controlled stochastic approximation, and exposes a new tradeoff axis (tail parameter $\delta$) for balancing fidelity and efficiency.

LCB's success further motivates improvements in soft-value function estimation and opens opportunities for prompt-conditioned, task-adaptive reward modeling. The empirical coverage results indicate that state-dependent adaptive baselines suffice for most practical generative tasks, where global upper bounds would be prohibitively loose.

Future Directions

The development of LCBs suggests several avenues in both inference-time alignment and reward-guided generation:

• Extension to prompt-conditioned rewards: Training value and baseline networks conditioned on prompt/context could yield universal alignment modules transferable across tasks.
• Robust reward design and adversarial safeguards: Frameworks for controlled baseline learning could help mitigate risk from pathological or adversarial reward functions in model alignment.
• Integration with discrete/continuous generative settings: LCBs' architectural agnosticism offers compatibility with new diffusion and autoregressive regimes, subject to further exploration in high-dimensional, multimodal domains.

Conclusion

Learnable Chernoff Baselines offer a provably efficient, adaptive mechanism for inference-time reward-guided alignment in generative models, achieving near-ideal distributional fidelity at a fraction of the proposal complexity of prior methods. By leveraging tight moment-generating bounds and empirical baseline fitting, LCBs enable scalable, accurate, and practical alignment—optimally balancing compute and approximation budgets—with broad applicability to both language and vision generative systems.