Training-Free Adaptation of Diffusion Models via Doob's $h$-Transform

Adaptation methods have been a workhorse for unlocking the transformative power of pre-trained diffusion models in diverse applications. Existing approaches often abstract adaptation objectives as a reward function and steer diffusion models to generate high-reward samples. However, these approaches can incur high computational overhead due to additional training, or rely on stringent assumptions on the reward such as differentiability. Moreover, despite their empirical success, theoretical justification and guarantees are seldom established. In this paper, we propose $\texttt{DOIT}$ ($\textbf{D}$oob-$\textbf{O}$riented $\textbf{I}$nference-time $\textbf{T}$ransformation), a training-free adaptation method for generic, non-differentiable rewards. We develop two variants of this method: $\texttt{DOIT-Proto}$, a full-rollout simulation-based algorithm and $\texttt{DOIT-Practical}$, an efficient surrogate algorithm. The key idea is a measure transport formulation that seeks to transport the pre-trained generative distribution to a high-reward target distribution. We leverage Doob's $h$-transform to realize this transport, which induces a dynamic correction to the diffusion sampling process and admits simulation-based computation without modifying the pre-trained model. Theoretically, we establish a high-probability convergence guarantee to the target high-reward distribution for $\texttt{DOIT-Proto}$ by characterizing the approximation error in the dynamic Doob correction. Empirically, $\texttt{DOIT-Practical}$ consistently outperforms state-of-the-art baselines on D4RL offline RL benchmarks while preserving sampling efficiency.

Modern generative models, such as diffusion models and flow-based models, can create realistic images and useful robot actions by gradually turning random noise into structured outputs. These models, are powerful, but people often want them to do more than imitate the data they were trained on. In many applications, we want generated outputs to score highly on a specific goal, such as producing a more appealing image or choosing a better robot action. Existing ways to steer these models often require costly extra training, or they only work when the goal can be smoothly differentiated, which is not true for many real-world reward signals. We propose $\texttt{DOIT}$, a method that adapts a pre-trained diffusion model while it is generating samples, without changing the model itself. The method gently redirects the model’s sampling path toward outputs that are more likely to satisfy the desired goal, while still staying close to what the original model can produce. $\texttt{DOIT}$ is inspired by a classical probability tool called Doob’s $h$-transform, which tells us how to modify a random process when we want it to end in a desirable event. We develop a principled prototype version of $\texttt{DOIT}$ with theoretical guarantees, as well as a practical version of $\texttt{DOIT}$ that keeps sampling efficient. In experiments on image generation and offline reinforcement learning, the practical version shifts samples toward higher rewards and improves average performance compared with baselines. This makes it easier to reuse powerful pre-trained diffusion models in new tasks where retraining is expensive and reward signals may be black-box or non-differentiable.