The tractability landscape of diffusion alignment: regularization, rewards, and computational primitives

TL;DR

This paper explores the computational primitives needed for reward alignment in diffusion models, analyzing how different distance measures like KL and Wasserstein influence the algorithms and reward classes that can be efficiently implemented.

Contribution

It introduces a primitive-based framework for reward alignment, characterizing the necessary algorithms for different distribution distances and reward types.

Findings

KL-based alignment uses exponential tilts that are efficiently sampleable.

Wasserstein-based alignment employs proximal transport oracles for certain reward classes.

Choice of distance measure determines the computational primitive and reward class tractability.

Abstract

Inference-time reward alignment asks how to turn a pre-trained diffusion model with base law $p$ into a sampler that favors a reward $r$ while remaining close to $p$. Since there is no canonical distributional distance for this closeness constraint, different choices lead to different "reward-aligned" laws and, just as importantly, different algorithmic problems. We develop a primitive-based approach to reward alignment: rather than assuming arbitrary reward-aligned laws can be sampled, we ask which simple algorithmic primitives suffice to implement alignment for non-trivial reward classes. If closeness is measured in KL distance, the target law is $q(x) \propto p(x) \exp(\lambda^{-1}r(x))$. For this setting, we show that linear exponential tilts of the form $q(x)\propto p(x)\exp(\langle \theta, x \rangle)$ -- which according to recent work [MRR26] can be efficiently sampled from -- are a sufficient primitive for aligning to a very broad class of convex low-dimensional rewards. If closeness is measured in Wasserstein distance, the corresponding primitive is a proximal transport oracle: given $x$, solve $\mbox{argmax}_y \{r(y)- \lambda c(x,y)\}$. This oracle can be efficiently implemented for concave or low-dimensional Lipschitz rewards $r(x)=f(Ax)$. Together, these results illustrate that the choice of distribution distance for alignment affects the computational primitive and the tractable reward class.