Research map / Image generation / Fast diffusion sampling

Learnable Chernoff Baselines for Provable Inference-Time Alignment

NeurIPS 2026 · Image generation › Fast diffusion sampling · approach: Distillation for few-step generation

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes Learnable Chernoff Baselines (LCBs), adaptive rejection-sampling envelopes learned from soft-value estimates, to approximately sample reward-tilted diffusion processes with fewer base-model queries and TV guarantees.

Problem

Inference-time reward-guided alignment of pretrained (diffusion or masked language) generative models with only black-box sampling access; existing approaches like best-of-N and SMC are costly or lack non-asymptotic guarantees, and naive rejection sampling needs exponentially many proposals at low t

Approach

Learn state-dependent baseline functions (Chernoff-certified via an MGF objective) that serve as adaptive acceptance envelopes for per-step rejection sampling of exponentially tilted transition kernels, giving a tunable approximation-compute tradeoff parameter delta and TV bounds to the ideal aligned model.

Key results

• On the 2-d Gaussian mixture, LCB matches rejection sampling distribution while using about 5.5 effective proposals versus 69 for RS and 40 for Best-of-N.
• On LLaDA-8B at alpha=0.3, LCB achieves reward 0.885 with about 38.7 proposals versus 64 for Best-of-2 (reward 0.893), a roughly 40% reduction in queries.
• On LLaDA at alpha=0.1, LCB reaches reward 0.917 and 81.8% perfect rate with 51 proposals, versus RS 0.923 with 403 proposals.

Surprising findings

• At high temperature (alpha=0.5), LCB uses 16.3% more queries than unguided sampling while still improving reward metrics.
• Naive rejection sampling proposal counts blow up at low temperature (about 403 proposals at alpha=0.1), motivating adaptive baselines.

Limitations

• LCB quality is bounded by soft-value estimation error, which the TV bounds carry through.
• Worst-case coverage assumption and sub-Gaussian assumptions are needed for the theoretical guarantees; empirical LLaDA evaluation uses a single simple reward and prompt.
• Text quality is only checked via reward statistics and simple metrics; no distance between text distributions can be exactly computed.

Experimental setup

Datasets: LLaDA-8B generated trajectories (first-letter constrained three-sentence stories) 2-d Gaussian mixture DDPM samples

Benchmarks: 2-d Gaussian mixture tail reward LLaDA first-letter constraint task

Metrics: average reward perfect-sample rate effective number of proposals histogram match to ideal rejection sampling

Baselines: Rejection Sampling Best-of-N Best-of-2 Unguided LLaDA

Builds on: classifier guidance best-of-N sequential Monte Carlo tilted empirical risk minimization DDPM

Closest papers by meaning

• InfAlign: Inference-aware language model alignment (ICML 2025 · similarity 0.90)
• Is Best-of-N the Best of Them? Coverage, Scaling, and Optimality in Inference-Time Alignment (ICML 2025 · similarity 0.90)
• Inference-time Alignment in Continuous Space (NeurIPS 2025 · similarity 0.89)
• $\Psi$-Sampler: Initial Particle Sampling for SMC-Based Inference-Time Reward Alignment in Score Models (NeurIPS 2025 · similarity 0.88)
• Diffusion Tree Sampling: Scalable inference‑time alignment of diffusion models (NeurIPS 2025 · similarity 0.87)
• On the Limits of Test-Time Compute: Sequential Reward Filtering for Better Inference (ICML 2026 · similarity 0.87)