[论文解读] TAPE: Task-Agnostic Prior Embedding for Image Restoration

Lin Liu, Lingxi Xie+6 more|arXiv (Cornell University)|Mar 11, 2022

Advanced Image Processing Techniques被引用 4

一句话总结

该论文提出了一种任务无关先验嵌入（TAPE），一种基于Transformer架构的两阶段框架，通过在多样化干净图像数据上进行预训练，并在特定任务上微调，学习一种通用的、与任务无关的图像恢复先验。TAPE在多个恢复任务中将PSNR提升最高达1.45dB，并能有效泛化到未见过的退化类型，优于特定任务的模型。

ABSTRACT

Learning a generalized prior for natural image restoration is an important yet challenging task. Early methods mostly involved handcrafted priors including normalized sparsity, l_0 gradients, dark channel priors, etc. Recently, deep neural networks have been used to learn various image priors but do not guarantee to generalize. In this paper, we propose a novel approach that embeds a task-agnostic prior into a transformer. Our approach, named Task-Agnostic Prior Embedding (TAPE), consists of two stages, namely, task-agnostic pre-training and task-specific fine-tuning, where the first stage embeds prior knowledge about natural images into the transformer and the second stage extracts the knowledge to assist downstream image restoration. Experiments on various types of degradation validate the effectiveness of TAPE. The image restoration performance in terms of PSNR is improved by as much as 1.45dB and even outperforms task-specific algorithms. More importantly, TAPE shows the ability of disentangling generalized image priors from degraded images, which enjoys favorable transfer ability to unknown downstream tasks.

研究动机与目标

• 为解决学习通用且可迁移的图像先验的挑战，这些先验不针对任何单一恢复任务。
• 以一种能够将干净图像先验与退化伪影分离的方式，实现向未见图像恢复任务的迁移。
• 开发一个统一框架，在无需针对特定任务重新训练的情况下，提升多个低层次视觉任务的性能。
• 探索单一预训练先验模块是否能有效支持多种下游恢复任务。

提出的方法

• TAPE采用两阶段训练流程：在干净图像上进行与任务无关的预训练，随后在退化数据上进行特定任务的微调。
• 通过像素级对比损失训练先验学习模块（PLM），从干净图像中学习通用特征，并将其编码为查询嵌入。
• Transformer解码器利用这些查询嵌入关注输入中的相关特征，从而实现对退化图像的有效恢复。
• 预训练阶段使用多个恢复数据集（如去雨、去噪、去摩尔纹等），但不使用真实退化标注，仅依赖干净图像数据。
• 该方法通过学习与特定退化类型无关的先验表示，显式地将干净图像先验与退化伪影分离。
• 该框架兼容多种Transformer主干网络，可针对特定任务进行端到端微调。

实验结果

研究问题

• RQ1 能否从干净图像中学习一个单一的、与任务无关的先验，并有效迁移至多个图像恢复任务？
• RQ2 在干净图像上使用像素级对比损失进行预训练，是否能提升对未见退化类型的泛化能力？
• RQ3 将干净图像先验与退化伪影分离，对未知任务上的性能有何影响？
• RQ4 所提出的先验嵌入在PSNR和泛化能力方面，与特定任务模型相比，优势有多大？

主要发现

• 在Rain200L数据集上，TAPE将PSNR提升1.45dB，优于特定任务模型。
• 在去雪、去阴影、超分辨率和去模糊等未见任务上，TAPE分别实现0.91dB、0.29dB、0.41dB和0.48dB的PSNR增益。
• 在SIDD数据集上，即使未访问真实噪声图像，仅通过干净图像进行任务无关预训练，仍使PSNR提升0.31dB，证明了干净图像先验的价值。
• 消融实验表明，预训练网络θ和PLM模块比预训练主干φ更为关键，若省略则PSNR分别下降0.22dB和0.34dB。
• 预训练阶段的像素级对比损失增强了特征泛化能力，从而提升了在多样化任务上的性能。
• 可视化结果表明，PLM输出的特征在保留干净图像的结构和边缘信息的同时，有效抑制了退化伪影。