Iterative Importance Fine-tuning of Diffusion Models

Diffusion models are an important tool for generative modelling, serving as effective priors in applications such as imaging and protein design. A key challenge in applying diffusion models for downstream tasks is efficiently sampling from resulting posterior distributions, which can be addressed using Doob’s $h$-transform. This work introduces a self-supervised algorithm for fine-tuning diffusion models by learning the optimal control, enabling amortised conditional sampling. Our method iteratively refines the control using a synthetic dataset resampled with path-based importance weights. We demonstrate the effectiveness of this framework on class-conditional sampling, inverse problems and reward fine-tuning for text-to-image diffusion models.

💡 Research Summary

The paper tackles the problem of efficiently sampling from a tilted distribution p̃(x) ∝ p_data(x)·exp(r(x)/λ) using diffusion models. This tilted distribution arises in many downstream tasks such as Bayesian inverse problems, class‑conditional generation, and reward‑based fine‑tuning. Theoretically, sampling from p̃ can be achieved by adding the optimal control u*_t = ∇_x ln h_r_t(x) to the reverse SDE of a diffusion model (Doob’s h‑transform). However, u*_t is intractable in practice.

The authors propose a self‑supervised, iterative importance‑fine‑tuning algorithm that learns an approximation of u*_t without ever needing samples from p̃. The procedure consists of three steps repeated until convergence:

1. Sampling with the current control – Starting from an initial control u_0 (often zero), they generate N trajectories x^{(n)}