Exploiting Diffusion Prior for Task-driven Image Restoration¶
会议: ICCV 2025
arXiv: 2507.22459
代码: 无
领域: 图像修复
关键词: 任务驱动图像复原, 扩散先验, 高层视觉任务, 部分扩散, 短步去噪
一句话总结¶
提出 EDTR 方法，通过预修复+部分扩散和短步去噪策略，有效利用扩散模型先验恢复与高层视觉任务相关的细节，在复杂退化场景下显著提升分类、分割和检测性能。
研究背景与动机¶
真实世界图像经常受到多种退化因素的影响（下采样、模糊、噪声、JPEG 压缩等），导致分类、分割、检测等高层视觉任务性能严重下降。图像复原（IR）看似可以作为预处理方案，但研究表明简单地将 IR 作为前处理并不能有效恢复任务关键信息。
任务驱动图像复原（Task-driven Image Restoration, TDIR）由此萌生——其目标是以提升下游任务性能为导向进行图像恢复。然而现有 TDIR 方法面临三大挑战：
单一退化局限：以往 TDIR 方法通常针对单一退化类型（如仅超分辨率或仅去噪），难以应对真实场景中多种退化并存的情况
退化严重时线索不足：当图像被多种复杂因素严重退化时，可供恢复的线索极少，传统方法难以从中挖掘任务相关细节
扩散先验的误用：虽然 StableDiffusion 提供了强大的自然图像先验，但标准的 SD-based IR 流程（从纯噪声出发、50 步去噪）会生成视觉上看似合理但缺乏任务相关细节的结果——例如生成了错误的眼睛形状导致鸟类分类错误
本文识别出两个核心问题：(1) 从纯噪声生成阻碍了对 LQ 图像中有用线索的直接利用；(2) 长步去噪引入冗余细节，稀释了关键的任务相关信息。
方法详解¶
整体框架¶
EDTR 由三个核心组件构成：预修复+部分扩散、短步去噪、以及专用于任务驱动训练的损失函数。使用 StableDiffusion-2.1 作为骨干，ControlNet 作为可训练适配器，SwinIR 作为像素误差预修复网络。
关键设计¶
-
预修复与部分扩散（Pre-restoration and Partial Diffusion）:
- 功能：将 LQ 图像通过像素级修复后直接注入扩散过程，而非从纯噪声开始生成
- 核心思路：先用 SwinIR 对 LQ 图像进行预修复 \(z_{\text{pre-res}} = \mathcal{E}(\mathcal{R}_{\text{pix}}(I_{\text{LQ}}))\)，然后在预修复结果上加入轻微噪声启动扩散过程：\(z_{t,\text{partial}} = \sqrt{\bar\alpha_t} \cdot z_{\text{pre-res}} + \sqrt{1-\bar\alpha_t} \cdot \epsilon\)，其中时间步 \(t_p = 200\)（远小于 \(T=1000\)）
- 设计动机：SD 并非为处理复杂退化而训练，直接输入退化图像效果差。预修复消除了大部分退化伪影，使得 SD 能在相对干净的基础上补充任务相关的高频细节。部分扩散保留了 LQ 图像中的有用信息，不会像纯噪声起点那样完全依赖 SD 的生成能力
-
短步去噪（Short-step Denoising）:
- 功能：在推理时仅使用极少步数（1 步或 4 步）完成去噪，而非传统的 50 步
- 核心思路：对于 \(n\) 步去噪，构建时间步序列 \(\mathcal{T} = [t_p, \lfloor t_p/n \cdot (n-1)\rfloor, \ldots, \lfloor t_p/n \rfloor]\)。1 步去噪时简化为：\(z_{\text{diff-res}} = \frac{z_{t_p,\text{partial}} - \sqrt{1-\bar\alpha_{t_p}} \epsilon_\theta(z_{t_p,\text{partial}}, t_p, z_{\text{pre-res}})}{\sqrt{\bar\alpha_{t_p}}}\)
- 设计动机：实验发现增加去噪步数虽然提升视觉质量，但反而降低任务性能。每一步去噪都引入扩散先验中的细节，反复添加会生成与任务无关的冗余纹理，稀释任务关键信息。短步去噪保留了预修复的有用信息，仅通过扩散先验做最小限度的细节补充
-
任务驱动训练损失:
- 功能：设计专门的损失函数引导扩散先验恢复任务相关细节
- 核心思路：高层特征（HLF）损失 \(\mathcal{L}_{\text{HLF}} = \frac{1}{2}(\|\mathcal{H}^f(I_{\text{EDTR}}) - \mathcal{H}^f(I_{\text{HQ}})\|_1 + \|\mathcal{H}^f_{\text{HQ}}(I_{\text{EDTR}}) - \mathcal{H}^f_{\text{HQ}}(I_{\text{HQ}})\|_1)\)，使用两个任务网络（当前训练的 \(\mathcal{H}\) 和预训练在 HQ 上的 \(\mathcal{H}_{\text{HQ}}\)）的中间特征空间来度量复原质量
- 设计动机：噪声预测损失 \(\mathcal{L}_\epsilon\) 并非为任务驱动设计；HLF 损失直接在任务相关的特征空间中约束复原质量，实现了"以任务为导向"的修复
损失函数 / 训练策略¶
- 交替训练：每个迭代中先用 \(\mathcal{L}_\text{HLF}\) 更新 EDTR，再用 \(\mathcal{L}_\text{task} + \alpha \mathcal{L}_\text{FM}\) 更新任务网络 \(\mathcal{H}\)
- 任务损失：\(\mathcal{L}_\text{task} = f_\text{task}(\mathcal{H}(I_{\text{EDTR} \oplus \text{HQ}}), y)\)，使用半批 EDTR 复原图像和半批 HQ 图像联合训练，稳定训练过程
- 特征匹配损失：\(\mathcal{L}_\text{FM} = \|\mathcal{H}^f(I_{\text{EDTR} \oplus \text{HQ}}) - \mathcal{H}^f_\text{HQ}(I_\text{HQ})\|_1\)，相当于跨质量的知识蒸馏
- Wavelet 色彩校正：解码后用小波变换分离高低频，高频取自扩散结果，低频取自预修复结果：\(I_\text{EDTR} = \mathbf{H}(\mathcal{D}(z_\text{diff-res})) + \mathbf{L}(\mathcal{R}_\text{pix}(I_\text{LQ}))\)
实验关键数据¶
主实验¶
图像分类（CUB200）+ 语义分割（VOC2012）+ 目标检测（VOC2012）
退化类型 Mixture-A: SR(s=8) + JPEG(q=75)
| 方法 | 分类 Acc↑ (%) | 分割 mIoU↑ (%) | 检测 mAP↑ (%) |
|---|---|---|---|
| Oracle (HQ) | 82.5 | 67.0 | 36.9 |
| No restoration | 60.8 | 47.4 | 14.5 |
| SwinIR | 70.0 | 56.2 | 21.0 |
| SR4IR | 71.5 | 56.2 | 22.3 |
| EDTR-1 step | 74.4 | 64.1 | 31.2 |
| EDTR-4 step | 74.1 | 65.3 | 33.4 |
消融实验¶
Mixture-B (更复杂退化) 上的任务性能
| 方法 | 分类 Acc↑ (%) | 分割 mIoU↑ (%) | 检测 mAP↑ (%) |
|---|---|---|---|
| No restoration | 47.6 | 40.2 | 0.0 |
| SwinIR | 60.7 | 50.4 | 22.1 |
| SR4IR | 63.4 | 51.0 | 22.6 |
| EDTR-1 step | 68.8 | 60.4 | 30.6 |
| EDTR-4 step | 68.4 | 62.9 | 33.4 |
视觉质量指标对比（Mixture-A）
| 方法 | NIQE↓ | Q-Align↑ | PSNR↑ |
|---|---|---|---|
| SwinIR | 9.97 | 2.55 | 25.21 |
| SR4IR | 5.41 | 3.37 | 24.07 |
| EDTR-1 step | 4.68 | 3.53 | 23.63 |
| EDTR-4 step | 4.26 | 3.81 | 22.46 |
关键发现¶
- EDTR 在检测任务上提升最为显著：Mixture-A 上从无修复的 14.5% 提升到 33.4%（+18.9%），Mixture-B 上从 0.0% 提升到 33.4%
- PSNR 与任务性能呈负相关：EDTR 的 PSNR 低于 SwinIR，但所有任务指标均大幅领先，说明像素忠实度不等于任务有效性
- 短步去噪对任务至关重要：1 步和 4 步效果相近，但大幅优于传统 50 步
- 视觉质量（NIQE/Q-Align）和任务性能可以同时提升，打破了两者必须权衡的认知
亮点与洞察¶
- PSNR 与任务性能的悖论：首次在 TDIR 框架中清晰展示了像素级指标和任务级指标的不一致性，具有重要方法论价值
- 部分扩散的优雅设计：预修复 → 加轻噪 → 短步去噪的流程既保留了 LQ 图像的有用信息，又利用了扩散先验补充高频细节
- 1 步去噪即可：颠覆了扩散模型"步数越多越好"的常识，在 TDIR 场景下 1 步已足够
- 通用性：同一框架适用于分类、分割、检测三种截然不同的任务
局限与展望¶
- 依赖预训练 StableDiffusion，计算开销较大（需要 SD 推理 + ControlNet + 任务网络联合训练）
- 需要 HQ 图像进行训练，在完全无监督的真实场景下不适用
- 仅验证了三种视觉任务，对于更细粒度的任务（如关键点检测、实例分割）有效性未知
- 退化类型在训练时需要预定义，面对未知退化的泛化能力有待验证
- 预修复网络 SwinIR 的质量直接影响最终结果
相关工作与启发¶
- SR4IR：提出了 TDIR 的联合训练范式和任务驱动感知损失，是本文的直接前置工作
- StableSR / DiffBIR：展示了 SD 作为图像先验的强大能力，但未考虑任务驱动
- ControlNet：提供了对冻结 SD 进行条件控制的工具
- 启发：扩散模型的生成能力需要被"驯服"才能为特定任务服务，盲目追求视觉质量反而可能损害任务性能
评分¶
- 新颖性: ⭐⭐⭐⭐ 首次在 TDIR 中有效利用扩散先验，部分扩散 + 短步去噪设计巧妙
- 实验充分度: ⭐⭐⭐⭐ 三种任务 × 两种退化类型，消融全面，但数据集多样性可进一步扩展
- 写作质量: ⭐⭐⭐⭐ 动机阐述清晰，方法流程明了，但符号偶有冗余
- 价值: ⭐⭐⭐⭐ 为如何正确使用扩散先验提供了有价值的指引，检测性能提升尤为突出