DEFT: Efficient Fine-Tuning of Diffusion Models by Learning the Generalised $h$-transform

TL;DR

DEFT introduces a unified Doob's $h$-transform framework for conditional diffusion models and shows that a small, trainable $h$-transform network can be fine-tuned to learn the conditional score $\nabla_{\bm{x}} \ln p_t(\bm{y}|\bm{x})$ while keeping a large unconditional diffusion model fixed. It provides a simulation-free, score-matching objective and a stochastic-control interpretation, enabling both amortised and zero-shot conditioning strategies. Empirically, DEFT delivers faster conditional sampling and strong perceptual and reconstruction performance across image inpainting, super-resolution, HDR, phase retrieval, non-linear deblurring, CT, and protein motif scaffolding, often outperforming state-of-the-art baselines. The approach is particularly advantageous when direct backpropagation through the large base model is infeasible (e.g., API-restricted settings) and demonstrates substantial speedups and data-efficient conditioning with practical impact for inverse problems and bio-design.

Abstract

Generative modelling paradigms based on denoising diffusion processes have emerged as a leading candidate for conditional sampling in inverse problems. In many real-world applications, we often have access to large, expensively trained unconditional diffusion models, which we aim to exploit for improving conditional sampling. Most recent approaches are motivated heuristically and lack a unifying framework, obscuring connections between them. Further, they often suffer from issues such as being very sensitive to hyperparameters, being expensive to train or needing access to weights hidden behind a closed API. In this work, we unify conditional training and sampling using the mathematically well-understood Doob's h-transform. This new perspective allows us to unify many existing methods under a common umbrella. Under this framework, we propose DEFT (Doob's h-transform Efficient FineTuning), a new approach for conditional generation that simply fine-tunes a very small network to quickly learn the conditional $h$-transform, while keeping the larger unconditional network unchanged. DEFT is much faster than existing baselines while achieving state-of-the-art performance across a variety of linear and non-linear benchmarks. On image reconstruction tasks, we achieve speedups of up to 1.6$\times$, while having the best perceptual quality on natural images and reconstruction performance on medical images. Further, we also provide initial experiments on protein motif scaffolding and outperform reconstruction guidance methods.

Figures (15)

• Figure 1: DEFT reverse diffusion setup. The pre-trained unconditional diffusion model $s_t^\theta$ and the fine-tuned $h$-transform $h_t^\phi$ are combined at every sampling step. We propose a special network to parametrise the $h$-transform including the guidance term $\nabla_{\hat{{\bm{x}}}_0} \mathop{\mathrm{ln}}\nolimits p({{\bm{y}}}|\hat{{\bm{x}}}_0)$ as part of the architecture. Here $\hat{{\bm{x}}}_0$ denotes the unconditional denoised estimate given $s_t^\theta({\bm{x}}_t)$. During training, we only need to fine-tune $h_t^{\phi}$ (usually $4$-$9\%$ the size of $s_t^{\theta}$) using a small dataset of paired measurements, keeping $s_\theta^t$ fixed. During sampling, we do not need to backpropagate through either model, resulting in speed-ups during evaluation.
• Figure 2: Results for inpainting. We show the ground truth with the inpainting mask superimposed.
• Figure 3: Results for non-linear deblurring. We show both the ground truth, the measurements and samples for DPS, RED-diff and DEFT. DEFT is able to reconstruct high-quality images.
• Figure 4: Reconstructions for computed tomography on LoDoPab-CT
• Figure 5: Comparison of DPS, DEFT and amortised training for motif scaffolding for 12 contiguous targets. 4% and 9% are the relative sizes of the h-transform compared to the unconditional model.

Theorems & Definitions (13)

• Proposition 2.1
• Proposition 2.2
• Theorem 3.1
• Corollary C.1
• proof
• proof
• Proposition G.1
• proof
• Remark 1
• Proposition G.2