Inference-Time Alignment of Diffusion Models: Efficient Sampling and Tail-Risk Control

Abstract

Diffusion generative models produce realistic images, text, and other complex data, but downstream applications require specific desirable properties. Inference-time alignment steers pretrained diffusion models without fine-tuning, yet can be computationally expensive and leave undesirable tail outcomes uncontrolled. This talk addresses two complementary goals for inference-time alignment of diffusion models: efficiency and reliability. For efficiency, we introduce Learnable Chernoff Baseline (LCB) for sampling from the exponentially tilted transitions arising in Kullback-Leibler-regularized reward (or cost) alignment. Using black-box sampling access to the pretrained model, LCB adapts rejection-sampling thresholds to balance computational cost and accuracy. We establish total-variation error bounds relative to the ideal aligned distribution. Experiments on synthetic Gaussian mixtures and text generation show that LCB achieves alignment performance comparable to exact rejection sampling while requiring substantially fewer queries to the pretrained model. For reliability, we develop a risk-averse framework balancing expected cost and Conditional Value-at-Risk (CVaR). Characterizing the optimal aligned distribution reduces the problem to a one-dimensional threshold optimization. We study its convexity and introduce quadratic regularization for efficient derivative-free optimization. Under suitable assumptions, we further derive finite-sample guarantees for threshold and value estimation, which yield an end-to-end sample-complexity guarantee relative to the regularized target under exact tilted sampling. Experiments on synthetic Cauchy mixtures and image generation demonstrate improved tail performance, with the mean-tail trade-off varying across samplers and experimental settings.