跳转至

Continuous Expert Assembly: Instance-Conditioned Low-Rank Residuals for All-in-One Image Restoration

会议: NeurIPS 2026
arXiv: 2605.06127
OpenReview: https://openreview.net/forum?id=pQjphOzAGt
领域: 图像恢复
关键词: All-in-One Image Restoration, Continuous Expert Assembly, Dynamic Low-Rank Residuals, Cross-Attention Hyper-Adapter, Linear Attention Formulation

一句话总结

针对全能型图像恢复中全局提示特征瓶颈与离散 MoE 路由不稳定的痛点,本文提出 Continuous Expert Assembly (CEA),通过轻量交叉注意力超适配器按样本动态合成低秩路由基与残差方向,并在空间 token 级以有符号点积连续汇聚残差更新,显著提升了复合退化与空间异质退化的恢复质量。

研究背景与动机

真实场景下的图像退化极少以单一、均一的形式孤立出现。低光照、雾霾、雨条、运动模糊与传感器噪声往往在同一张拍摄画面中共存,具有显著的空间异质性与复合交互性,且在推理阶段完全未知具体的退化类型与严重程度。全能型(All-in-One, AiO)图像恢复的核心目标便是利用一套统一的网络权重,在无需测试期退化先验标签的前提下,自适应应对异构、局域多变且相互重叠的复合退化。

现有 AiO 恢复模型通常采用两类策略调节共享主干网络。第一类依托全局条件信号,例如显式退化表征、任务先验嵌入、可学习提示词(Prompts)或语言指令;这类方法虽然轻量紧凑,但本质上是将整幅图像的退化线索压缩为一个全局特征向量,在面对雨雾共存、强弱不均的空间局部退化时,极易丢失细粒度的空间差异信息。第二类则引入条件计算机制,如混合专家(Mixture-of-Experts, MoE)或低秩适配分支;然而传统稀疏 MoE 依赖离散的 Top-\(K\) 路由决策,不仅容易引起训练优化震荡与负载均衡开销,且其专家池在样本间是静态固定的,本质上仅是在离散算子空间中做硬切换或软混合,难以拟合连续多变的复合退化空间。

这两类范式的根本矛盾在于:全局提示压缩了局部退化细节,而静态专家网络池又受制于离散路由与样本无关的专家基底。本文的切入角度是打破“从静态专家库中离散挑专家”的传统思维,转为“根据当前输入图像动态生成低秩残差子空间,并由每个空间 token 自主完成连续残差汇聚”。核心 idea:通过轻量级 Cross-Attention Hyper-Adapter 探测局部退化上下文,动态合成样本自适应的低秩路由基 \(A(x)\) 与残差方向 \(B(x)\),各空间 token 通过无 Softmax 的稠密有符号点积亲和度完成 token 级连续残差装配,形成兼具空间自适应与线性注意力高效特性的参数残差注入机制。

方法详解

整体框架

Continuous Expert Assembly (CEA) 构建于非对称 U 形 Transformer 主干之上,主要作用于解码器块中的多头自注意力投影层(默认调制 Query 与 Key 投影)。整个恢复过程分为三个核心阶段:局部退化上下文压缩与空间探测、样本条件化低秩分量合成与 RankNorm 归一化、以及空间 Token 级的稠密有符号点积连续残差装配。

对于输入退化图像 \(x\),主干编码器提取多尺度特征。在解码器的特定层中,中间空间特征图 \(X^{(l)} \in \mathbb{R}^{N \times d_{\text{in}}}\) 首先经过轻量级卷积缩减空间分辨率并聚合局部结构模式;随后由一组可学习的紧凑查询向量通过交叉注意力探针提取退化证据,分别解码为该样本专属的低秩路由基 \(A_\tau^{(l)}(x)\) 与残差方向 \(B_\tau^{(l)}(x)\)。随后,各空间 token \(X_n^{(l)}\) 直接与路由基计算有符号内积亲和度,以连续加权方式组合对应的残差方向,生成专属的残差偏移量 \(\Delta Y_{\tau, n}^{(l)}\) 注入到基础线性投影中。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    In["退化输入图像 x<br/>(含未知/复合退化)"] --> Enc["U-Transformer 编码器<br/>多尺度特征抽取"]
    Enc --> Dec["解码器中间空间特征 X"]

    subgraph CEA_Module["Continuous Expert Assembly 流程"]
        direction TB
        Dec --> Cond["空间特征下采样压缩<br/>DW-PW Conv (stride=2)"]
        Cond --> Probe["交叉注意力超适配器<br/>可学习 Queries 空间探测"]
        Probe --> Head["解耦线性解码头<br/>生成 A(x) 与 B(x)"]
        Head --> Norm["秩方向归一化 RankNorm<br/>消除动态分解尺度模糊"]
        Norm --> Assembly["连续专家装配<br/>稠密有符号内积线性注意力"]
    end

    Dec --> BaseProj["基础注意力投影<br/>X * W (Q/K 投影)"]
    Assembly --> Residual["动态残差增量 ΔY"]
    BaseProj --> Add["残差注入: Y = X*W + ΔY"]
    Residual --> Add
    Add --> OutBlocks["后续解码与多尺度重建"]
    OutBlocks --> CleanOut["高质量恢复图像 y^"]

关键设计

1. 交叉注意力超适配器:解耦空间探测与样本自适应低秩因子生成

针对全局池化丢失局部退化证据、而纯局部超网络参数爆炸的矛盾,CEA 提出了 Cross-Attention Hyper-Adapter。在第 \(l\) 个解码器块中,先通过核大小为 \(3 \times 3\)、步长 \(s_c=2\) 的深度可分离卷积(DW-Conv)配合 \(1 \times 1\) 逐点卷积(PW-Conv)将特征图 \(X^{(l)} \in \mathbb{R}^{N \times d_{\text{in}}}\) 降采样为紧凑特征 \(\hat{X}^{(l)}\)。随后引入两组可学习查询向量 \(R_A^{(l)}, R_B^{(l)}\)(各 8 个 query),通过交叉注意力机制探测降采样特征中的退化线索: $\(T_\phi^{(l)}(x) = R_\phi^{(l)} + \text{CrossAttn}\left(R_\phi^{(l)}, \hat{X}^{(l)}, \hat{X}^{(l)}\right), \quad \phi \in \{A, B\}\)$ 探测得到的上下文向量 \(T_A^{(l)}(x)\) 与 \(T_B^{(l)}(x)\) 在不同目标投影间共享,再由轻量全连接头解码生成调制目标 \(\tau \in \{Q, K\}\) 所需的动态路由基 \(A_\tau^{(l)}(x) \in \mathbb{R}^{d_{\text{in}} \times r}\) 与残差方向矩阵 \(B_\tau^{(l)}(x) \in \mathbb{R}^{r \times d_{\text{out}}}\)。这种设计既保留了长程空间退化感知能力,又将计算开销约束在极低水平。

2. 连续专家装配机制:有符号亲和度驱动的免 Softmax 稠密残差合成

传统 MoE 依靠固定专家库和稀疏离散的 Top-\(K\) 路由(如公式 \(Y_n = X_n W + \sum g_i(X_n) \mathcal{E}_i(X_n)\)),容易引发路由硬截断和训练抖动。CEA 将专家参数化为动态生成的秩分量对 \((a_{\tau,k}^{(l)}(x), b_{\tau,k}^{(l)}(x))\),令每个空间 token \(X_n^{(l)}\) 通过连续的有符号点积计算对第 \(k\) 个秩分量的亲和度,并以此连续加权合成残差方向: $\(\Delta Y_{\tau,n}^{(l)} = \alpha \sum_{k=1}^r \langle X_n^{(l)}, a_{\tau,k}^{(l)}(x) \rangle b_{\tau,k}^{(l)}(x)\)$ 其中 \(\alpha = 1/r\) 为稳定缩放因子。矩阵形式下,全序列更新可表示为 \(\Delta Y_\tau^{(l)} = \alpha X^{(l)} A_\tau^{(l)}(x) B_\tau^{(l)}(x)\)。这在代数形式上完全等价于免 Softmax 的线性注意力 \(\alpha \text{Attn}_{\text{lin}}(X^{(l)}, A_\tau^{(l)}(x)^\top, B_\tau^{(l)}(x))\),其中输入 token 扮演 Query,生成的路由基充当 Key,残差方向充当 Value。每个 token 无需经过概率归一化即可自主决定各分量的正负激励与贡献大小,且计算可通过两次低秩矩阵乘法实现,复杂度仅为 \(\mathcal{O}(N d_{\text{in}} r + N r d_{\text{out}})\)。

3. 秩方向归一化 (RankNorm):消除动态分解的尺度漂移与训练失稳

在动态低秩分解中,矩阵乘积存在固有的标量尺度模糊:任意非零标量 \(c\) 均满足 \(A B = (c A)(B / c)\)。在反向传播过程中,超适配器容易导致 \(A\) 或 \(B\) 的范数单向发散或塌陷,使得 token 与路由基之间的内积亲和度数值分布剧烈漂移。为保证训练稳定性,CEA 提出了实例级 Rank-wise Factor Normalization (RankNorm)。在残差装配前,对生成的每一个秩分量独立做 \(\ell_2\) 范数归一化: $\((A_\tau^{(l)\prime}(x))_{:,k} = \frac{(A_\tau^{(l)}(x))_{:,k}}{\|(A_\tau^{(l)}(x))_{:,k}\|_2 + \epsilon}, \qquad (B_\tau^{(l)\prime}(x))_{k,:} = \frac{(B_\tau^{(l)}(x))_{k,:}}{\|(B_\tau^{(l)}(x))_{k,:}\|_2 + \epsilon}\)$ 此归一化在单个样本内执行,杜绝了批次内跨样本耦合;同时配合比例缩放系数 \(\alpha = 1/r\),将每个秩分量对残差更新的贡献尺度锚定在稳定数值区间内,有效避免了深层特征弥散与优化发散。

4. 注意力投影注入:针对性重塑退化自适应特征交互场

CEA 将生成的动态残差定向注入到 Transformer 块的多头注意力 Query (\(Q\)) 与 Key (\(K\)) 线性投影中: $\(Q = X^{(l)} W_Q + \Delta Y_Q^{(l)}, \qquad K = X^{(l)} W_K + \Delta Y_K^{(l)}\)$ 通过显式调制 \(Q\) 与 \(K\) 投影,CEA 能够在保持主干权重的通用表达能力的同时,直接重塑特征之间的注意力相似度图(Attention Map),引导模型在雨条边缘、雾浓度梯度过渡带或低照度暗区动态调整局部与长程上下文聚合权重。消融实验证实,同时调制 \(Q+K\) 能够最大化几何退化自适应增益,其表现明显优于仅调制 Value (\(V\)) 或前馈网络(FFN)。

损失函数 / 训练策略

模型在端到端训练过程中无需任何退化类型标签或辅助分类损失,仅采用空间域与频域联合重构损失: $\(\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{rec}} + \lambda_f \mathcal{L}_{\text{fft}}\)$ 其中 \(\mathcal{L}_{\text{rec}} = \|\hat{y} - y\|_1\) 为像素级 \(\ell_1\) 基础重建损失;\(\mathcal{L}_{\text{fft}} = \| |\mathcal{F}(\hat{y})| - |\mathcal{F}(y)| \|_1\) 为快速傅里叶变换幅值损失,用于约束高频纹理与频域一致性,权重设定为 \(\lambda_f = 0.10\)。优化器选用 Adam(\(\beta_1=0.9, \beta_2=0.999\)),初始学习率 \(5 \times 10^{-4}\),配合余弦退火策略;AIO 训练 120 个 epoch,CDD-11 训练 200 个 epoch,Batch Size 为 64。

实验关键数据

主实验

评估涵盖标准 AIO-3、AIO-5 单一/混合退化基准,以及 CDD-11 复合退化协议(包含单退化、双退化及三重退化交叠场景)。所有测试均在盲恢复设定下进行。

基准协议 模型规模 方法 参数量 平均 PSNR (dB) 平均 SSIM 核心增益说明
AIO-3 Light MoCE-IR-S 11M 32.57 0.916 基线对比
AIO-3 Light CEA-IR-S (本文) 10M 32.89 0.918 +0.32 dB / +0.002,去雾去雨显著占优
AIO-3 Heavy MoCE-IR 25M 32.73 0.917 重型基线
AIO-3 Heavy CEA-IR (本文) 22M 33.07 0.922 +0.34 dB / +0.005,降噪与天气恢复均衡提升
AIO-5 Light MoCE-IR-S 11M 30.08 0.913 基线对比
AIO-5 Light CEA-IR-S (本文) 10M 30.80 0.919 +0.72 dB / +0.006,去雨/微光大幅领先
AIO-5 Heavy MoCE-IR 25M 30.58 0.919 重型基线
AIO-5 Heavy CEA-IR (本文) 22M 31.13 0.922 +0.55 dB / +0.003,去雾达 31.96 dB

在极具挑战的 CDD-11 复合退化基准上,CEA 在各重数退化交叠下均取得显著突破:

方法 单退化 (Single) 双重退化 (Double) 三重退化 (Triple) 全局综合 (Overall) 参数量 / 运行耗时
AirNet 26.09 / .885 23.11 / .776 22.02 / .717 23.75 / .814 8.93M / 42.13 ms
PromptIR 28.88 / .920 24.46 / .834 23.54 / .750 25.90 / .850 35.59M / 40.86 ms
OneRestore 31.68 / .938 27.35 / .866 24.84 / .790 28.47 / .878 6.00M / -
MoCE-IR-S 32.54 / .941 27.73 / .869 25.40 / .790 29.05 / .881 11.47M / 41.63 ms
CEA-IR-S (本文) 33.36 / .943 28.51 / .875 25.90 / .798 29.79 / .886 9.71M / 29.92 ms

消融实验

在 CDD-11 复合退化基准上对因子生成源、路由规则、空间探测方式及注入目标进行了细致的析因分析(统一在轻量级骨干与 \(r=8\) 下测试):

实验组别 探究变量配置 Single PSNR Double PSNR Triple PSNR Overall PSNR/SSIM 关键机制推论
因子与路由析因 静态因子 + Top-2 Softmax 32.15 27.50 25.16 28.77 / .879 静态离散基底恢复复合退化能力严重不足
静态因子 + 稠密有符号点积 32.35 27.77 25.40 29.01 / .881 稠密装配略有增益但受限于静态参数
动态因子 + Top-2 Softmax 32.50 27.94 25.76 29.20 / .882 动态因子带来基础改善,离散门控限制表达
动态因子 + 稠密有符号点积 (CEA) 33.36 28.51 25.90 29.79 / .886 双重协同发挥最大效能,较基线提升 1.02 dB
空间探测方式 全局平均池化 (GAP) + MLP 32.53 27.87 25.53 29.14 / .882 全局压缩抹除局域退化证据(-0.65 dB)
交叉注意力查询探测 (CEA) 33.36 28.51 25.90 29.79 / .886 空间细粒度证据被有效保留并编码为低秩因子
RankNorm 影响 移除 RankNorm 归一化 - - - 29.59 / .885 尺度漂移导致亲和度抖动,下降 0.20 dB
保留 RankNorm (CEA) - - - 29.79 / .886 保障低秩分量数值良态与稳定装配
注入位置 纯主干(无注入) 32.35 27.74 25.42 29.00 / .882 基线特征表达瓶颈
仅注入 Query (\(Q\)) 32.88 28.12 25.62 29.39 / .884 单侧调制改善注意力交互
同时注入 \(Q + K\) (CEA) 33.36 28.51 25.90 29.79 / .886 协同重塑相关性矩阵,取得最佳增益 (+0.79 dB)
注入 \(Q + K + V\) 32.94 28.18 25.68 29.46 / .884 过度参数化引入扰动,性能反而回落

关键发现

  • 退化越复合,CEA 的优势越明显:在单一退化中 CEA-IR-S 领先 MoCE-IR-S 0.82 dB,而在双重与三重退化中分别领先 0.78 dB 与 0.50 dB,在低光+雨/雾复合场景下峰值提升超过 1.2 dB,充分证实空间局域连续适配对于多退化叠加至关重要。
  • 算力与运行效率全面占优:在 RTX 4090 相同基准下,CEA-IR-S 参数量仅为 9.71M,FLOPs 仅 28.03G,推理延迟压至 29.92 ms,显著优于 MoCE-IR-S (41.63 ms) 和 PromptIR (40.86 ms);其根本原因在于两次低秩矩阵乘法直接规避了分支分支跳转或大矩阵动态卷积的访存瓶颈。
  • 低秩瓶颈饱和于 \(r=8\):实验显示秩维度从 4 升至 8 带来 +0.39 dB 增益,而继续扩大到 16 性能微降至 29.70 dB,表明极低维度的残差子空间足以捕捉退化调控流形,且能有效防止过拟合。

亮点与洞察

  • 从离散选择到连续装配的范式跃迁:传统 MoE 关注于“在多个专家网络中硬挑哪几个”,而 CEA 转向“针对输入生成连续流形基底并在 token 级做自由组合”,消除了路由坍塌和负载均衡损失等繁琐工程设计。
  • 巧借线性注意力实现低秩路由:将空间 token、路由基与残差方向巧妙映射为线性注意力的 Query、Key 与 Value,利用结合律将计算复杂度从序列平方级降至线性,数学形式优雅且硬件执行效率极高。
  • 通用的参数高效自适应机制:CEA 的超适配器与 RankNorm 设计不依赖特定主干,可直接迁移至底层视觉的其他基础模型(如 MambaIR、Restormer 或扩散模型),作为即插即用的动态条件调制单元。

局限与展望

  • 主干与任务覆盖边界:本文主要基于非对称 U-Transformer 架构并在中等分辨率上验证,尚未探索在超大规模扩散逆问题骨干或极高分辨率(如 4K)实时视频恢复上的动态适配开销。
  • 注入位置的理论最优性:目前消融确立了 \(Q+K\) 的经验最优,但对于 FFN 中跨通道动态混合与自注意力空间动态混合的最优容量分配比例,尚缺理论层面的解析度量。
  • 未来方向:探索将 CEA 的动态因子生成扩展为多尺度层级装配,或结合零样本开放世界大模型引导更复杂的野外退化解耦。

相关工作与启发

  • vs MoCE-IR [50] & AdaptIR [16]: 后者依赖预定义的静态异构专家网络库与离散路由策略,容易在复合退化下产生空间均质更新或路由抖动;CEA 彻底摒弃静态专家库,改用输入动态生成的低秩残差子空间与稠密点积连续装配。
  • vs PromptIR [37] & OneRestore [18]: PromptIR 与 OneRestore 主要依赖全局提示向量或解耦退化嵌入调制主干,面对局部强弱不等的退化易产生全局表征瓶颈;CEA 依托空间查询探测保留了局域退化证据,实现了逐 token 的精细化自适应。
  • vs LoRA-IR [1] & HAIR [2]: LoRA-IR 采用静态低秩矩阵的分支组合,HAIR 利用全局超网络生成权重矩阵;CEA 创新性地将超网络拆解为“空间探测生成基底 + 序列线性注意力装配”,兼具极高参数效率与局部动态表达力。

评分

  • 新颖性: ⭐⭐⭐⭐⭐ 巧妙融合线性注意力与动态低秩自适应,彻底重塑了全能恢复的参数条件化范式。
  • 实验充分度: ⭐⭐⭐⭐⭐ 涵盖 AIO-3、AIO-5、CDD-11 及零样本水下泛化,包含成对 Bootstrap 统计检验与严谨析因消融。
  • 写作质量: ⭐⭐⭐⭐⭐ 形式化推导清晰严密,代数解释直观透彻,图表组织层次分明。
  • 价值: ⭐⭐⭐⭐⭐ 为底层视觉中的条件计算、动态神经架构与参数高效微调提供了兼具理论自洽性与极高落地效率的经典方案。