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

TL;DR

提出DEFT,通过学习广义h变换快速微调扩散模型,提升条件采样效率与性能。

cs.LG 🔴 高级 2024-06-04 50 次浏览
Alexander Denker Francisco Vargas Shreyas Padhy Kieran Didi Simon Mathis Vincent Dutordoir Riccardo Barbano Emile Mathieu Urszula Julia Komorowska Pietro Lio
生成模型 扩散模型 条件采样 微调 贝叶斯变换

核心发现

方法论

本文基于Doob的h变换,提出DEFT框架,将预训练的无条件扩散模型通过微调小型网络学习条件h变换,从而实现高效条件采样。核心算法包括利用score匹配和贝叶斯推断,结合特定网络结构,快速适应不同条件任务。该方法无需反向传播大模型参数,显著提升训练与推理速度。具体实现上,采用离散化的DDPM模型,利用小样本数据进行微调,优化目标为学习条件引导项。通过连接随机控制理论,提出多种变换表示,确保模型在多任务中的泛化能力。

关键结果

  • 在图像重建任务中,DEFT实现了最高的感知质量和重建性能,速度提升达1.6倍,且在自然图像和医学图像上表现优异。在图像修复、超分辨率、相位检索等非线性任务中,DEFT均优于传统方法,尤其在感知指标(LPIPS、KID)上表现出色。蛋白质骨架重建中,DEFT也超越了指导重建方法,验证其多领域适用性。
  • 在大规模实验中,DEFT微调仅需少量样本(如200张ImageNet图像),即能达到接近全数据训练的性能,显示出极高的样本效率。与现有的条件扩散方法(如DPS、RED-diff)相比,DEFT在训练时间和推理速度上具有明显优势,适合实际应用场景。
  • 通过对不同任务的消融分析,验证了微调样本量、网络结构和目标函数对性能的影响,进一步证明了DEFT的灵活性和鲁棒性。

研究意义

该研究突破了利用大规模预训练无条件扩散模型进行条件采样的瓶颈,提供了理论上统一的框架,解决了现有方法对超参数敏感、训练成本高、依赖API限制等问题。其快速微调机制极大降低了条件生成的门槛,为医学影像、蛋白质设计等领域带来新机遇。通过引入Doob的h变换,模型在保持高质量生成的同时,显著提升了效率和适应性,有望推动扩散模型在实际应用中的普及与创新。

技术贡献

本文在理论上将Doob的h变换引入扩散模型条件采样,提出微调小网络学习条件引导项,结合随机控制与贝叶斯推断,建立了统一的数学框架。技术上,设计了高效的score匹配损失,支持少样本微调,避免大模型反向传播,显著降低计算成本。架构创新包括融合无条件模型的denoising估计与条件引导,增强模型的非线性表达能力。实验验证了在多任务、多模态场景中的优越性能,推动了条件生成的理论与工程发展。

新颖性

首次将Doob的h变换系统性引入扩散模型条件采样,提出微调小网络学习条件引导,突破了依赖大模型反向传播的限制。相较于传统的指导方法(如梯度指导、后验引导),该框架提供了理论基础和统一视角,兼容多种任务,具有较强的泛化能力。创新点在于结合随机控制、贝叶斯推断与score匹配,提出高效的训练与推理流程,极大提升了条件采样的速度与质量。

局限性

  • DEFT依赖少量标注样本,对于极端少样本或无标注场景可能效果有限,且微调样本选择可能影响性能。
  • 在高维复杂任务中,微调网络可能不足以捕获所有条件信息,存在泛化瓶颈。
  • 模型在极端非线性或噪声极大情况下的鲁棒性仍需验证,未来需结合更强的正则化策略。

未来方向

未来将探索多任务、多条件联合微调,提升模型在复杂场景中的适应性。还计划结合自监督学习和强化学习,增强模型对少样本和无标签数据的利用能力。此外,将扩展到更广泛的模态(如视频、三维数据)和多任务学习,推动条件扩散模型的工业化应用。

AI 总览摘要

扩散模型作为生成领域的核心技术之一,近年来在图像、蛋白质设计等多领域展现出巨大潜力。然而,现有条件采样方法普遍面临训练成本高、超参数敏感、依赖大模型反向传播的挑战。本文提出的DEFT框架,基于Doob的h变换,创新性地将预训练的无条件扩散模型微调为条件生成器。通过学习条件引导项的小型网络,DEFT实现了高速、精确的条件采样,显著优于传统方法。

在图像重建、超分辨率、医学影像等任务中,DEFT不仅提升了感知质量,还实现了1.6倍的速度提升。蛋白质骨架重建实验也验证了其跨领域的适用性。该方法的核心在于结合随机控制理论,利用score匹配和贝叶斯推断,建立了统一的数学框架,极大简化了条件采样的复杂性。

此外,DEFT的样本效率极高,仅需少量标注样本便能达到优异性能,为实际应用提供了可能。未来,结合多任务学习和自监督技术,DEFT有望在更多复杂场景中实现快速、精准的条件生成,推动扩散模型的广泛落地。

深度分析

研究背景

扩散模型作为生成技术的前沿,经过逐步发展,从最早的噪声逐步去除到高质量图像生成,代表性工作包括Denoising Diffusion Probabilistic Models (DDPM)和Score-based Models。近年来,条件生成成为热点,诸如DPS、Denoising Diffusion Restoration Models (DDRM)等方法尝试引入引导机制,但存在训练成本高、参数敏感等问题。随着大规模预训练模型的兴起,如何高效利用这些模型进行条件采样成为研究焦点。现有方法多依赖反向传播或昂贵的后处理,限制了实际应用的普及。

核心问题

核心问题在于如何在保持预训练模型性能的基础上,实现高效、灵活的条件采样。传统方法通常需要大量标注数据或反向传播大模型参数,计算成本高且不适用API封闭环境。此外,现有引导技术对超参数敏感,难以在多任务、多模态场景中泛化。如何设计一种理论上统一、计算上高效的框架,满足实际应用需求,是亟待解决的问题。

核心创新

本研究的创新点包括:1)引入Doob的h变换,建立条件扩散的数学基础,提供理论支撑;2)设计微调小网络学习条件引导项,显著降低训练成本;3)结合随机控制和贝叶斯推断,提出多种变换表示,增强模型泛化能力;4)实现无需反向传播大模型参数的快速微调,适应API封闭环境。此框架兼容多任务、多模态,突破了传统方法的局限。

方法详解

  • �� 以预训练无条件扩散模型为基础,定义逆扩散过程。• 利用Doob的h变换,将条件采样转化为学习条件引导项的微调问题。• 设计score匹配损失,训练小型网络以学习条件引导。• 采用离散化的DDPM模型,利用少量样本进行微调。• 结合贝叶斯推断,推导多种变换表示,确保模型在不同任务中的适应性。• 在训练中只优化小网络参数,无需反向传播大模型。• 在推理阶段,将学习到的条件引导与预训练模型结合,实现快速条件采样。

实验设计

在图像重建、医学成像、蛋白质设计等多个任务中,采用ImageNet、LoDoPab-CT等公开数据集。对比DPS、RED-diff等先进方法,指标包括PSNR、SSIM、LPIPS、KID。设置不同的样本量(如200张图像)进行微调,验证样本效率。采用不同的采样步数(如100步)进行推理,评估速度与质量。还进行了消融实验,分析网络结构和样本数量对性能的影响。

结果分析

DEFT在多项任务中表现优异,重建任务中速度提升1.6倍,感知指标优于对比方法。蛋白质骨架重建中超越指导方法。少样本微调(200图像)即可达到接近全数据训练的性能,验证了极高的样本效率。在非线性任务(如HDR、相位检索)中,也展现出优越性能,显示出强大的泛化能力。消融分析证实微调样本量和网络设计对性能的关键作用。

应用场景

该方法适用于医学影像重建、工业检测、蛋白质设计等领域,尤其在数据有限或API封闭环境中表现出巨大优势。可快速适应不同任务,无需大规模标注或反向传播,极大降低部署门槛。未来有望结合多模态数据、多任务学习,推动智能制造、药物研发等行业的创新。

局限与展望

目前DEFT在极端噪声或复杂非线性任务中表现仍有限,模型鲁棒性有待提升。微调样本依赖于任务相关性,泛化能力在高维空间仍需验证。此外,虽然速度提升显著,但在超大模型或极端场景中仍存在计算瓶颈。未来需结合正则化和多任务优化策略,增强模型的适应性与鲁棒性。

通俗解读 非专业人士也能看懂

想象你是一位厨师,手里有一份食谱(预训练模型),可以做出各种美味菜肴(生成内容)。但有时候你需要根据客人的特殊要求(条件)调整菜肴,比如少盐或多辣。传统方法就像每次都要重新学习整个食谱,既麻烦又慢。而这篇研究提出的方法,就像只微调厨房里的一小块调料罐(小网络),让它学会根据不同需求调整味道。这样,你就可以快速满足不同客人的口味,同时保持菜肴的高品质。这种方式既节省时间,又灵活多变,适合快节奏的餐厅(应用场景)。

简单解释 像给14岁少年讲一样

你知道在游戏里调角色装备吗?有时候你想让角色变得更强或更快,但不想重新开始一遍。传统的方法就像每次都要重新打怪升级,非常耗时间。而这项新技术就像给角色装上一个神奇的背包,只需要调整一下里面的小按钮,就能让角色变得更厉害,速度也快很多。它用一种叫做h变换的神奇魔法,让模型学会根据不同的需求快速变化。这样一来,不管是修复图片、增强细节,还是设计蛋白质,都能更快更好地完成。就像给游戏角色加了个超级技能,既省时又灵活!

原文摘要

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.

cs.LG