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

TL;DR

This paper introduces DOIT, a training-free, efficient method for adapting pre-trained diffusion models to high-reward distributions using Doob's $h$-transform, with theoretical guarantees and empirical success on RL benchmarks.

Contribution

We propose a novel, training-free adaptation technique for diffusion models that leverages Doob's $h$-transform, enabling efficient, non-differentiable reward optimization with theoretical guarantees.

Findings

Outperforms state-of-the-art baselines on D4RL benchmarks.

Provides theoretical convergence guarantees for the adaptation process.

Maintains sampling efficiency while achieving high-reward distribution alignment.

Abstract

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 DOIT (Doob-Oriented Inference-time Transformation), a training-free and computationally efficient adaptation method that applies to generic, non-differentiable rewards. The key framework underlying our method 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 enables efficient simulation-based computation without modifying the pre-trained model. Theoretically, we establish a high probability convergence guarantee to the target high-reward distribution via characterizing the approximation error in the dynamic Doob's correction. Empirically, on D4RL offline RL benchmarks, our method consistently outperforms state-of-the-art baselines while preserving sampling efficiency.