1. SLER-IR框架概述:全能图像修复的新范式
图像修复技术作为计算机视觉领域的基石任务,其核心目标是从各种退化输入(如噪声、雨雾、模糊等)中恢复出高质量图像。传统方法通常采用"一个任务一个模型"的设计思路,例如专门针对去噪的DnCNN、针对去雨的PReNet等。这类方法虽然在特定任务上表现优异,但在面对真实世界中复杂多变的退化场景时,往往显得力不从心。更关键的是,当遇到训练数据中未出现的新退化类型时,传统方法需要重新训练整个模型,这严重限制了其实际应用价值。
近年来,全能(All-in-One)图像修复框架逐渐成为研究热点。这类方法试图通过单一模型处理多种退化类型,其技术路线主要分为两大流派:特征调制方法(如PromptIR、AirNet)通过提示向量或对比学习来调整共享参数空间;架构自适应方法(如MoCE-IR)则采用混合专家(Mixture of Experts,MoE)机制激活任务特定的处理模块。然而,现有方法仍面临三个关键挑战:
- 特征干扰问题:共享主干网络中的特征在处理不同退化类型时会产生冲突,例如抑制噪声与恢复高频细节之间的目标矛盾。
- 专家利用不足:现有路由机制难以充分利用退化先验信息,专家模块往往局限于局部处理,缺乏全局协调。
- 空间非均匀性:实际退化(如雨线、局部雾霾)通常具有空间变化特性,而现有方法多采用全局统一处理策略。
针对这些挑战,我们提出了SLER-IR(Spherical Layer-wise Expert Routing for Image Restoration)框架,其核心创新体现在三个方面:
球形分层专家路由机制:将传统网络中的每个处理层扩展为多个参数独立的专家节点,通过超球面退化表示学习实现精确的专家选择。如图1所示,这种设计允许模型为不同退化类型动态组合专家路径,形成"渐进式专业化"的处理流程。
超球面退化表示学习:通过将退化特征映射到单位超球面并进行对比学习优化,有效消除了线性嵌入空间中的几何偏差,使相似度计算更准确地反映真实的退化关联性。
全局-局部粒度融合(GLGF):创新性地整合全局语义(CLS token)与局部退化线索(patch token),解决了基于patch的训练与全图推理之间的粒度差异问题,特别适合处理空间变化的退化模式。
实验结果表明,SLER-IR在PSNR和SSIM指标上全面超越现有方法。例如在五任务(去雾、去雨、去噪、去模糊、低光增强)设定下,平均PSNR达到31.73dB,比之前最佳方法(MoCE-IR)提升1.15dB。更重要的是,这种优势在不同退化类型和不同强度下都保持稳定,证明了框架的鲁棒性。
核心设计哲学:SLER-IR的核心理念是将"退化感知"与"内容适应"解耦。通过球形路由专注于退化特性识别,通过GLGF模块保持内容一致性,两者协同工作实现了"对症下药"式的智能修复。
2. 核心技术解析:球形分层专家路由
2.1 专家网络架构设计
SLER-IR采用4级编码器-解码器结构,每级包含不同数量的Transformer块(分别为4、6、6、8个)。与传统架构的关键区别在于:我们将每个标准处理块替换为三个参数独立的专家节点,形成多专家主干网络。这些专家共享相同的输入输出接口,但内部参数和特征表示各不相同,专门针对不同类型的退化特性进行优化。
专家多样性保障:为避免专家同质化,我们采用两种策略:(1)初始化阶段使用正交初始化保证专家参数的差异性;(2)训练过程中引入专家差异损失(expert diversity loss),鼓励不同专家关注不同的特征维度。具体实现是通过计算专家输出特征的余弦相似度矩阵,并最小化非对角线元素:
$$ \mathcal{L}{div} = \frac{1}{C(C-1)}\sum{i\neq j}|\cos(E_i(x),E_j(x))| $$
其中$C$为专家数量,$E_i$表示第$i$个专家。
2.2 两阶段路由策略
路由机制是MoE架构的核心,SLER-IR采用创新的两阶段路由策略:
阶段一:概率路由(训练前15epoch)
- 通过退化提取器获取输入图像的退化特征
- 使用MLP将特征投影到路由空间并进行L2归一化
- 计算归一化特征与专家中心的余弦相似度
- 应用行方向Softmax得到选择概率分布
这个阶段采用"软路由"策略,即所有专家都参与前向计算并按概率加权组合输出。虽然计算开销较大,但能保证所有专家都获得有效的梯度信号,避免某些专家因初期表现不佳而被完全忽略的"专家死亡"问题。
阶段二:确定路由(后续80epoch)
- 冻结路由器的参数
- 对每层选择概率最高的专家(argmax)
- 仅激活被选中的专家进行前向计算
这种"硬路由"策略虽然在训练初期会导致性能下降,但在路由器优化充分后,能显著降低计算成本(实际仅需约1.5倍单专家计算量,而非全专家计算的3倍)。
2.3 超球面退化表示学习
传统路由方法存在两个根本性局限:(1)将退化类型硬性分类到有限专家,忽略细粒度相似性;(2)线性嵌入空间的几何偏差会扭曲相似度计算。如图2所示,SLER-IR通过超球面嵌入解决了这些问题。
超球面投影:对原始退化特征$f$进行L2归一化: $$ \hat{f} = \frac{f}{|f|_2} $$ 这使得所有特征向量都落在单位超球面上,相似度计算转化为纯角度比较: $$ s_j = \cos(\hat{f},\hat{c}_j) = \hat{f}^\top\hat{c}_j $$
对比学习优化:采用三元组约束的对比损失函数: $$ \mathcal{L}_{HC,l} = \max(\mathbb{E}N[S{i,j,l}] - \mathbb{E}P[S{i,j,l}] + \lambda, 0) $$ 其中$\mathbb{E}_P$和$\mathbb{E}_N$分别表示正样本对和负样本对的平均相似度,$\lambda$为间隔超参数(默认0.5)。该损失函数推动同类退化样本在球面上聚集,异类样本相互远离。
专家中心均匀化:为避免专家中心在球面上聚集,我们引入蒙特卡洛正则化:定期从训练集中采样特征,计算其与专家中心的最小夹角,并通过梯度惩罚鼓励中心均匀分布: $$ \mathcal{L}_{uniform} = \exp(-\min_j \angle(\hat{f}_i,\hat{c}_j)) $$
3. 全局-局部粒度融合(GLGF)模块
3.1 动机与挑战
真实场景中的退化往往具有空间不均匀性,例如:
- 雨线通常只出现在局部区域
- 雾霾浓度随景深变化
- 运动模糊在不同物体区域表现不同
同时,现有方法多采用patch级训练(如256×256裁剪)但需处理全图推理,这导致训练-测试粒度差异(granularity gap)。GLGF模块的设计正是为了应对这些挑战。
3.2 双路径特征提取
GLGF包含两个并行分支:
内容分支(全局语义):
- 使用ViT架构处理完整输入图像
- 提取CLS token作为全局场景表征$c$
- 提取patch token$T_c \in \mathbb{R}^{H×W×d}$作为局部内容特征
- 通过FiLM调制校准局部特征: $$ \gamma, \beta = \text{MLP}_c(c), \quad \hat{T}_c = \text{LN}(\gamma \odot T_c + \beta) $$
退化分支(局部证据):
- 将图像划分为4×4网格,每格独立输入退化提取器
- 获取16个区域级退化CLS token${d_i}_{i=1}^{16}$
- 重排为低分辨率退化图$E \in \mathbb{R}^{4×4×d}$
- 双线性上采样得到DSP map$D \in \mathbb{R}^{H×W×d}$
3.3 内容引导的退化融合
通过交叉注意力机制整合两条路径: $$ F' = \hat{T}_c + \text{Attn}(Q=\hat{T}_c, K=D, V=D) $$ 其中注意力计算采用缩放点积形式: $$ \text{Attn}(Q,K,V) = \text{Softmax}(\frac{QK^\top}{\sqrt{d}})V $$
这种设计使得全局语义可以指导局部退化特征的解读,例如:
- 在天空区域,模糊更可能是雾霾而非运动模糊
- 在文字区域,锐化处理需要更谨慎以避免伪影
3.4 分层特征注入
将融合后的先验图$F'$注入到修复主干的各个层级:
- 通过线性投影对齐$F'$与当前层特征$F_l$的维度
- 以先验特征为query,主干特征为key/value进行交叉注意力
- 残差连接融合结果: $$ F_{l}^{out} = F_l + \text{MLP}(\text{Attn}(Q=F', K=F_l, V=F_l)) $$
这种注入方式允许网络在不同深度利用先验信息:浅层侧重局部纹理修复,深层关注全局结构一致性。
4. 实验分析与工程实践
4.1 数据集与评估协议
我们采用学术界广泛使用的五个基准数据集:
- 去雾:SOTS(O-HAZE)
- 去雨:Rain100L
- 去噪:CBSD68(σ=15,25,50)
- 去模糊:GoPro
- 低光增强:LOL
训练策略:
- 三任务设定:去雾、去雨、去噪联合训练
- 五任务设定:加入去模糊和低光增强
- 输入patch尺寸:256×256
- 批量大小:阶段一为10,阶段二为20
- 优化器:Adam(β1=0.9, β2=0.999)
- 初始学习率:2e-4,余弦衰减
4.2 性能对比实验
表1展示了五任务设定下的定量结果(PSNR/dB):
| 方法 | 去雾 | 去雨 | 去噪(σ=25) | 去模糊 | 低光 | 平均 |
|---|---|---|---|---|---|---|
| Restormer | 24.09 | 34.81 | 31.49 | 27.22 | 20.41 | 27.60 |
| AirNet | 21.04 | 32.98 | 30.91 | 24.35 | 18.18 | 25.49 |
| PromptIR | 26.54 | 36.37 | 31.47 | 28.71 | 22.68 | 29.15 |
| MoCE-IR | 30.48 | 38.04 | 31.34 | 30.05 | 23.00 | 30.58 |
| SLER-IR | 33.43 | 38.47 | 31.50 | 31.27 | 23.96 | 31.73 |
关键发现:
- 在去雾任务上优势最显著(+2.95dB vs MoCE-IR)
- 即使对于传统强项如去噪(σ=25),仍保持领先(31.50 vs 31.34)
- 计算开销仅比单任务模型高约50%,远低于激活所有专家的理论值
4.3 消融实验分析
路由机制对比:
- 分类器路由:28.76dB
- 球形对比路由:31.73dB 说明几何一致的相似度度量对专家选择至关重要。
专家数量影响:
- 2专家:31.22dB
- 3专家:31.73dB
- 4专家:31.32dB 表明适度增加专家数量有益,但需平衡计算成本。
GLGF模块有效性:
- 移除GLGF:30.64dB
- 完整模型:31.73dB 验证了全局-局部融合对复杂退化的适应性。
4.4 实际部署建议
计算资源权衡:
- 边缘设备:可使用2专家配置,牺牲少量性能换取更低延迟
- 云端部署:推荐3专家配置,利用并行计算优势
新任务扩展:
- 冻结主干,仅训练新的专家中心
- 采用线性探测(linear probing)策略逐步解冻参数
混合精度训练:
- 专家内部使用FP16
- 路由计算保持FP32以确保稳定性
缓存优化:
- 对常用退化类型缓存专家路径
- 实现"预热"推理,避免动态路由的开销波动
5. 局限性与未来方向
尽管SLER-IR取得了显著进展,但仍存在一些值得改进的方面:
动态计算分配:当前所有专家参数量相同,未来可探索异构专家设计,为不同复杂度退化分配不同容量的专家。
跨模态退化:当前主要处理视觉退化,对跨模态问题(如文本叠加、二维码遮挡)处理能力有限。
极端退化恢复:对于严重退化(如90%像素丢失),性能仍有提升空间。
未来工作可能沿着以下方向展开:
- 可解释路由:可视化专家选择决策过程,增强模型透明度
- 终身学习:在不遗忘旧任务的前提下持续学习新退化类型
- 3D扩展:将球形路由思想应用于视频修复,利用时间维度信息
一个特别有前景的方向是退化-内容协同路由,即同时考虑图像内容和退化类型来动态调整网络结构。初步实验表明,这种双重注意力机制可以进一步提升对复杂场景的适应能力。