论文状态:已完成

Restora-Flow: Mask-Guided Image Restoration with Flow Matching

发表:2025/11/25
原文链接PDF 下载
价格:0.100000
价格:0.100000
已有 3 人读过
本分析由 AI 生成,可能不完全准确,请以原文为准。

TL;DR 精炼摘要

Restora-Flow是一种新型的图像修复方法,通过引导流匹配采样并结合退化掩码和轨迹校正机制,避免了昂贵的训练过程。该方法在自然及医学数据集上表现出超越现有扩散和流匹配方法的感知质量与处理速度优势,适用于图像修复任务。

摘要

Flow matching has emerged as a promising generative approach that addresses the lengthy sampling times associated with state-of-the-art diffusion models and enables a more flexible trajectory design, while maintaining high-quality image generation. This capability makes it suitable as a generative prior for image restoration tasks. Although current methods leveraging flow models have shown promising results in restoration, some still suffer from long processing times or produce over-smoothed results. To address these challenges, we introduce Restora-Flow, a training-free method that guides flow matching sampling by a degradation mask and incorporates a trajectory correction mechanism to enforce consistency with degraded inputs. We evaluate our approach on both natural and medical datasets across several image restoration tasks involving a mask-based degradation, i.e., inpainting, super-resolution and denoising. We show superior perceptual quality and processing time compared to diffusion and flow matching-based reference methods.

思维导图

论文精读

中文精读

1. 论文基本信息

1.1. 标题

Restora-Flow: Mask-Guided Image Restoration with Flow Matching

1.2. 作者

Arnela Hadzic1, Franz Thaler1,2,3, Lea Bogensperger4, Simon Johannes Joham1, Martin Urschler1

隶属机构: 1 医科大学格拉茨医学信息学、统计学和文献研究所 (Institute for Medical Informatics, Statistics and Documentation, Medical University of Graz, Graz, Austria) 2 医科大学格拉茨医学物理与生物物理学部 (Division of Medical Physics and Biophysics, Medical University of Graz, Graz, Austria) 3 格拉茨科技大学视觉计算 (Visual Computing, Graz University of Technology, Graz, Austria) 4 苏黎世大学定量生物医学系 (Department of Quantitative Biomedicine, University of Zurich, Zurich, Switzerland)

1.3. 发表期刊/会议

该论文作为预印本发表于 arXiv。

1.4. 发表年份

2025年(根据 arXiv 发布日期 2025-11-25)。

1.5. 摘要

Flow Matching (流匹配) 已经成为一种前景广阔的生成方法,它解决了最先进的扩散模型 (diffusion models) 采样时间长的问题,并能够实现更灵活的轨迹设计,同时保持高质量的图像生成。这一能力使其适合作为图像修复 (image restoration) 任务的生成先验 (generative prior)。尽管目前利用流模型 (flow models) 的方法在修复中取得了可喜的成果,但有些方法仍然面临处理时间长或产生过度平滑 (over-smoothed) 结果的问题。为了应对这些挑战,本文引入了 Restora-Flow,这是一种无需训练 (training-free) 的方法,它通过退化掩码 (degradation mask) 引导流匹配采样,并结合了轨迹校正 (trajectory correction) 机制以确保与退化输入的一致性。本文在自然和医学数据集上,针对涉及基于掩码的退化(即图像修复 (inpainting)、超分辨率 (super-resolution) 和去噪 (denoising))的多种图像修复任务对该方法进行了评估。结果显示,与基于扩散和流匹配的参考方法相比,Restora-Flow 在感知质量 (perceptual quality) 和处理时间方面表现出优越性。

1.6. 原文链接

2. 整体概括

2.1. 研究背景与动机

图像修复 (image restoration) 是计算机视觉领域的一个重要任务,旨在从受损或退化的图像中恢复出原始清晰的图像。许多常见的图像修复任务,如图像去噪 (image denoising)、超分辨率 (super-resolution) 和图像修复 (inpainting),都可以被建模为涉及基于掩码的退化 (mask-based degradation) 的逆问题 (inverse problems)。其目标是从一个退化观测 zz 中恢复原始图像 xx,其中 z=Hx+ξz = Hx + \xiHH 是退化操作符,mm 是掩码,ξ\xi 是噪声。修复后的图像需要同时具备高视觉质量和与输入数据的高度一致性。

近年来,扩散模型 (diffusion models) 因其强大的生成能力和生成高质量图像的能力而成为生成建模的标准方法。这促使研究人员探索将其作为无监督先验 (unsupervised priors) 应用于逆问题中,并取得了显著成功。然而,扩散模型存在一个显著的局限性:由于其采样轨迹高度弯曲,导致采样时间较长。

Flow Matching (流匹配) 作为一种新兴的生成建模方法,以其在数据生成过程中保持更直的轨迹而脱颖而出,从而实现更快的训练和采样速度,同时仍能生成高质量的图像。这使得 Flow Matching 成为解决图像修复任务的潜力方向。然而,现有的基于流模型的图像修复方法仍面临一些挑战:部分方法处理时间较长,有些则产生过度平滑 (over-smoothed) 的结果,甚至可能引入伪影 (artifacts)。这些问题限制了它们在实际应用中的性能和效率。

因此,本文的动机在于开发一种能够克服这些局限性的 Flow Matching 图像修复方法,特别是在处理基于掩码的退化任务时,既能保证高质量的重建,又能显著缩短处理时间。

2.2. 核心贡献/主要发现

本文提出了 Restora-Flow,一种无需训练的 (training-free) 基于 Flow Matching 的图像修复方法,其核心贡献和主要发现如下:

  • 引入 Restora-Flow 算法: 提出了一种用于解决基于掩码的逆问题 (inverse problems) 的无需训练算法,它利用无条件流先验模型 (unconditional flow prior models) 和掩码引导融合 (mask-guided fusion) 进行 Flow Matching 采样。这使得模型能够在采样过程中有效整合退化信息。

  • 创新的轨迹校正机制: 为增强修复过程的保真度 (fidelity),引入了一种新颖的校正机制。该机制通过在每次 ODE (Ordinary Differential Equation) 迭代中执行单个校正步骤,将流轨迹引导至与观测数据更好地对齐,显著改善了恢复图像与已知区域的一致性,同时避免了过长的处理时间。

  • 全面的性能评估: 在计算机视觉和医学数据集(CelebA, AFHQ-Cat, COCO, X-ray Hand)上对多种任务(去噪、盒式图像修复、超分辨率和随机图像修复,以及医学图像上的遮挡移除)进行了全面评估。

  • 卓越的性能表现: 实验结果表明,Restora-Flow 在感知质量 (LPIPS) 和处理时间方面优于所有比较的基于扩散和流匹配的基线方法。在失真度量 (SSIM, PSNR) 上也表现出色,通常达到最佳或接近最佳水平。

  • 高效和简洁性: 该方法无需额外的超参数调优,除了标准 Flow Matching 中使用的 ODE 步数外,只需要一个固定的校正步数 C=1C=1,这极大地简化了其应用和部署。

    这些贡献共同使得 Restora-Flow 成为一种高效、高质量且易于使用的图像修复解决方案,特别适用于基于掩码的退化任务。

3. 预备知识与相关工作

3.1. 基础概念

理解本文需要以下基本的技术和理论知识:

  • 图像修复 (Image Restoration): 指从退化观测 zz 中恢复原始清晰图像 xx 的过程。退化通常由退化操作符 HH(如模糊、下采样、掩码)和附加噪声 ξ\xi 引起,表示为 z=Hx+ξz = Hx + \xi。图像修复的目标是同时实现高视觉质量和与观测数据的一致性。

  • 生成模型 (Generative Models): 一类机器学习模型,旨在学习数据分布,并能够生成与训练数据相似的新样本。本文主要关注两类生成模型:扩散模型 (diffusion models) 和流匹配 (flow matching)。

  • 扩散模型 (Diffusion Models, DMs): 一种通过模拟扩散过程来生成数据的生成模型。它包含两个主要过程:

    1. 前向扩散过程 (Forward Diffusion Process): 逐渐向数据中添加高斯噪声,直到数据完全变为纯噪声。
    2. 逆向去噪过程 (Reverse Denoising Process): 学习一个神经网络来反转前向过程,逐步从噪声中恢复原始数据。 扩散模型的采样需要多次迭代去噪步骤,这导致其采样时间通常较长。
  • 流匹配 (Flow Matching, FM): 一种学习连续时间生成模型的方法。它直接学习一个速度场 (velocity field) vθ,tv_{\theta, t},该速度场定义了一个从简单基础分布(如标准高斯分布)到目标数据分布的连续变换轨迹。通过对这个速度场进行常微分方程 (Ordinary Differential Equation, ODE) 积分,可以高效地生成高质量样本。相比扩散模型,Flow Matching 的轨迹通常更直,从而实现更快的采样。

  • 常微分方程 (Ordinary Differential Equation, ODE):Flow Matching 中,ODE 用于描述数据在连续时间 tt 上的演变轨迹。通过积分速度场定义的 ODE,可以从初始噪声样本逐步得到最终的数据样本。

  • 最大后验估计 (Maximum A Posteriori, MAP): 在图像修复中,MAP 估计旨在找到最有可能的原始图像 x^\hat{x},给定退化观测 zz。它通过优化一个目标函数来实现,该目标函数通常包括一个数据保真项 (data fidelity term) D(Hx,z)\mathcal{D}(Hx, z)(衡量重建图像与观测数据的一致性)和一个先验项 (prior term) Rθ(x)\mathcal{R}_\theta(x)(编码图像的先验知识或统计特性): x^=argminxD(Hx,z)+Rθ(x) \hat { x } = \arg \operatorname* { m i n } _ { x } \mathcal { D } ( H x , z ) + \mathcal { R } _ { \theta } ( x ) 其中,HH 是退化操作符,D\mathcal{D} 是数据保真度损失函数,Rθ\mathcal{R}_\theta 是由参数 θ\theta 定义的图像先验。

  • 退化操作符 HH 和掩码 mm (Degradation Operator HH and Mask mm): HH 表示对原始图像 xx 进行退化操作的数学算子,例如下采样、模糊或部分遮挡。掩码 mm 是一个二值矩阵,用于指示图像中已知(未退化)和未知(退化)的区域。在基于掩码的退化任务中,掩码 mm 在恢复过程中扮演着关键角色,它指导模型只关注需要修复的区域,同时保留已知区域。

3.2. 前人工作

3.2.1. 传统图像修复方法

早期和当前的图像修复方法包括学习从退化测量到清晰图像的直接映射,通常通过最小化重建损失 (reconstruction loss) 来实现。

  • [9] Dong et al. (2016) 提出了使用深度卷积网络进行图像超分辨率。
  • [32] Tao et al. (2018) 提出了用于图像去模糊的尺度循环网络。
  • [35] Zamir et al. (2021) 提出了多阶段渐进式图像修复方法。
  • [38] Zhang et al. (2017) 提出了学习深度 CNN 去噪先验用于图像修复。 局限性: 这些方法通常需要包含配对的退化图像和清晰图像的数据集,并且每当任务改变时都需要重新训练模型。

3.2.2. 基于扩散模型的图像修复

随着深度生成先验 (deep generative priors) 的兴起,扩散模型被广泛应用于图像修复任务,因为它们不需要配对数据集,并且能够建模复杂的数据统计。

  • DDRM [18]: 通过对退化操作符 HH 进行奇异值分解 (Singular Value Decomposition),解决线性逆问题。
  • DDNM+ [34]: 利用范围-零空间分解 (range-null space decomposition) 作为指导函数,以零样本 (zero-shot) 方式处理逆问题。
  • RePaint [23]: 通过使用未被遮蔽的区域来引导扩散过程,主要用于图像修复 (inpainting)。
  • ΠIGDM [29]: 引入向量-雅可比积 (vector-Jacobian product) 作为额外指导,确保去噪结果与退化测量的一致性。
  • RED-Diff [24]: 将图像修复表述为一个优化问题,最小化测量一致性损失 (measurement consistency loss),同时应用分数匹配正则化 (score-matching regularization)。 局限性: 扩散模型固有的长采样时间是其显著缺点,尤其是在需要多次迭代的修复任务中。

3.2.3. 基于流匹配的图像修复

近期,Flow Matching 作为一种具有更快采样速度的替代方案,开始被应用于图像修复。

  • OT-ODE [26]: 借鉴 ΠIGDM 的梯度校正项来指导基于流的生成过程,在感知质量方面表现出优势。
  • D-Flow [2]: 将图像修复表述为源点优化问题 (source point optimization problem),通过最小化流匹配框架中初始点的成本函数实现。 局限性: D-Flow 需要通过 ODE 求解器进行反向传播,导致处理时间相对较长(每张样本 5 到 15 分钟)。
  • Flow-Priors [40]: 将流的轨迹分解为多个局部目标,并利用 Tweedie 公式顺序优化这些目标,从而减少了采样时间。
  • PnP-Flow [25]: 将即插即用 (Plug-and-Play, PnP) 方法与 Flow Matching 结合,无需反向传播。 局限性: PnP-Flow 倾向于产生过度平滑的结果。

3.3. 技术演进

图像修复领域的技术演进经历了从早期的基于模型和信号处理方法(如小波变换、稀疏表示),到基于深度学习的端到端映射方法(如卷积神经网络),再到近年来利用深度生成模型作为强大先验的范式转变。扩散模型和 Flow Matching 代表了当前最先进的生成先验方法。

最初,深度学习方法需要大量配对的退化-清晰图像数据进行训练,且泛化能力有限。生成模型的引入,特别是扩散模型,通过学习数据分布本身,使得无需配对数据也能进行修复,大大提升了通用性和质量。然而,扩散模型固有的高计算成本(尤其是采样时间)成为了一个瓶颈。

Flow Matching 正是为了解决扩散模型在采样效率上的痛点而出现的。它通过学习一个更平滑、更直的轨迹,实现了更快的生成速度。本文的 Restora-Flow 进一步推动了这一进展,它将 Flow Matching 的采样效率与掩码引导机制和轨迹校正结合,以 training-free 的方式,在保持甚至提升图像修复质量的同时,显著降低了处理时间,并解决了现有流模型可能存在的过度平滑或伪影问题。这标志着图像修复技术在效率和质量之间平衡方面取得了重要进步。

3.4. 差异化分析

Restora-Flow 与上述相关工作的主要区别和创新点在于:

  • 训练自由 (Training-free): 不同于许多需要针对特定任务或退化模式进行微调的深度学习方法,Restora-Flow 是一种无需训练的方法,它利用预训练的无条件 Flow Matching 模型作为通用先验,通过巧妙的采样引导和校正机制来解决图像修复问题。这大大降低了其应用成本和复杂性。

  • 掩码引导融合 (Mask-guided Fusion): Restora-Flow 采用了一种基于掩码的引导策略,类似于 RePaint [23] 在扩散模型中的应用。它在采样过程中将时间依赖变量与退化图像的已知部分进行融合,确保已知区域的保真度。

  • 轨迹校正机制 (Trajectory Correction Mechanism): 这是 Restora-Flow 的一个核心创新点。它解决了简单掩码融合可能导致样本偏离 Flow Matching 模型训练分布的问题。通过先进行预测性外推 (extrapolation) 到终点,再缩放并重新引入噪声,该机制能够将样本重新拉回到数据流形 (data manifold) 上,纠正掩码融合引起的偏差,从而提高生成样本的质量和一致性。

  • 效率和质量的平衡: 相比于扩散模型(如 RePaintDDNM+DDNM+)的长采样时间,Restora-Flow 显著加快了处理速度。与现有的流匹配方法(如 D-Flow 的长处理时间或 PnP-Flow 的过度平滑结果)相比,Restora-Flow 在保持速度优势的同时,提供了卓越的感知质量,并避免了伪影。

  • 简洁的超参数: Restora-Flow 仅需调整 ODE 步数这一个超参数(与标准 Flow Matching 相同),并在所有实验中固定校正步数 C=1C=1,这使得其易于使用和泛化。

    总结来说,Restora-Flow 在无需训练的前提下,通过创新的掩码引导和轨迹校正机制,有效结合了 Flow Matching 的高效率和高质量生成能力,为基于掩码的图像修复任务提供了一个更优的解决方案。

4. 方法论

4.1. 方法原理

Restora-Flow 的核心思想是利用预训练的无条件 Flow Matching 模型作为生成先验,并通过两个关键机制将其应用于图像修复任务:掩码引导采样 (mask-guided sampling) 和 轨迹校正 (trajectory correction)。

其基本原理在于,图像修复问题可以被公式化为最大后验 (MAP) 估计问题,旨在找到一个在保持与退化观测一致性的同时,也符合图像先验分布的图像。Flow Matching 模型能够学习复杂的图像先验,生成高质量图像。然而,直接从 Flow Matching 模型采样无法保证与特定退化观测的一致性。

为了解决这个问题,Restora-FlowFlow Matching 的连续时间 ODE 采样过程中,周期性地将当前样本与退化图像的已知(未遮蔽)区域进行融合(掩码引导)。这种融合能够强制样本与观测数据保持一致。然而,简单的融合可能使样本偏离 Flow Matching 模型所学习的有效数据流形。为了纠正这种偏离,Restora-Flow 引入了一个轨迹校正机制:在每次 ODE 迭代后,通过将当前样本外推到流轨迹的终点(生成最终图像),然后重新引入噪声并调整到当前时间步,从而将样本“拉回”到符合生成模型分布的轨迹上。这个过程通过单次校正步骤即可实现高质量和高效率的平衡。

4.2. 核心方法详解

4.2.1. 流匹配 (Flow Matching, FM)

在流匹配 [20, 21] 中,其核心思想是学习一个速度场 vθ,tv_{\theta, t},该速度场描述了概率流 Ψt\Psi_t 的变化。这个速度场控制着一个简单的基准分布(在时间 t=0t=0)如何转化为目标分布 p(x)\mathrm{p}(x)(在时间 t=1t=1)的连续轨迹。

为了实现无模拟训练 (simulation-free training),Flow Matching 使用了条件 FM 损失函数: minθEt,x1,x0[12vθ,t(Ψt(x0))(x1x0)2] \operatorname* { m i n } _ { \theta } \mathbb { E } _ { t , x _ { 1 } , x _ { 0 } } \Big [ \frac { 1 } { 2 } \Big \lVert v _ { \theta , t } \big ( \Psi _ { t } ( x _ { 0 } ) \big ) - \big ( x _ { 1 } - x _ { 0 } \big ) \Big \rVert ^ { 2 } \Big ] 符号解释:

  • minθ\operatorname* { m i n } _ { \theta }:表示通过优化模型参数 θ\theta 来最小化期望。
  • Et,x1,x0\mathbb { E } _ { t , x _ { 1 } , x _ { 0 } }:表示对时间 tt、目标样本 x1x_1 和基准样本 x0x_0 求期望。
  • tU[0,1]t \sim \mathcal { U } [ 0 , 1 ]:时间 tt 从均匀分布 U[0,1]\mathcal{U}[0, 1] 中采样。
  • x1p(x)x _ { 1 } \sim \mathrm { p } ( x ):目标样本 x1x_1 从真实数据分布 p(x)\mathrm{p}(x) 中采样。
  • x0N(0,I)x _ { 0 } \sim \mathcal { N } ( 0 , I ):基准样本 x0x_0 从标准正态分布 N(0,I)\mathcal{N}(0, I) 中采样,其中 II 是单位矩阵。
  • Ψt(x0)=(1t)x0+tx1\Psi _ { t } ( x _ { 0 } ) = ( 1 - t ) x _ { 0 } + t x _ { 1 }:这是条件流 (conditional flow),表示在时间 tt 时,从基准样本 x0x_0 到目标样本 x1x_1 的线性插值。
  • vθ,t(Ψt(x0))v _ { \theta , t } \big ( \Psi _ { t } ( x _ { 0 } ) \big ):由神经网络 vθv_\theta 预测的在时间 tt 处、状态为 Ψt(x0)\Psi_t(x_0) 时的速度向量。
  • ( x _ { 1 } - x _ { 0 } ):这是从 x0x_0x1x_1 的真实速度向量(因为 Ψt(x0)\Psi_t(x_0) 是线性插值,其导数即为 x1x0x_1-x_0)。 这个损失函数的目标是训练神经网络 vθv_\theta 来准确预测在给定时间 tt 和状态 Ψt(x0)\Psi_t(x_0) 时,样本从 x0x_0 迁移到 x1x_1 所需的速度。

学习到的速度场 vθ,tv_{\theta, t} 描述了分布随时间的变化。通过积分对应的常微分方程 (ODE),我们可以从基准分布采样得到目标分布的样本: ddtΨt(x)=vθ,t(Ψt(x)) \frac { \mathrm { d } } { \mathrm { d } t } \Psi _ { t } ( x ) = v _ { \theta , t } ( \Psi _ { t } ( x ) ) 符号解释:

  • ddtΨt(x)\frac { \mathrm { d } } { \mathrm { d } t } \Psi _ { t } ( x ):表示样本 Ψt(x)\Psi_t(x) 随时间 tt 的变化率,即其速度。
  • vθ,t(Ψt(x))v _ { \theta , t } ( \Psi _ { t } ( x ) ):神经网络 vθv_\theta 在时间 tt 处、状态为 Ψt(x)\Psi_t(x) 时预测的速度向量。 这个 ODE 描述了样本如何沿着速度场从 t=0t=0 演变到 t=1t=1

例如,使用显式欧拉积分 (explicit Euler integration) 方案,其中 t=0,,1Δtt = 0, \ldots, 1 - \Delta_t,样本的估计值通过以下方式更新: xt+Δt=xt+Δtvθ,t(xt) x _ { t + \Delta _ { t } } = x _ { t } + \Delta _ { t } v _ { \theta , t } ( x _ { t } ) 符号解释:

  • xtx_t:当前时间步 tt 的样本。
  • Δt\Delta_t:时间步长。
  • xt+Δtx_{t+\Delta_t}:下一个时间步 t+Δtt+\Delta_t 的样本估计。
  • vθ,t(xt)v_{\theta, t}(x_t):神经网络 vθv_\theta 在当前时间 tt 和当前样本 xtx_t 处预测的速度向量。 这个公式表示在每个时间步,样本通过沿着速度场方向移动一个小的距离 (Δtvθ,t(xt)\Delta_t v_{\theta, t}(x_t)) 来更新。

4.2.2. Restora-Flow 算法

图像修复 (image restoration) 中,从退化观测 zz 恢复未知图像 xx 可以被公式化为最大后验 (MAP) 估计问题: x^=argminxD(Hx,z)+Rθ(x) \hat { x } = \arg \operatorname* { m i n } _ { x } \mathcal { D } ( H x , z ) + \mathcal { R } _ { \theta } ( x ) 符号解释:

  • x^\hat{x}:估计的原始图像。

  • argminx\arg \operatorname* { m i n } _ { x }:寻找使目标函数最小的 xx

  • HH:退化操作符 (degradation operator)。

  • D(Hx,z)\mathcal { D } ( H x , z ):数据保真项 (data fidelity term),衡量 Hx(重建图像经过退化操作后)与观测 zz 之间的一致性。

  • Rθ(x)\mathcal { R } _ { \theta } ( x ):先验项 (prior term),编码了通过参数 θ\theta 学习到的关于图像的先验知识。

    当使用 Flow Matching 作为先验 Rθ(x)\mathcal { R } _ { \theta } ( x ) 时,仅仅通过 Eq. (3)(即 xt+Δt=xt+Δtvθ,t(xt)x_{t+\Delta_t} = x_t + \Delta_t v_{\theta, t}(x_t))生成样本并不能直接得到 Eq. (4) 的最小化结果。因此,必须通过将退化观测 zz 整合到采样过程中来引导采样朝 MAP 解决方案发展。

本文选择使用掩码引导 (mask-guidance),因为它与图像修复 (inpainting) 和其他涉及基于掩码的退化的图像修复任务具有内在联系。

4.2.2.1. 引入掩码引导 (Incorporating Mask-guidance)

掩码引导由 RePaint [23] 引入,它将时间依赖变量 xtx_t 与原始图像 zz 的未遮蔽部分融合。 首先,通过凸组合 (convex combination) 将 zz 调整以匹配流估计 xtx_t 中包含的噪声水平: z=tz+(1t)ϵ,ϵN(0,I) z ^ { \prime } = t z + ( 1 - t ) \epsilon , \quad \epsilon \sim \mathcal { N } ( 0 , I ) 符号解释:

  • zz':调整后的退化观测,其噪声水平与当前时间步 tt 的生成样本相匹配。
  • tt:当前时间步。
  • zz:原始退化观测。
  • ϵN(0,I)\epsilon \sim \mathcal { N } ( 0 , I ):从标准正态分布中采样的噪声。 这个步骤的目的是确保 zz' 在不同时间步具有合适的噪声水平,以便与 Flow Matching 过程中不同噪声水平的 xtx_t 进行融合。

然后,将 zz'xtx_t 进行融合: xt=mz+(1m)xt x _ { t } ^ { \prime } = m \odot z ^ { \prime } + ( 1 - m ) \odot x _ { t } 符号解释:

  • xtx_t':融合后的样本,其中已知区域来自 zz',未知区域来自 xtx_t
  • mm:二值掩码 (binary mask),指示 zz 中已知(1)和未知(0)的区域。
  • \odot:表示元素级乘法 (element-wise multiplication)。
  • zz':上述调整后的退化观测。
  • xtx_t:当前时间步 tt 的生成样本。 这个公式表示在掩码 mm 为1的区域(已知区域),使用 zz' 的信息;在掩码 mm 为0的区域(未知区域),保留 xtx_t 的信息。

一种朴素 (naive) 的方法是在 Flow Matching 更新 (Eq. (3)) 之前,将每个时间依赖的样本 xtx_t 与掩码进行融合。这将产生一个修改后的更新方程: xt+Δt=xt+Δtvθ,t(xt) x _ { t + \Delta _ { t } } = x _ { t } ^ { \prime } + \Delta _ { t } v _ { \theta , t } ( x _ { t } ^ { \prime } ) 符号解释:

  • xt+Δtx_{t+\Delta_t}:下一个时间步的样本估计。
  • xtx_t':根据 Eq. (6) 融合后的样本。
  • Δt\Delta_t:时间步长。
  • vθ,t(xt)v_{\theta, t}(x_t'):神经网络 vθv_\theta 在时间 tt 和融合样本 xtx_t' 处预测的速度向量。 然而,这种朴素方法会导致更新后的样本 xt+Δtx_{t+\Delta_t} 在已知和未知区域之间出现明显的错位,尤其是在掩码边界处。这是因为 xtx_t' 很可能不在 Flow Matching 模型训练分布所学习的轨迹上,导致模型无法可靠地将 xtx_t' 映射到高质量的图像。

4.2.2.2. 轨迹校正 (Trajectory Correction)

为了解决上述问题并提高生成样本的质量,Restora-Flow 引入了一种轨迹校正机制。在初始的掩码引导更新 (Eq. (7)) 之后,执行以下步骤:

  1. 前向外推 (Forward Extrapolation): 使用学习到的速度场将样本 xt+Δtx_{t+\Delta_t} 外推到生成路径的终点 t=1t=1x~1=xt+Δt+(1(t+Δt))vθ,t+Δt(xt+Δt) \widetilde { x } _ { 1 } = x _ { t + \Delta _ { t } } + \left( 1 - \left( t + \Delta _ { t } \right) \right) v _ { \theta , t + \Delta _ { t } } ( x _ { t + \Delta _ { t } } ) 符号解释:

    • x~1\widetilde{x}_1:外推到时间 t=1t=1 的样本估计。
    • xt+Δtx_{t+\Delta_t}:经过朴素更新后的样本。
    • (1(t+Δt))(1 - (t + \Delta_t)):从当前时间 t+Δtt+\Delta_t 到终点 t=1t=1 的剩余时间步长。
    • vθ,t+Δt(xt+Δt)v_{\theta, t+\Delta_t}(x_{t+\Delta_t}):神经网络 vθv_\theta 在时间 t+Δtt+\Delta_t 和样本 xt+Δtx_{t+\Delta_t} 处预测的速度向量。 这个外推步骤充当一个学习到的去噪器,帮助将样本投影到更接近数据流形的位置,从而校正掩码融合引起的错位。
  2. 重新引入噪声 (Reintroduction of Noise): 为了将样本放置在生成轨迹上的正确位置,对 x~1\widetilde{x}_1 进行缩放并重新引入噪声: xt=tx~1+(1t)η,ηN(0,I). \begin{array} { r } { x _ { t } = t \widetilde { x } _ { 1 } + ( 1 - t ) \eta , \quad \eta \sim \mathcal { N } ( 0 , I ) . } \end{array} 符号解释:

    • xtx_t:经过校正后,用于下一个 ODE 迭代的样本。
    • tt:当前时间步。
    • x~1\widetilde{x}_1:外推到时间 t=1t=1 的样本。
    • (1t)(1 - t):当前时间步 tt 到基准时间 t=0t=0 的距离。
    • ηN(0,I)\eta \sim \mathcal { N } ( 0 , I ):从标准正态分布中采样的噪声。 这个步骤重新引入了随机性到生成过程中,允许流模型生成多样化且真实的图像。经验表明,即使每次 ODE 步长使用单个校正 (C=1) 也能显著改善修复图像与观测内容之间的一致性。

4.2.2.3. 算法总结

Restora-Flow 算法针对基于掩码的图像修复任务(如图像修复和超分辨率)总结在算法 1 中。对于图像去噪任务,它在算法 2 中有描述。

Algorithm 1 Mask-Guided Restora-Flow Sampling 输入: 已学习的流网络 vθv_{\theta}, 退化观测 zRdz \in \mathbb{R}^d, ODE 步数 NN (其中 Δt1N\Delta_t \gets \frac{1}{N}), 校正步数 C>0C > 0, 掩码 mm

  1. N(0,I)\mathcal{N}(0, I) 采样 xx (初始化样本)。
  2. 对于 t=0,Δt,,1Δtt = 0, \Delta_t, \ldots, 1 - \Delta_t 执行:
  3. 对于 c=0,,Cc = 0, \ldots, C 执行:
  4. N(0,I)\mathcal{N}(0, I) 采样 ϵ\epsilon
  5. 如果 t>0t > 0,则 ztz+(1t)ϵz' \gets t z + (1 - t) \epsilon,否则 z=0z' = 0
  6. xmz+(1m)xx' \gets m \odot z' + (1 - m) \odot x (掩码引导融合)。
    
  7. xx+Δtvθ,t(x,t)x \gets x' + \Delta_t v_{\theta, t}(x', t) (朴素 ODE 更新)。
    
  8. 如果 c>0c > 0t<1Δtt < 1 - \Delta_t 执行:
    
  9. N(0,I)\mathcal{N}(0, I) 采样 η\eta
  10.  x~1x+(1(t+Δt))vθ,t+Δt(x,t+Δt)\widetilde{x}_1 \gets x + (1 - (t + \Delta_t)) v_{\theta, t + \Delta_t}(x, t + \Delta_t) (前向外推)。
    
  11.  xtx~1+(1t)ηx \gets t \widetilde{x}_1 + (1 - t) \eta (重新引入噪声并调整)。
    
  12. 否则 (当 c=0c=0t=1Δtt=1-\Delta_t 时,不执行校正):
  13.  tt+Δtt \gets t + \Delta_t (更新时间)。
    
  14. 返回 xx

Algorithm 2 Restora-Flow Sampling for Denoising 输入: 退化观测 zRdz \in \mathbb{R}^d 带有噪声水平 σ\sigma, ODE 步数 NN (其中 Δt1N\Delta_t \gets \frac{1}{N})

  1. N(0,I)\mathcal{N}(0, I) 采样 x0x_0 (初始化样本)。
  2. 对于 t=0,Δt,,1Δtt = 0, \Delta_t, \ldots, 1 - \Delta_t 执行:
  3. N(0,I)\mathcal{N}(0, I) 采样 ϵ\epsilon
  4. z(1σ)zz' \gets (1 - \sigma) z。 (注意: 此处的 zz' 公式与 Algorithm 1 不同)
  5. xt+Δtxt+Δtvθ,t(xt)x_{t + \Delta_t}' \gets x_t + \Delta_t v_{\theta, t}(x_t) (ODE 更新)。
  6. xt+Δt1{t<1σ}z+1{t1σ}xt+Δtx_{t + \Delta_t} \gets \mathbf{1}_{\{t < 1 - \sigma\}} z' + \mathbf{1}_{\{t \geq 1 - \sigma\}} x_{t + \Delta_t}' (时间依赖的掩码引导)。
  7. 返回 x1x_{1} (最终的去噪图像)。

Algorithm 2 的特点: 与算法 1 不同,算法 2 使用了一个时间依赖的全局掩码 m(t)。这里的掩码是隐式的,通过指示函数 1{.}\mathbf{1}_{\{.\}} 来实现。当 t<1σt < 1 - \sigma 时,指示函数为 1,这意味着噪声观测 zz(通过 zz' 形式)作为采样的初始化和引导。当 t1σt \geq 1 - \sigma 时,指示函数为 0,此时 ODE 演化解,不再受 zz 的进一步影响。这意味着在早期阶段(噪声水平较高时),模型更多地依赖观测 zz;而在后期阶段(噪声水平降低后),模型更多地依赖于学习到的流先验来生成清晰图像。

5. 实验设置

5.1. 数据集

为了评估 Restora-Flow 的性能,本文使用了四个不同特征的数据集,涵盖了自然图像和医学图像:

  • CelebA [22]: 包含 162k 张名人面部训练图像,大小调整为 128×128128 \times 128 像素。
  • AFHQ-Cat [5]: 包含 5k 张猫脸训练图像,大小调整为 256×256256 \times 256 像素。
  • COCO [19]: 包含 118k 张各种物体类型的训练图像,大小调整为 128×128128 \times 128 像素。
  • X-ray Hand [13, 36]: 包含 895 张手部 X 射线图像,大小调整为 256×256256 \times 256 像素。

测试集:

  • CelebA: 100 张测试图像。

  • AFHQ-Cat: 100 张测试图像。

  • COCO: 100 张验证图像。

  • X-ray Hand: 298 张测试图像。

    以下是 Restora-Flow 在 CelebA, AFHQ-Cat, COCO 和 X-ray Hand 数据集上的定性结果示例图,展示了去噪、盒式图像修复、超分辨率和随机图像修复的效果:

    该图像是一个比较不同图像修复方法的插图,展示了去噪、框内修补、超分辨率以及随机修补等任务下的结果。各行展示了不同方法的恢复效果,包括原始图像和使用 Restora-Flow 制作的结果。 该图像是一个比较不同图像修复方法的插图,展示了去噪、框内修补、超分辨率以及随机修补等任务下的结果。各行展示了不同方法的恢复效果,包括原始图像和使用 Restora-Flow 制作的结果。

图注: Restora-Flow 在CelebA数据集上的定性结果。从上到下:原始图像、去噪 (row 1)、盒式图像修复 (row 2)、超分辨率 (row 3) 和随机图像修复 (row 4)。差异在PDF版本中可看得更清楚。

该图像是一个比较不同图像恢复方法效果的示意图。展示了去噪、盒子填充、超分辨率和随机填充的结果,并分别与原始图像及其他方法(如RePaint、DDNM、Flow-Priors等)进行了对比,最后展示了Restora-Flow的效果。 该图像是一个比较不同图像恢复方法效果的示意图。展示了去噪、盒子填充、超分辨率和随机填充的结果,并分别与原始图像及其他方法(如RePaint、DDNM、Flow-Priors等)进行了对比,最后展示了Restora-Flow的效果。

图注: Restora-Flow 在AFHQ-Cat数据集上的定性结果。从上到下:原始图像、去噪 (row 1)、盒式图像修复 (row 2)、超分辨率 (row 3) 和随机图像修复 (row 4)。差异在PDF版本中可看得更清楚。

5.2. 评估指标

本文采用以下评估指标来衡量图像修复任务的性能:

  • 结构相似性指数 (Structural Similarity Index, SSIM):

    1. 概念定义: SSIM 是一种衡量两幅图像相似度的指标,它从亮度、对比度和结构三个方面进行评估。与传统的 PSNR 关注像素级误差不同,SSIM 更符合人类视觉感知系统对图像质量的判断。其取值范围通常在 -1 到 1 之间,1 表示两幅图像完全相同,0 表示完全不相似。
    2. 数学公式: SSIM(x,y)=[l(x,y)]α[c(x,y)]β[s(x,y)]γ \mathrm{SSIM}(x, y) = [l(x, y)]^{\alpha} \cdot [c(x, y)]^{\beta} \cdot [s(x, y)]^{\gamma} 其中,通常 α=β=γ=1\alpha = \beta = \gamma = 1,且: l(x,y)=2μxμy+C1μx2+μy2+C1 l(x, y) = \frac{2\mu_x\mu_y + C_1}{\mu_x^2 + \mu_y^2 + C_1} c(x,y)=2σxσy+C2σx2+σy2+C2 c(x, y) = \frac{2\sigma_x\sigma_y + C_2}{\sigma_x^2 + \sigma_y^2 + C_2} s(x,y)=σxy+C3σxσy+C3 s(x, y) = \frac{\sigma_{xy} + C_3}{\sigma_x\sigma_y + C_3}
    3. 符号解释:
      • xx: 原始(参考)图像。
      • yy: 修复后的图像。
      • μx\mu_x: 图像 xx 的像素平均值。
      • μy\mu_y: 图像 yy 的像素平均值。
      • σx\sigma_x: 图像 xx 的像素标准差。
      • σy\sigma_y: 图像 yy 的像素标准差。
      • σxy\sigma_{xy}: 图像 xxyy 的协方差。
      • C1,C2,C3C_1, C_2, C_3: 避免分母为零或数值不稳定的常数,通常为小正数。
      • l(x, y): 亮度比较函数。
      • c(x, y): 对比度比较函数。
      • s(x, y): 结构比较函数。
      • α,β,γ\alpha, \beta, \gamma: 用于调整三个分量重要性的参数,通常设为1。 指标趋势: SSIM 值越高越好。
  • 峰值信噪比 (Peak Signal-to-Noise Ratio, PSNR):

    1. 概念定义: PSNR 是一个广泛用于衡量图像质量的客观指标,通常以分贝 (dB) 为单位。它通过比较原始图像和修复图像的像素值差异来量化重建质量。PSNR 值越高表示图像失真越小,质量越好。它对像素级误差敏感。
    2. 数学公式: PSNR=10log10(MAXI2MSE) \mathrm{PSNR} = 10 \cdot \log_{10} \left( \frac{\mathrm{MAX}_I^2}{\mathrm{MSE}} \right) 其中,MSE (Mean Squared Error) 为: MSE=1MNi=1Mj=1N[I(i,j)K(i,j)]2 \mathrm{MSE} = \frac{1}{MN} \sum_{i=1}^{M} \sum_{j=1}^{N} [I(i,j) - K(i,j)]^2
    3. 符号解释:
      • MAXI\mathrm{MAX}_I: 图像中像素的最大可能值。对于 8 位灰度图像,通常为 255。
      • MSE\mathrm{MSE}: 原始图像和修复图像之间的均方误差。
      • I(i,j): 原始图像在坐标 (i,j) 处的像素值。
      • K(i,j): 修复图像在坐标 (i,j) 处的像素值。
      • M, N: 图像的行数和列数。 指标趋势: PSNR 值越高越好。
  • 学习感知图像块相似度 (Learned Perceptual Image Patch Similarity, LPIPS) [39]:

    1. 概念定义: LPIPS 是一种基于深度学习的图像相似度度量,旨在更好地匹配人类的感知判断。它通过计算两幅图像在预训练深度神经网络(如 AlexNet、VGG 等)提取的特征空间中的距离来评估它们之间的相似性。LPIPS 值越低表示两幅图像在感知上越相似,质量越高。它在评估生成模型输出的感知真实感方面比 PSNRSSIM 更有效。
    2. 数学公式: LPIPS(x,y)=l1HlWlh,wwl(ϕl(x)h,wϕl(y)h,w)22 \mathrm{LPIPS}(x, y) = \sum_{l} \frac{1}{H_l W_l} \sum_{h, w} \| w_l \odot (\phi_l(x)_{h,w} - \phi_l(y)_{h,w}) \|_2^2
    3. 符号解释:
      • xx: 原始(参考)图像。
      • yy: 修复后的图像。
      • ϕl()\phi_l(\cdot): 预训练神经网络在第 ll 层提取的特征图。
      • wlw_l: 在第 ll 层上学习到的权重,用于对不同通道进行加权。
      • Hl,WlH_l, W_l: 第 ll 层特征图的高度和宽度。
      • \odot: 元素级乘法。
      • 22\|\cdot\|_2^2: L2 范数的平方。 指标趋势: LPIPS 值越低越好。
  • 处理时间 (Processing Time): 以秒 (s) 为单位,表示恢复单张图像所需的平均时间。 指标趋势: 处理时间越低越好。

5.3. 对比基线

本文将 Restora-Flow 与以下流匹配 (flow-based) 和扩散模型 (diffusion-based) 方法进行比较:

流匹配基线:

  • OT-ODE [26]: 利用梯度校正来指导流生成过程。
  • Flow-Priors [40]: 通过迭代腐败轨迹匹配 (Iterative Corrupted Trajectory Matching) 来优化局部目标。
  • D-Flow [2]: 将图像修复视为源点优化问题。
  • PnP-Flow [25]: 结合了即插即用 (Plug-and-Play) 方法与 Flow Matching超参数设置: 在 CelebA 和 AFHQ-Cat 数据集上,沿用了 [25] 中为每个基线通过网格搜索找到的最优超参数。OT-ODE、Flow-Priors 和 PnP-Flow 各需调优两个超参数,而 D-Flow 需调优三个。这些值会根据任务和数据集而异。COCO 数据集使用与 CelebA 相同的超参数,X-ray Hand 数据集使用与 AFHQ-Cat 相同的超参数。

扩散模型基线:

  • RePaint [23]: 利用去噪扩散概率模型进行图像修复。 超参数设置: 使用其提出的超参数值:跳跃长度 (jump length) 为 10,重采样步数 (resampling steps) 为 10。
  • DDNM+ [34]: 使用去噪扩散零空间模型进行零样本图像修复。 超参数设置: 通过时间旅行技巧 (time-travel trick) [23, 34] 达到最佳性能,参数设置为 s=1s=1l=5l=5

5.4. 实验任务

本文评估了 Restora-Flow 在以下图像修复任务上的性能:

  • 去噪 (Denoising):
    • CelebA, COCO, AFHQ-Cat: 噪声水平 σ=0.2\sigma = 0.2
    • X-ray Hand: 噪声水平 σ=0.08\sigma = 0.08
  • 盒式图像修复 (Box Inpainting):
    • CelebA, COCO: 40×4040 \times 40 像素的中心掩码。
    • AFHQ-Cat: 80×8080 \times 80 像素的中心掩码。
    • X-ray Hand: 128×128128 \times 128 像素的中心掩码。
  • 超分辨率 (Super-resolution):
    • CelebA, COCO, X-ray Hand: 2×2 \times 超分辨率。
    • AFHQ-Cat: 4×4 \times 超分辨率。
  • 随机图像修复 (Random Inpainting):
    • CelebA, AFHQ-Cat, COCO: 掩码覆盖 70% 的像素。
  • 遮挡移除 (Occlusion Removal):
    • X-ray Hand: 临床动机的遮挡移除任务,合成添加遮挡,如 [14] 所述。

      通用设置: 除非另有说明,所有实验均包含高斯测量噪声水平 σ=0.01\sigma = 0.01

5.5. 实现细节

  • 预训练模型:
    • CelebA 和 AFHQ-Cat 数据集使用 [25] 提供的预训练流模型。
    • 为了与扩散模型进行公平比较,为 RePaint [23] 和 DDNM+DDNM+ [34] 从头开始训练 DDPMs,使用与预训练流模型相同的 U-Net 架构和训练参数。
      • CelebA: 学习率 1e-4, 批大小 128, 训练 200 epoch。
      • AFHQ-Cat: 学习率 1e-4, 批大小 64, 训练 400 epoch。
      • 扩散时间步数 (diffusion time steps) 设为 250。
    • COCO: 从头开始训练 DDPM 和流模型,学习率 1e-4, 批大小 64, 训练 300 epoch。
    • X-ray Hand: 使用 [14] 提供的预训练流模型和扩散模型。
  • 基线实现: 流模型的基线 (OT-ODE, Flow-Priors, D-Flow, PnP-Flow) 使用 [25] 提供的实现。RePaintDDNM+DDNM+Restora-Flow 也在同一框架内实现。
  • 医学数据集: 对于 X-ray Hand 数据集,所有基线和 Restora-Flow 都在 MEDDDPM 框架 [10, 14] 内实现。
  • Restora-Flow 超参数:
    • 校正步数 C=1C=1 (所有实验和数据集)。
    • ODE 步数:
      • 自然图像数据集 (CelebA, AFHQ-Cat, COCO):
        • 去噪和盒式图像修复: 64 ODE 步。
        • 2×2 \times 超分辨率和随机图像修复: 128 ODE 步。
        • 4×4 \times 超分辨率: 256 ODE 步。
      • 医学 X-ray Hand 数据集:
        • 去噪和 2×2 \times 超分辨率: 64 ODE 步。
        • 盒式图像修复和遮挡移除: 32 ODE 步。
  • 硬件:
    • 生成先验模型在 NVIDIA A100 GPU 上训练。
    • 256×256256 \times 256 分辨率的修复实验在 NVIDIA A100 GPU 上进行。
    • 128×128128 \times 128 分辨率的修复实验在 NVIDIA GeForce RTX 3090 上进行。

6. 实验结果与分析

6.1. 核心结果分析

本文通过定量和定性分析,全面评估了 Restora-Flow 在多个数据集和任务上的性能,并与多种基于扩散和流匹配的基线方法进行了比较。

6.1.1. 定量结果分析 (Quantitative Results Analysis)

以下是原文 Table 1 的结果,展示了 Restora-Flow 与其他方法在 CelebA, AFHQ-Cat, COCO 和 X-ray Hand 数据集上的 LPIPS (\downarrow 表示越低越好), SSIM (\uparrow 表示越高越好), PSNR (\uparrow 表示越高越好) 和处理时间 (\downarrow 表示越低越好) 表现。

以下是原文 Table 1 的结果:

Model Denoising σ = 0.2 Box inpainting 40 × 40 Super-resolution 2× Random inpainting 70%
LPIPS SSIM PSNR Time in s LPIPS SSIM PSNR Time in s LPIPS SSIM PSNR Time in s LPIPS SSIM PSNR Time in s
CelebA
RePaint [23] N/A N/A N/A N/A 0.016 0.967 30.81 32.89 0.014 0.946 32.59 32.89 0.014 0.945 32.37 32.89
DDNM+ [34] 0.076 0.885 30.70 11.57 0.019 0.969 31.05 11.57 0.046 0.905 30.02 11.57 0.031 0.920 30.83 11.57
OT-ODEDE [26] 0.033 0.858 30.36 2.95 0.022 0.954 29.85 3.68 0.055 0.870 28.65 3.76 0.051 0.871 28.41 3.76
Flow-Priors [40] 0.132 0.767 29.27 26.22 0.020 0.969 31.17 26.22 0.110 0.722 28.52 26.22 0.019 0.944 32.34 26.22
D-Flow [2] 0.099 0.695 24.64 22.73 0.041 0.907 29.77 65.81 0.031 0.894 31.30 71.43 0.021 0.931 32.48 131.78
PnP-Flow [25] 0.056 0.910 32.12 4.60 0.045 0.941 30.48 4.60 0.058 0.908 31.37 4.60 0.022 0.954 33.55 4.60
Restora-Flow 0.019 0.922 33.09 0.58 0.018 0.964 30.91 2.06 0.014 0.952 33.59 3.63 0.015 0.947 32.71 3.63
AFHQ-Cat
RePaint [23] N/A N/A N/A N/A 0.043 0.939 26.26 86.23 0.139 0.701 24.71 86.23 0.034 0.897 30.93 86.23
DDNM+ [34] 0.170 0.818 29.06 13.74 0.048 0.942 25.16 13.74 0.462 0.534 19.69 13.74 0.065 0.876 30.12 13.74
OT-ODE [26] 0.078 0.814 29.73 5.54 0.048 0.924 24.36 6.94 0.285 0.565 21.85 7.28 0.094 0.839 28.87 7.28
Flow-Priors [40] 0.153 0.771 29.43 67.10 0.054 0.942 26.04 67.05 0.271 0.565 23.50 67.30 0.046 0.909 31.82 67.69
D-Flow [2] 0.184 0.648 24.98 44.45 0.112 0.839 26.17 126.09 0.123 0.707 25.34 261.84 0.056 0.878 30.97 266.18
PnP-Flow [25] 0.165 0.864 31.10 9.86 0.124 0.904 26.18 9.86 0.180 0.790 26.95 46.26 0.042 0.930 33.07 19.15
Restora-Flow 0.051 0.899 32.35 0.72 0.047 0.939 25.96 3.96 0.158 0.761 26.33 14.48 0.034 0.914 31.99 7.48
COCO
RePaint [23] N/A N/A N/A N/A 0.093 0.922 21.20 32.89 0.046 0.856 25.84 32.89 0.038 0.876 26.82 32.89
DDNM+ [34] 0.162 0.805 27.04 11.57 0.112 0.925 21.71 11.57 0.257 0.682 19.05 11.57 0.069 0.845 25.80 11.57
OT-ODEDE [26] 0.066 0.810 27.52 2.95 0.073 0.914 23.40 3.68 0.146 0.745 23.83 3.76 0.130 0.763 23.98 3.76
Flow-Priors [40] 0.116 0.751 27.08 26.22 0.084 0.927 23.58 26.22 0.112 0.698 24.93 26.22 0.055 0.855 25.97 26.22
D-Flow [2] 0.252 0.552 21.19 22.73 0.115 0.825 23.46 65.81 0.083 0.778 24.80 71.43 0.053 0.840 26.29 131.78
PnP-Flow [25] 0.128 0.855 28.97 4.60 0.121 0.892 24.56 4.60 0.118 0.827 26.73 4.60 0.053 0.896 28.13 4.60
Restora-Flow 0.026 0.905 30.57 0.58 0.084 0.929 24.80 2.06 0.044 0.877 27.44 3.63 0.040 0.881 27.37 3.63
X-ray Hand
RePaint [23] N/A N/A N/A N/A 0.046 0.821 23.90 17.02 0.074 0.767 20.04 17.02 0.032 0.898 29.66 17.02
DDNM+ [34] 0.057 0.819 23.78 13.35 0.059 0.801 22.76 13.35 0.143 0.635 14.10 13.35 0.047 0.884 26.57 13.35
OT-ODEE [26] 0.026 0.853 27.83 8.73 0.038 0.801 23.58 11.17 0.076 0.684 22.01 11.17 0.029 0.845 26.55 11.17
Flow-Priors [40] 0.033 0.885 28.58 68.54 0.035 0.882 25.74 68.59 0.162 0.460 20.38 68.62 0.023 0.933 27.07 68.59
D-Flow w [2] 0.077 0.630 24.09 101.66 0.145 0.588 13.61 285.55 0.127 0.639 15.26 361.22 0.110 0.587 22.23 361.22
PnP-Flow [25] 0.052 0.843 25.17 20.48 0.054 0.822 23.67 20.35 0.029 0.884 25.88 102.29 0.045 0.889 26.83 20.35
Restora-Flow 0.021 0.912 31.34 0.50 0.035 0.846 24.67 4.03 0.037 0.857 24.66 7.95 0.017 0.935 33.51 4.03

关键发现:

  • CelebA 数据集:

    • Restora-Flow 在感知质量 (LPIPS) 上表现最佳,并在去噪、超分辨率任务中取得最高的 SSIM 和 PSNR。在盒式图像修复和随机图像修复任务中,LPIPS 也是最佳,SSIM 和 PSNR 接近最佳(分别为 Flow-PriorsPnP-Flow 的第二名)。
    • 在处理时间方面,Restora-Flow 显著优于所有其他方法,通常比扩散模型快约 6-10 倍,也比大多数流匹配基线快数倍。例如,去噪任务仅需 0.58秒,而 RePaint32.89秒
  • AFHQ-Cat 和 COCO 数据集:

    • Restora-Flow 在这些自然图像数据集上也展现了其多功能性。在几乎所有任务的 LPIPS 指标上都表现最佳,即使不是最佳,也仅次于表现最好的方法。
    • 在 SSIM 和 PSNR 指标上,Restora-Flow 也取得了领先或接近领先的性能。
    • 处理时间方面,Restora-Flow 在大多数设置下都是最快的。例如,在 AFHQ-Cat 去噪任务中,Restora-Flow 仅需 0.72秒,而 OT-ODE5.54秒Flow-Priors67.10秒
  • X-ray Hand 数据集 (医学图像):

    • 在医学数据集上,Restora-Flow 同样表现出色,在 LPIPS、SSIM 和 PSNR 指标上 consistently 取得优异结果,并且处理时间最短。
    • 在具有临床意义的遮挡移除任务中,Restora-Flow 也取得了最佳的感知质量和最短的处理时间。
  • 与扩散模型基线对比:

    • RePaintDDNM+DDNM+ 在某些任务上能提供有竞争力的重建质量,但其处理时间显著高于 Restora-Flow。例如,Restora-FlowRePaint 快约 10 倍,比 DDNM+DDNM+ 快约 6 倍,且能提供同等甚至更好的重建质量。
  • 与流匹配基线对比:

    • D-Flow 尽管有时能产生真实的视觉结果,但其处理时间非常长,在一些任务中甚至超过 100 秒。
    • PnP-Flow 尽管在 SSIM 和 PSNR 上表现不错,但其结果常常显得过度平滑,感知质量 (LPIPS) 不如 Restora-Flow
    • OT-ODE 在某些情况下速度较快,但重建质量通常不如 Restora-Flow

6.1.2. 效率与质量权衡的可视化 (Visualization of Efficiency and Quality Trade-off)

以下是 CelebA 数据集上定量结果的可视化图,展示了 LPIPS (\downarrow), SSIM (\uparrow), PSNR (\uparrow) 与处理时间 (\downarrow) 的关系。

Figure 4. Visual representation of quantitative results on CelebA. Restora-Flow \(( \\bigcirc )\) is compared to related work methods (other shapes) on four different tasks (colors). The plots show LPIPS \(\\downarrow\) (left), \({ \\bf S S I M \\uparrow }\) (center) and PSNR \(\\uparrow\) (right) on the y-axis, and processing time \(\\downarrow\) (all plots) on the \(\\mathbf { X }\) ai.For better visualizatin and comparison each plot is separateinto two parts with different scales in the \(\\mathbf { X }\) -axis. 该图像是图表,展示了Restora-Flow与相关方法在CelebA数据集上的定量比较结果。图表中包含三个子图,分别展示了LPIPS值(左)、SSIM值(中)和PSNR值(右)与处理时间的关系。每个子图的y轴显示评分,x轴表示处理时间,使用不同形状和颜色标识不同方法及任务类型,包括去噪、超分辨率和图像修复。

图注: CelebA 数据集定量结果的可视化表示。Restora-Flow (\bigcirc) 与相关工作方法(其他形状)在四个不同任务(颜色)上进行比较。图表显示 y 轴上的 LPIPS (\downarrow)(左)、SSIM (\uparrow)(中)和 PSNR (\uparrow)(右),以及 x 轴上的处理时间 (\downarrow)(所有图表)。为了更好的可视化和比较,每个图表分为两部分,x 轴具有不同的刻度。

从图 4 可以清晰地看出,Restora-Flow 在 LPIPS、SSIM 和 PSNR 指标上都处于左上角,这意味着它在实现最佳重建质量(低 LPIPS,高 SSIM/PSNR)的同时,也保持了最短的处理时间。这直观地证明了 Restora-Flow 在质量和效率之间的优越平衡。

6.1.3. 定性结果分析 (Qualitative Results Analysis)

定性结果(如 Figures 3, 6, 7, 9, 10, 11)进一步支持了定量结果的发现:

  • OT-ODE: 倾向于产生伪影 (artifacts),尤其是在 CelebA 和 AFHQ-Cat 的盒式图像修复以及 COCO 的随机图像修复任务中。
  • Flow-Priors: 在 CelebA 的超分辨率输出中产生嘈杂的重建结果,并在 AFHQ-Cat 和 COCO 的超分辨率和随机图像修复任务中产生伪影。
  • D-Flow: 尽管通常产生逼真的结果,但处理时间很慢,并且在 COCO 数据集的去噪任务中难以重建某些对象。
  • PnP-Flow: 尽管 SSIM 和 PSNR 分数较高,但在所有数据集上经常产生过度平滑 (over-smoothed) 的结果,缺乏细节和纹理。
  • 扩散模型基线:
    • DDNM+DDNM+ 倾向于在超分辨率任务中引入伪影。
    • RePaint 产生视觉上逼真的结果,但处理速度非常慢。
  • Restora-Flow: 能够生成无伪影、逼真的图像,并保持纹理,同时在所有实验中都确保了快速处理。

6.1.4. 医学数据集的特殊观察

在医学 X-ray Hand 数据集上,不同方法之间的总体变异性较低,这可能是由于手部 X 射线图像与自然图像相比变化较小。尽管如此,Restora-Flow 仍然 consistent 地提供了高质量的重建。特别是在临床上动机的遮挡移除任务中,一些基线方法无法完全移除遮挡,留下部分可见的伪影,而 Restora-Flow 则能高质量地完成移除,这对于下游的分割或分类任务具有重要意义。

6.2. 消融实验/参数分析

本文通过消融实验探究了 ODE 步数和校正步数 CCRestora-Flow 性能的影响。

以下是 ODE 步数和校正步数 CC 对 CelebA 数据集上 2×2 \times 超分辨率性能影响的图表:

Figure 5. Ablation of ODE steps (indicated by markers) and correction steps \(C\) for \(2 \\times\) super-resolution on CelebA comparing LPIPS \(\\downarrow\) (top), \({ \\bf S S I M \\uparrow }\) (middle) and \(\\mathrm { P S N R \\uparrow }\) (bottom) to processing time \(\\downarrow\) .ODE steps increase from left to right and represent 4, 8, 16, 32, 64, 128 and 256, respectively. For better visualization, ODE steps 4 and 8 when using \(C = 0\) are omitted. The circle indicates the selected hyperparameters. Time is per image and displayed on a logarithmic scale. 该图像是图表,展示了在 CelebA 数据集上进行 2 imes 超分辨率的实验结果。图中分别显示了使用不同修正步骤 CC 值下的 LPIPS、SSIM 和 PSNR 指标相对于处理时间的变化,横轴为时间,纵轴为各指标的值。OD步骤的变化在图中标记。

图注: CelebA 上 2×2 \times 超分辨率的 ODE 步数(由标记表示)和校正步数 CC 的消融实验,比较了 LPIPS (\downarrow)(顶部)、SSIM (\uparrow)(中部)和 PSNR (\uparrow)(底部)与处理时间 (\downarrow)。ODE 步数从左到右增加,分别代表 4、8、16、32、64、128 和 256。为了更好的可视化,使用 C=0C=0 时的 ODE 步数 4 和 8 被省略。圆圈表示所选的超参数。时间以对数刻度显示,单位为每张图像。

分析:

  • 校正步数 CC 的影响:
    • 从图中可以看出,增加校正步数 CC 会导致更长的评估时间(对于相同数量的 ODE 步数)。
    • C=0C=0 时(即不进行轨迹校正),LPIPS 值显著更高(更差),SSIM 和 PSNR 值更低(更差),尤其是在 ODE 步数较少时。这表明轨迹校正机制对于确保重建质量至关重要。
    • C=0C=0 增加到 C=1C=1 时,性能有显著提升,同时处理时间增加相对较小。
    • 进一步增加 CCC=2C=2C=4C=4 时,虽然 LPIPS、SSIM 和 PSNR 可能会有轻微改善,但处理时间也明显增加。性能提升的边际效益递减,例如,从 C=1C=1C=2C=2 的性能提升远不如从 C=0C=0C=1C=1
  • ODE 步数的影响:
    • 对于给定的 CC,增加 ODE 步数通常会提高 LPIPS、SSIM 和 PSNR,但也伴随着处理时间的增加。
    • 在 ODE 步数足够多时,性能会趋于饱和。
  • 权衡与最优选择:
    • 图表清晰地展示了快速采样和更好分数之间的权衡。

    • 较低的校正步数 (如 C=1C=1) 在相同性能水平下更具优势,因为它需要的 ODE 步数更少,从而更快。

    • 作者经验性地发现,C=1C=1 在高重建质量和快速处理之间提供了最佳的权衡。因此,在所有实验中,Restora-Flow 都将 CC 固定为 1。

    • 这意味着 Restora-Flow 只需要优化一个超参数:ODE 步数,这与标准的 Flow Matching 相同,简化了模型的使用和部署。

      消融实验有力地证明了轨迹校正机制的有效性,并支持了 C=1C=1 作为 Restora-Flow 默认设置的合理性。

7. 总结与思考

7.1. 结论总结

本文成功引入了 Restora-Flow,一种基于 Flow Matching 的、无需训练的图像修复算法,专门为掩码引导的图像修复任务设计。通过在 Flow Matching 模型的生成过程中有效地整合了掩码引导融合和新颖的轨迹校正机制,Restora-Flow 展现了在多种任务和数据集上的卓越性能。实验结果表明,与基于扩散和流匹配的基线方法相比,Restora-Flow 在感知质量 (LPIPS) 和处理速度方面均取得了显著优势。此外,该方法具有极简的超参数设置,除了标准 Flow Matching 所需的 ODE 步数外,仅需固定一个校正步数 C=1C=1,极大地简化了其应用和部署。Restora-Flow 为高效、高质量的图像修复提供了一个强大的新工具。

7.2. 局限性与未来工作

作者指出了当前 Restora-Flow 的主要局限性在于其专注于涉及基于掩码的退化操作符 (mask-based degradation operators) 的图像修复任务。未来的工作计划将该算法扩展到处理不涉及掩码的退化操作符的图像修复任务。这将使其能够应对更广泛的逆问题类型,例如去除运动模糊、图像压缩伪影等。

7.3. 个人启发与批判

7.3.1. 个人启发

  1. 效率与质量的平衡: Restora-Flow 的成功再次强调了在生成模型应用中,效率(特别是采样时间)与生成质量之间寻求最佳平衡的重要性。Flow Matching 相较于扩散模型在速度上的优势,以及 Restora-Flow 进一步通过巧妙机制达成的加速,为未来的生成模型研究提供了方向。
  2. “训练自由”的潜力: 无需训练 (training-free) 的方法极大地降低了应用门槛和计算成本。通过充分利用预训练的通用生成先验,并结合任务特定的引导机制,可以在不重新训练模型的情况下解决新的下游任务。这种范式对于资源受限或需要快速部署的场景具有巨大吸引力。
  3. 机制设计的精妙: 轨迹校正 (trajectory correction) 机制是 Restora-Flow 的核心创新点。它解决了在生成过程中强制引入外部信息(如掩码引导)可能导致样本偏离模型学习的流形这一关键问题。这种先“大胆融合”,再“巧妙修正”的思路,值得在其他生成模型引导方法中借鉴。
  4. 跨领域适用性: Restora-Flow 在自然图像和医学图像数据集上的优异表现,证明了其方法的通用性和鲁棒性。特别是在医学图像领域,高效且高质量的修复对于诊断和分析具有重要价值。

7.3.2. 批判

  1. “训练自由”的语境: 尽管 Restora-Flow 自身是训练自由的,但它高度依赖于一个高质量的、预训练的无条件 Flow Matching 模型。这个基础模型的训练成本和数据需求并未消失,只是被转移到了上游。如果基础模型质量不高,Restora-Flow 的性能也会受限。
  2. 轨迹校正的理论基础: 尽管轨迹校正机制在经验上表现良好,但其背后是否存在更严格的理论推导或收敛性保证,以解释为何这种外推和重新加噪的组合能有效将样本拉回流形,可能值得进一步探讨。目前来看,其成功更多是基于经验观察。
  3. ODE 步数和 CC 值的鲁棒性: 论文虽然指出 C=1C=1 是一个好的经验选择,且只需要调整 ODE 步数,但这些参数在面对全新的、差异巨大的退化类型或数据集时,是否依然能保持其鲁棒性和最优性,仍需更多验证。例如,当退化程度非常剧烈时,单一的 C=1C=1 是否足够?
  4. 非掩码退化的挑战: 论文自身也指出了其局限性在于主要处理基于掩码的退化。将方法扩展到非掩码退化(例如,复杂的非线性模糊、去雨、去雾等)将是巨大的挑战,因为此时“已知区域”的概念不再明确,数据保真项的构建和引导策略需要完全不同的设计。
  5. 与其他先进扩散模型方法的对比: 尽管论文对比了 DDNM+DDNM+RePaint,但扩散模型领域仍在快速发展,出现了许多更先进的采样策略(如 DDIM、DDPM++ 等)和针对逆问题的更精细的引导技术。与这些最前沿的扩散模型方法进行更全面的对比,可能会提供更全面的视角。

相似论文推荐

基于向量语义检索推荐的相关论文。

暂时没有找到相似论文。