本文提出了一种用于统一图像恢复(all-in-one image restoration, AIR)的退化感知残差条件最优传输(Degradation-Aware Residual-Conditioned Optimal Transport, DA-RCOT)方法。该研究由西安交通大学数学与统计学院的Xiaole Tang、Xiang Gu、Xiaoyi He、Xin Hu以及Jian Sun(通讯作者)共同完成,发表在《IEEE Transactions on Pattern Analysis and Machine Intelligence》2025年8月第47卷第8期。该工作是其ICML会议版本的扩展,通过引入多尺度残差嵌入和统一最优传输框架,显著增强了对多种退化类型同时处理的能力。
在学术背景方面,图像恢复是低层计算机视觉中的基础任务,传统方法通常针对单一退化类型训练特定模型,例如去噪、去雨、去雾、去模糊或低光增强。然而,真实应用场景中退化常常未知且混合出现,单一任务模型难以实用。因此,一体化图像恢复(AIR)应运而生,旨在使用一个统一模型处理多种退化。现有AIR方法主要通过在多域退化-干净图像对上拟合回归模型,并利用从退化图像中提取的退化嵌入来指导恢复。然而,由于异质退化反演的严重不适定性,这些方法往往对退化语义感知不足,且依赖成对数据进行监督训练,导致恢复结果结构受损、缺乏对真实世界或不成对数据的适应性。为突破这些限制,本研究从最优传输(optimal transport, OT)视角出发,将AIR建模为将多域退化图像分布传输到干净图像分布的最优传输问题,并引入传输残差(transport residual)作为退化特异性线索,同时用于传输代价和传输映射。
在方法流程上,DA-RCOT首先提出傅里叶残差引导的最优传输(Fourier residual-guided optimal transport, FROT)目标。该目标在传输代价中加入了残差正则项,利用传输残差在频域中的稀疏性来刻画不同类型退化的特异性模式。具体而言,对于除噪声外的退化,使用残差傅里叶变换的L1范数作为正则;对于去噪任务,则使用L2范数。该方法在Kantorovich问题的对偶形式下推导出极小极大优化问题,并证明该优化问题的鞍点可提供最优传输映射。在FROT框架中,仅需不成对数据即可求解;为进一步利用成对数据,又提出了成对信息增强的FROT目标,在传输代价中加入成对距离项。整个框架通过对抗训练两个神经网络来求解,其中传输网络作为生成器,势函数网络作为判别器。
DA-RCOT的核心创新在于将传输映射设计为两阶段(two-pass)结构。第一阶段,生成器无条件地生成中间恢复结果,并计算中间传输残差。第二阶段,该残差经过一个多尺度残差嵌入生成模块(Residual Embeddings Generation Module, REGM)编码为多尺度残差嵌入,这些嵌入在生成器的不同尺度上作为条件,指导更精细的恢复。REGM包含基于门控的Transformer模块,包括多深度卷积头转置注意力(MDTA)和门控深度卷积前馈网络(GDFN),用于选择性地传播关键特征。为增强退化语义感知,研究者对最密集尺度的嵌入施加对比损失,使同一退化类型的嵌入在特征空间中聚类,同时保留不同噪声水平间的层级关系。通过t-SNE可视化,DA-RCOT的残差嵌入相比PromptIR的提示嵌入在任务间显示出更清晰的分离,并在噪声水平间呈现有序的层级分布。
在实验部分,DA-RCOT在合成与真实数据集上进行了全面评估,涵盖三种退化(去雾、去雨、去噪)和五种退化(增加去模糊、低光增强)的一体化设置。在三退化设置下,DA-RCOT在无成对训练条件下平均PSNR提升2.11 dB,LPIPS降低0.030;在成对设置下,相比采用相同骨干网络(Restormer)的PromptIR,平均PSNR提升0.8 dB,FID降低14.44。在五退化设置下,DA-RCOT相比InstructIR平均PSNR提升0.85 dB,相比DA-CLIP平均FID降低10.13。视觉对比显示,DA-RCOT能够更有效地去除退化并保留更真实的结构内容。
在任务特定设置下,DA-RCOT同样表现优异。在Rain100L去雨任务上,其PSNR超越PromptIR 1.94 dB;在SPANet真实去雨数据集上,超越DA-CLIP 2.49 dB;在SOTS和O-HAZE去雾任务上,相比DehazeFormer分别提升0.44 dB和2.22 dB。此外,DA-RCOT在鲁棒性实验中展现了良好的退化类型数量扩展能力:随着训练退化类型从1种增加到5种,DA-RCOT在去雨任务上的性能下降幅度最小,且始终优于其他方法。在未见的退化水平泛化实验中,DA-RCOT在跨雨强(Rain100L与Rain100H互换)和未见噪声水平(σ=60和σ=75)上均表现出优越的泛化能力。在真实世界场景泛化评估中,DA-RCOT在O-HAZE上提升1.33 dB,在SPANet上提升1.06 dB,并在混合退化图像上获得更低的NIQE和PIQE指标。
消融研究进一步验证了各模块的作用。残差嵌入条件机制(REC)相较于无条件或仅以中间结果或原始残差为条件的模型,在五退化设置下取得更优性能;FROT目标中的残差正则化也带来明显增益;对比损失在不成对设置下对任务语义解耦贡献显著,使无成对FROT训练的感知质量达到与成对L1训练相当的水平。REGM中的门控Transformer模块与对比损失共同作用,使残差嵌入在不同退化类型间形成清晰可分的簇。
本研究的价值在于将最优传输理论系统性地引入统一图像恢复问题,通过传输残差同时指导代价函数和传输映射,有效提升了对异质退化语义的感知能力和结构保持能力。该方法不仅在合成数据上取得了最先进性能,还在真实世界混合退化和未见退化水平场景中表现出优异的鲁棒性与泛化性,具有重要的理论意义和实际应用前景。其局限性在于,当噪声水平极端严重时,作为条件的残差可能会干扰原始特征,导致泛化性能下降。未来工作将探索将OT模型适配到多模态一对多生成和翻译问题中。