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

TL;DR

DEFT introduces a unifying framework for conditional diffusion model fine-tuning using Doob's h-transform, enabling faster training and state-of-the-art performance in various inverse problem applications.

Contribution

The paper proposes DEFT, a novel method that fine-tunes a small network to efficiently learn the conditional h-transform, unifying existing approaches and improving speed and performance.

Findings

DEFT achieves up to 1.6× speedup in image reconstruction tasks.

It attains state-of-the-art results on natural and medical images.

Initial experiments show promising results in protein motif scaffolding.

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.