[论文解读] DriftRec: Adapting diffusion models to blind JPEG restoration

Simon Welker, Henry N. Chapman+1 more|arXiv (Cornell University)|Nov 12, 2022

Advanced Neuroimaging Techniques and Applications参考文献 51被引用 5

一句话总结

DriftRec 通过修改前向随机微分方程（SDE），提出了一种针对无监督 JPEG 图像修复的新型得分扩散模型适配方法，使其更契合受损图像的分布。与以往方法不同，该方法仅需成对的干净-受损图像数据，且在无需了解 JPEG 质量因子或损坏过程的情况下仍能有效工作，在低质量因子（QF ≤ 30）下实现了更优的感知质量与分布建模效果。

ABSTRACT

In this work, we utilize the high-fidelity generation abilities of diffusion models to solve blind JPEG restoration at high compression levels. We propose an elegant modification of the forward stochastic differential equation of diffusion models to adapt them to this restoration task and name our method DriftRec. Comparing DriftRec against an $L_2$ regression baseline with the same network architecture and state-of-the-art techniques for JPEG restoration, we show that our approach can escape the tendency of other methods to generate blurry images, and recovers the distribution of clean images significantly more faithfully. For this, only a dataset of clean/corrupted image pairs and no knowledge about the corruption operation is required, enabling wider applicability to other restoration tasks. In contrast to other conditional and unconditional diffusion models, we utilize the idea that the distributions of clean and corrupted images are much closer to each other than each is to the usual Gaussian prior of the reverse process in diffusion models. Our approach therefore requires only low levels of added noise and needs comparatively few sampling steps even without further optimizations. We show that DriftRec naturally generalizes to realistic and difficult scenarios such as unaligned double JPEG compression and blind restoration of JPEGs found online, without having encountered such examples during training.

研究动机与目标

• 为解决在低质量因子（QF ≤ 30）下进行无监督 JPEG 修复的挑战，传统方法常产生模糊输出。
• 开发一种无需了解 JPEG 损坏过程（包括质量因子或量化矩阵）的扩散模型。
• 实现对复杂真实场景（如未对齐的双重 JPEG 压缩和处理历史未知的在线 JPEG）的泛化能力。
• 相比现有基于扩散模型与非扩散基线方法，提升感知图像质量与分布保真度。
• 证明修改后的 SDE 框架可在扩散框架内有效建模 JPEG 压缩带来的非高斯、非线性损坏。

提出的方法

• 通过修改扩散模型前向 SDE 的漂移项，使其反映干净图像与 JPEG 压缩图像之间的分布接近性，而非假设高斯先验。
• 使用（干净，JPEG 压缩）图像的成对数据集进行监督训练，训练与推理过程中均无需了解 JPEG 编码过程。
• 采用基于得分的生成模型（SGM），并使用去噪得分匹配（DSM）目标，学习从压缩图像到干净图像的逆过程。
• 使用 100 步采样次数的 Euler-Maruyama 采样器生成重建结果，实现图像分布的迭代优化。
• 引入后处理色彩校正步骤，将重建图像中每个 RGB 通道的均值与输入 JPEG 图像对齐，以减少色彩偏移。
• 利用连续 SDE 框架，支持灵活的采样策略，并为未来集成快速采样技术提供可能，尽管漂移项非标准。

实验结果

研究问题

• RQ1 能否在无需了解 JPEG 质量因子或量化矩阵的前提下，有效适配扩散模型用于无监督 JPEG 修复？
• RQ2 与标准扩散模型和 L2 回归基线相比，所提出的 SDE 修改在分布建模与感知质量方面有何改进？
• RQ3 DriftRec 在真实世界复杂 JPEG 痕迹（如未对齐的双重压缩或处理历史未知的在线图像）上的泛化能力如何？
• RQ4 在 DriftRec 中使用多次采样或采样结果平均时，感知质量与失真之间的权衡如何？
• RQ5 DriftRec 的运行效率与非扩散基线相比如何？近期的采样加速技术能否缩小这一差距？

主要发现

• DriftRec 在感知图像质量方面显著优于 L2 回归基线与最先进方法，避免了其他方法常见的模糊问题。
• 即使在低质量因子（QF ≤ 30）下，该方法也能忠实恢复干净图像的分布，表现为 FID、KID 和 LPIPS 指标更优。
• DriftRec 在未对齐双重 JPEG 压缩与处理历史未知的在线 JPEG 上表现出良好的泛化能力，证明其在真实世界无监督修复场景中的鲁棒性。
• 尽管 DNN 推理次数比非迭代方法（如 QGAC、FBCNN）多出 20–100 倍，DriftRec 的性能仍具竞争力，且有望通过未来快速采样技术进一步提升。
• 对 DriftRec 的多次采样结果进行平均，可实现感知质量与失真之间的可控权衡，为重建目标提供灵活性。
• 与对比方法相比，该模型参数量最少，且仅对 L2 回归基线进行了微小架构修改，表明其在模型设计上具有高度效率。