Risk-Averse Inference-Time Diffusion Model Alignment with Finite-Sample Guarantees
Yuchen Liang, Xian Yu
Abstract
Inference-time alignment steers diffusion models toward downstream objectives without costly fine-tuning, but existing methods typically optimize expected cost (or reward) without explicitly controlling tail risk. In this work, we propose a risk-averse framework for inference-time alignment of diffusion models that balances expected cost and Conditional Value-at-Risk (CVaR) under trajectory-level Kullback-Leibler regularization. We characterize the optimal path distribution and reverse transitions as functions of the CVaR threshold, reducing the alignment problem to a one-dimensional optimization problem over the threshold. We show that this one-dimensional optimization problem can be nonconvex in general, establish sufficient conditions for convexity, and introduce quadratic regularization to enable efficient derivative-free ternary search. Under suitable assumptions, we provide finite-sample guarantees for threshold optimization and value estimation, and combine them into a terminal total-variation bound relative to the regularized target under exact resampling. Experiments on synthetic data and image generation using Stable Diffusion demonstrate improved tail performance over risk-neutral baselines while maintaining competitive average cost.