以下是论文《Amortized Guidance for Image Inpainting with Pretrained Diffusion Models》(AID)的通俗解释,辅以生动的类比。
问题:“修补”困境
想象你有一位大师级画家(预训练扩散模型),他擅长从零开始创作壮丽的风景画。现在,有人递给你一张风景照,但照片上有一大块被撕掉了(掩码图像)。你的目标是填补这块缺失区域,使其完美无缺,并与画面其余部分自然融合。
目前主要有两种方法,但两者都有个大麻烦:
- “定制学徒”方案:你雇佣一位新的、专门化的学徒,只学习如何修补破损照片。
- 麻烦在于:训练这位学徒耗时且耗费大量资源。此外,如果你以后想修补另一种类型的破损,可能又需要另一位不同的学徒。
- “即时”方案:你请大师画家当场修补照片。但由于这位画家并未专门接受过“修补”训练,他必须停下来,绞尽脑汁,为每一张照片都进行复杂、逐步的计算,以找出如何填补空缺。
- 麻烦在于:速度太慢。如果你有 1,000 张照片,这位画家就必须进行 1,000 次独立的“深度思考”。
解决方案:AID(“智能向导”)
作者提出了一种折中方案,称为 AID(基于扩散模型的摊销图像修复)。
把 AID 想象成在大师画家身旁聘请了一位智能向导。
- 大师画家保持原样。我们不会重新训练他们,也不会改变他们的大脑。他们继续发挥最擅长的本领。
- 智能向导是一个微小、轻量级的助手。我们离线(提前)在成千上万张破损照片的示例上训练这位向导。
- 工作原理:当一张新的破损照片到来时,大师画家开始作画。智能向导则向画家耳语指令,例如:“嘿,看看照片中可见的部分;确保新画的部分与现有树木的风格相匹配。”
一旦向导训练完成,它就能瞬间帮助修复任何破损照片。我们无需为每张新照片停下来思考;向导只需按部就班地工作即可。
核心秘诀:“随机排练”
这篇论文最难的部分在于如何训练向导背后的数学原理。
通常,若要教导向导给出完美指令,你必须为每一张照片解决一个巨大且几乎不可能的数学谜题。作者意识到,直接尝试求解“完美”谜题过于困难。
于是,他们想出了一个巧妙的技巧:
- 类比:想象你想教导向导在一条特定道路上完美驾驶。与其尝试计算每一秒的确切转向角度(这很难),不如让向导在佩戴使方向盘随机晃动的眼罩的情况下练习驾驶这条道路。
- 魔力所在:作者在数学上证明,如果你训练向导应对这种“晃动、随机”的驾驶,那么所有这些晃动动作的平均值,将精确指向你最初想要的那条完美、平滑的路线。
- 结果:他们使用了一种名为连续时间演员 - 评论家(一种强化学习)的方法,利用这种“随机晃动”技术来训练向导。训练完成后,他们只需取“平均值”(即平滑路线),并将其作为向导的指令。这使得训练成为可能,且最终结果完美无缺。
发现(结果)
作者在著名的图像数据集(AFHQv2、FFHQ 和 ImageNet)上测试了该方法。结果如下:
- 速度与质量:AID 比“即时”方法(如 RePaint)快得多,因为它无需为每张新照片进行繁重的计算。事实上,它比现有最佳方法快约10 倍,同时生成的图像质量更高。
- 极小占用:“智能向导”(可训练部分)极其微小。它使原始“大师画家”的体积增加不到1%。这就像给图书馆的书贴一张便利贴,而不是重写整本书。
- 通用性:在同一类破损上训练的同一位向导,无需重新训练即可完美处理不同类型的破损(如中心孔洞、条状孔洞)。
总结
该论文提出了一种利用强大 AI 模型修复受损图像的方法,且不会拖慢其速度。他们通过训练一个微小、可复用的“向导”来实现这一点,该向导向主 AI 耳语指令。他们在数学上证明,可以利用“随机排练”技术训练该向导,从而确保最终指令完美无缺。其结果是一个比先前方法更快、更便宜且修复质量更高的系统。
技术摘要:基于预训练扩散模型的图像修复摊销引导(AID)
问题陈述
图像修复是一项条件生成任务,其目标是在保持与可见像素一致性的同时,重建图像的缺失区域。尽管预训练扩散模型已显著提升了该能力,但现有方法在效率与适应性方面面临两难困境:
- 任务特定训练:如 LaMa、MAT 和 BrushNet 等方法训练专用模型或大型适配模块。虽然这摊销了推理成本,但需要大量的训练时间,并显著增加了模型复杂度。
- 单实例优化:如 RePaint、DPS 和 MCG 等方法保持预训练扩散主干网络固定,但在测试时为每个掩码图像单独进行适配。这保留了预训练模型,但在部署期间产生了重复的优化成本,将每张图像视为一个新的推理问题。
目前缺乏一种“中间路线”的解决方案,既能摊销推理成本,又无需训练大型任务特定模型或在测试时进行单实例优化。
方法论:基于扩散的摊销图像修复(AID)
作者提出了AID,这是一个保持预训练扩散主干网络固定,并离线训练一个小型、可重用引导模块的框架。该模块在测试时可在所有掩码图像上重复使用,无需进一步优化。
1. 控制公式化
该问题被公式化为一个确定性控制问题。
- 动力学:反向扩散过程由受固定预训练得分函数 S^ 驱动的概率流常微分方程(ODE)支配。
- 引导:在漂移项中添加一个任务相关的控制项 u(t,x;ξ),其中 ξ=(M,y) 代表可观测的掩码和可见像素。该控制项不依赖于隐藏的干净图像 x†。
- 目标:目标是最小化终端损失 Ψ(平衡对可见像素和缺失区域的保真度)加上惩罚引导场幅度的运行成本。这导出了一个哈密顿 - 雅可比 - 贝尔曼(HJB)方程,其中最优控制 u∗ 是确定性的。
2. 理论桥梁:从确定性到随机性
直接求解高维确定性 HJB 方程是不可行的。作者引入了一个辅助随机化控制问题,使用高斯策略使问题可通过连续时间强化学习(CT-RL)进行学习。
- 辅助问题:系统不使用确定性控制,而是采用高斯策略 π(λ)∼N(μ,σ2),其中均值 μ 是候选引导律。
- 优化器保持桥梁:论文证明了两个关键定理:
- 充分性:通过将策略均值设置为 u∗,最优确定性引导 u∗ 可以提升为最优高斯策略。
- 必要性:辅助问题的任何最优高斯策略,其均值必须精确恢复最优确定性引导 u∗。
- 意义:这一桥梁确保了学习随机策略(从而实现基于梯度的更新)不会改变部署时使用的目标确定性引导。高斯策略的方差是固定的,仅作为学习工具,而非熵正则化项。
3. 摊销演员 - 评论家算法
基于该理论桥梁,作者开发了一种完全数据驱动的演员 - 评论家算法:
- 参数化:演员(引导均值 μ^ϕ)和评论家(价值函数 V^θ)由神经网络参数化。
- 学习:该算法利用从辅助 HJB 方程导出的鞅正交条件。它基于从任务分布中采样的轨迹,使用随机近似来更新演员和评论家。
- 部署:训练完成后,引导模块以确定性方式部署。对于新的掩码图像,系统在反向时间求解器中使用学习到的均值引导 μ^ϕ,并丢弃策略噪声。
主要贡献
- 方法论:AID 引入了一种受控引导框架,通过离线训练轻量级引导模块来摊销推理成本,该模块随后可在任务间重复使用,无需单实例优化。
- 理论:论文在确定性引导问题与辅助高斯公式化之间建立了一个严格的“策略等价桥梁”。它证明了辅助问题是优化器保持的,从而为使用随机演员 - 评论家方法学习确定性控制场提供了依据。
- 效率:该方法相对于预训练主干网络仅增加了不到 1% 的可训练参数,同时显著改善了质量与速度的权衡。
实证结果
作者在 AFHQv2、FFHQ(使用像素空间 EDM 流程)和 ImageNet(使用潜在空间 EDM2 流程)上评估了 AID。
- 与固定主干基线的性能对比:AID 在多个指标(PSNR、SSIM、LPIPS)和掩码类型(自由形状、中心、条带)上,始终优于强大的固定主干基线(无引导、替换、MCG、DPS)。
- 与 RePaint 的性能对比:AID 在质量上超越了 RePaint(一种强大的迭代优化基线),同时仅使用了 约 10% 的采样预算(例如,RePaint 需要 350 步,而 AID 仅需 35 步)。
- 与 LatentPaint 的性能对比:在 ImageNet 上,AID 在所有指标上均优于 LatentPaint(一种可重用的潜在空间适配方法),同时使用的引导模块可训练参数少于其 45%。
- 延迟:AID 支持低延迟工作点(12 步),在保持性能优于非 RePaint 基线的同时,进一步减少了推理时间,并与 RePaint 保持竞争力。
- 泛化性:同一训练好的引导模块无需重新训练,即可有效迁移至未见过的掩码类型和不同的采样预算。
意义与主张
论文声称 AID 是图像修复领域首个摊销受控引导框架,其特点包括:
- 保持预训练扩散主干网络冻结。
- 仅训练一个小型、可重用的引导模块。
- 提供从随机演员 - 评论家学习到部署所用确定性引导场的精确优化器保持桥梁。
作者将这项工作定位为迈向利用控制和基于强化学习的框架,为研究和改进扩散模型提供理论依据方法的步骤。他们强调,连续时间强化学习既是严谨分析的基础,也是实用算法设计的核心,为启发式单实例优化或繁重的任务特定重新训练提供了一种原则性的替代方案。
每周获取最佳 electrical engineering 论文。
受到斯坦福、剑桥和法国科学院研究人员的信赖。
请查收邮箱确认订阅。
出了点问题,再试一次?
无垃圾邮件,随时退订。