DEFT: Efficient Fine-Tuning of Diffusion Models by Learning the Generalised $h$-transform
DEFT leverages Doob's h-transform to fine-tune small networks for fast, high-quality conditional diffusion sampling, outperforming existing methods.
Key Findings
Methodology
This work introduces a unified framework based on Doob’s h-transform, enabling the pre-trained unconditional diffusion model to be conditioned by learning a small neural network that models the time-dependent guidance term. The approach employs score matching and stochastic control principles, allowing efficient fine-tuning with minimal data without backpropagating through the large model. The framework connects Bayesian posterior inference with diffusion processes, providing a theoretical foundation for various existing guidance methods. The small network learns the conditional h-transform, which is then combined with the unconditional model during sampling, resulting in faster and more flexible conditional generation.
Key Results
- On image reconstruction tasks, DEFT achieves up to 1.6× speedup while maintaining state-of-the-art perceptual quality (LPIPS, KID). It outperforms baseline methods such as DPS and RED-diff in natural and medical images, with significant improvements in perceptual metrics. In protein motif scaffolding, DEFT surpasses guidance-based methods, demonstrating versatility across domains.
- Experiments show that fine-tuning on as few as 200 images yields performance close to full-data training, indicating high sample efficiency. The method is robust across linear and non-linear inverse problems, including HDR, phase retrieval, and non-linear deblurring, consistently outperforming existing approaches.
- Ablation studies confirm that smaller training datasets and specific network architectures contribute to performance gains, validating the efficiency and adaptability of DEFT in diverse scenarios.
Significance
This research addresses critical bottlenecks in leveraging large pre-trained diffusion models for conditional tasks, providing a theoretically grounded, computationally efficient framework. It overcomes limitations of hyperparameter sensitivity, high training costs, and API restrictions, making conditional diffusion models more accessible and practical. By enabling rapid fine-tuning with minimal data, DEFT opens new avenues for applications in medical imaging, protein design, and beyond, fostering broader adoption of diffusion-based generative modeling in industry and academia.
Technical Contribution
The paper pioneers the integration of Doob’s h-transform into diffusion models for conditional sampling, establishing a rigorous theoretical basis. It introduces a small neural network architecture optimized via score matching to learn the conditional guidance term efficiently. The approach leverages stochastic control and Bayesian inference to unify various existing guidance methods, providing guarantees on convergence and performance. The framework significantly reduces computational costs by avoiding large-scale backpropagation, enabling fast adaptation to multiple tasks with minimal data, and broadening the scope of diffusion model applications.
Novelty
This is the first systematic application of Doob’s h-transform to diffusion-based conditional sampling, offering a unifying theoretical framework that subsumes existing guidance techniques. Unlike prior heuristic or ad hoc methods, DEFT provides a principled approach grounded in stochastic differential equations and Bayesian inference. Its ability to learn the guidance term via small networks and minimal data distinguishes it from traditional methods that rely on extensive retraining or large datasets, marking a significant step forward in the field.
Limitations
- DEFT’s performance depends on the quality and representativeness of the small fine-tuning dataset; in highly complex or noisy scenarios, the learned guidance may be insufficient.
- The method’s effectiveness in extremely high-dimensional or highly nonlinear tasks remains to be fully validated, and may require further regularization or architecture adjustments.
- While faster than traditional approaches, the initial fine-tuning step still incurs some computational overhead, especially for very large models or real-time applications.
Future Work
Future directions include extending DEFT to multi-condition and multi-task settings, integrating self-supervised learning to reduce label dependence, and exploring adaptive architectures for better generalization. Additionally, applying the framework to video, 3D data, and other modalities could broaden its impact. Further theoretical work on robustness and convergence guarantees, as well as real-world deployment in clinical and industrial settings, are promising avenues.
AI Executive Summary
Diffusion models have revolutionized generative modeling, producing high-quality images, proteins, and other complex data. However, conditioning these models efficiently remains a challenge, especially when large pre-trained models are inaccessible for retraining or fine-tuning. Existing guidance methods often rely on heuristic hyperparameters, extensive retraining, or API access, limiting their practicality. Addressing this, the authors propose DEFT, a novel framework grounded in Doob’s h-transform, which enables rapid, small-network fine-tuning to learn the conditional guidance term. This approach leverages the theoretical foundations of stochastic differential equations and Bayesian inference, unifying various existing methods under a rigorous mathematical umbrella.
The core innovation lies in training a tiny neural network to approximate the time-dependent likelihood gradient, which is then combined with a pre-trained unconditional diffusion model during sampling. This design allows for significant speedups—up to 1.6× in image reconstruction tasks—while maintaining or surpassing state-of-the-art perceptual quality. The experiments span diverse applications, including natural and medical image reconstruction, super-resolution, phase retrieval, HDR, and protein design, demonstrating broad versatility.
The significance of DEFT extends beyond performance metrics. It offers a practical, scalable solution for industries requiring fast, accurate conditional generation without retraining large models. Its sample efficiency, theoretical rigor, and flexibility position it as a foundational advance in diffusion-based generative modeling. Future work aims to expand its scope to multi-condition, multi-modal, and real-time applications, promising a transformative impact across scientific and industrial domains.
Deep Analysis
Background
扩散模型作为生成技术的核心，经历了从噪声逐步去除到高质量图像生成的演变。代表性工作如Denoising Diffusion Probabilistic Models (DDPM)和Score-based Models推动了该领域的发展。近年来，条件生成成为研究热点，诸如DPS、DDRM等方法尝试引入引导机制，但存在训练成本高、参数敏感等问题。随着大规模预训练模型的兴起，如何高效利用这些模型进行条件采样成为焦点。现有方法多依赖昂贵的反向传播或后处理，限制了实际应用的普及。
Core Problem
核心问题在于如何在保持预训练模型性能的基础上，实现高效、灵活的条件采样。传统方法通常需要大量标注数据或反向传播大模型参数，计算成本高且不适用API封闭环境。此外，现有引导技术对超参数敏感，难以在多任务、多模态场景中泛化。如何设计一种理论上统一、计算上高效的框架，满足实际应用需求，是亟待解决的问题。
Innovation
本研究的创新点包括：1）引入Doob的h变换，建立条件扩散的数学基础，提供理论支撑；2）设计微调小网络学习条件引导项，显著降低训练成本；3）结合随机控制和贝叶斯推断，提出多种变换表示，增强模型泛化能力；4）实现无需反向传播大模型参数的快速微调，适应API封闭环境。此框架兼容多任务、多模态，突破了传统方法的局限。
Methodology
- �� 以预训练无条件扩散模型为基础，定义逆扩散过程。• 利用Doob的h变换，将条件采样转化为学习条件引导项的微调问题。• 设计score匹配损失，训练小型网络以学习条件引导。• 采用离散化的DDPM模型，利用少量样本进行微调。• 结合贝叶斯推断，推导多种变换表示，确保模型在不同任务中的适应性。• 在训练中只优化小网络参数，无需反向传播大模型。• 在推理阶段，将学习到的条件引导与预训练模型结合，实现快速条件采样。
Experiments
在图像重建、医学成像、蛋白质设计等多个任务中，采用ImageNet、LoDoPab-CT等公开数据集。对比DPS、RED-diff等先进方法，指标包括PSNR、SSIM、LPIPS、KID。设置不同的样本量（如200张图像）进行微调，验证样本效率。采用不同的采样步数（如100步）进行推理，评估速度与质量。还进行了消融实验，分析网络结构和样本数量对性能的影响。
Results
DEFT在多项任务中表现优异，重建任务中速度提升1.6倍，感知指标优于对比方法。蛋白质骨架重建中超越指导方法。少样本微调（200图像）即可达到接近全数据训练的性能，验证了极高的样本效率。非线性任务（如HDR、相位检索）中也展现出优越性能，显示出强大的泛化能力。消融分析证实微调样本量和网络设计对性能的关键作用。
Applications
该方法适用于医学影像重建、工业检测、蛋白质设计等领域，尤其在数据有限或API封闭环境中表现出巨大优势。可快速适应不同任务，无需大规模标注或反向传播，极大降低部署门槛。未来有望结合多模态数据、多任务学习，推动智能制造、药物研发等行业的创新。
Limitations & Outlook
目前DEFT在极端噪声或复杂非线性任务中表现仍有限，模型鲁棒性有待提升。微调样本依赖于任务相关性，泛化能力在高维空间仍需验证。模型在超大模型或极端场景中仍存在计算瓶颈。未来需结合正则化和多任务优化策略，增强模型的适应性与鲁棒性。
Plain Language Accessible to non-experts
想象你是一位厨师，手里有一份食谱（预训练模型），可以做出各种美味菜肴（生成内容）。但有时候你需要根据客人的特殊要求（条件）调整菜肴，比如少盐或多辣。传统方法就像每次都要重新学习整个食谱，既麻烦又慢。而这篇研究提出的方法，就像只微调厨房里的一小块调料罐（小网络），让它学会根据不同需求调整味道。这样，你就可以快速满足不同客人的口味，同时保持菜肴的高品质。这种方式既节省时间，又灵活多变，适合快节奏的餐厅（应用场景）。
ELI14 Explained like you're 14
你知道在游戏里调角色装备吗？有时候你想让角色变得更强或更快，但不想重新开始一遍。传统的方法就像每次都要重新打怪升级，非常耗时间。而这项新技术就像给角色装上一个神奇的背包，只需要调整一下里面的小按钮，就能让角色变得更厉害，速度也快很多。它用一种叫做h变换的神奇魔法，让模型学会根据不同的需求快速变化。这样一来，不管是修复图片、增强细节，还是设计蛋白质，都能更快更好地完成。就像给游戏角色加了个超级技能，既省时又灵活！
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.