[TPAMI 2025] BaryIR：利用 Wasserstein 重心空间突破全能图像修复的泛化天花板

TL;DR

传统的图像修复模型通常是“专才”，而全能图像修复（All-in-One IR）目标是做一个“全才”。然而，现有的全能模型在面对没见过的退化（OOD）时经常“翻车”。BaryIR 通过数学上严谨的 Wasserstein Barycenter (WB) 理论，在特征空间找到了不同退化之间的“最大公约数”，实现了退化无关内容与退化噪声的完美解耦，不仅在已知任务上霸榜，在未知退化任务上的泛化能力更是远超前人。

背景定位：从“适配”到“发现不变性”

在全能图像修复领域，之前的 SOTA 路径大多在折腾 Prompt（提示词） 或者 Routing（路由专家系统）。这类方法的逻辑是：给模型一个信号告诉它现在是在去雨还是去噪。

但这带来了一个本质缺陷：过拟合训练分布。如果训练场景里只有雨、雾、噪，模型就会形成一种强烈的归纳偏置。当你丢给它一张带 JPEG 压缩伪影或者水下模糊的照片时，模型由于找不到匹配的“专家”，往往效果惨不忍睹。BaryIR 的作者敏锐地意识到：泛化的关键不在于更好地适配每种退化，而在于从所有退化中提取出那个“退化无关”的原始图像分布。

核心直觉：寻找分布的“重心”

作者认为，不管图像是被雨淋了还是被雾罩了，其底层的内容结构（Content）是不变的，变化的只是退化产生的偏移（Shift）。通过最优传输理论，可以找到一个 Wasserstein Barycenter (WB)，它是到多个退化分布距离之和最小的那个“中心分布”。

方法论拆解：

1. 连续重心映射 (Barycenter Map)：通过一个基于 Transformer 的映射网络 $T_{h}eta$，将多源退化特征 $z_k$ 转换为重心特征 $b_k$。
2. 双重空间解耦：
   • WB 空间：编码图像的本质内容（退化无关）。
   • 残留子空间 (Residual Subspace)：捕捉特定退化的知识 $r_k=z_k-b_k$。
3. 正交性约束：引入 BRO Loss，强制重心特征与残留特征正交，确保内容和退化信息互不干扰。

图 1：BaryIR 整体架构，展示了如何通过 WB 映射和残留空间实现特征解耦。

实验战绩：全线霸榜与极强的 OOD 韧性

在包含去雾、去雨、降噪、去模糊、低光增强的五项全能测试中，BaryIR 展现了统治级的实力。

1. 突破未见之境 (Generalization)

最令人惊艳的是在 未见退化类型 上的表现。即使模型从没学过如何处理水下图像或 JPEG 压缩伪影，BaryIR 依然能修复出清晰的纹理。在水下增强任务上，对比目前的强基线 MoCE-IR，BaryIR 的 PSNR 提升了超过 2dB。

图 2：在未见退化类型（如 JPEG 压缩和水下增强）上的视觉对比，BaryIR 恢复了更好的颜色和纹理。

2. 特征可视化的魔力

通过 t-SNE 降维可以看到，BaryIR 的重心特征（WB Embeddings）在空间中高度重合（说明确实做到了退化无关），而残留特征（Residual Embeddings）则按退化类型清晰地聚类（说明有效捕捉了退化特性）。这种分布上的整齐划一，正是其强大泛化能力的源头。

图 3：特征空间可视化，重心特征实现了完美的跨退化对齐。

深度洞察：为什么 Wasserstein 重心有效？

传统的特征对齐（如简单的 L2 距离或对比学习）往往忽略了特征空间的几何结构。Wasserstein 距离（又称推土机距离）考虑了概率分布之间的水平位移。在图像退化建模中，退化通常表现为像素或特征的系统性偏移。BaryIR 利用最优传输的几何直觉，找到各个退化分布的几何中心，这比简单的“特征平均”更能保留图像的高频结构和对比度。

总结与局限

BaryIR 成功地将复杂的数学理论转化为实用的深度学习框架。它不仅是一个更强的修复模型，更提供了一套处理多源异构数据的通用方法论。

局限性：在面对极其严重的复合退化（如暴雨+重雾+极低光）且存在大量信息丢失时，模型可能会为了维持全局的一致性而牺牲局部细节，产生一定的平滑感。

未来启示：这一思路完全可以扩展到多模态领域（文字、语音、图像），通过寻找跨模态的“重心”，或许能诞生更通用的多模态理解与生成模型。