[CVPR 2026] MiM-DiT：分层 MoE 赋能扩散 Transformer，重塑全能图像修复

TL;DR

传统的图像修复往往“顾此失彼”，难以用一个模型完美解决从去雾到去噪的所有问题。本文提出的 MiM-DiT 巧妙地将“专家混合机制 (MoE)”嵌套化（MoE in MoE），并嵌入到预训练的 Diffusion Transformer (DiT) 中。它不仅能根据退化类型自动选择最合适的神经网络架构（如空间注意力 vs 通道注意力），还能针对退化的严重程度进行精细调节。

背景定位

在图像修复（Image Restoration）领域，全能模型（All-in-one）是终极目标。然而，不同退化的数学本质迥异：

• 模糊 (Blur) 需要恢复缺失的高频局部结构。
• 低光 (Low-light) 需要全局的亮度重映射和颜色校正。
• 噪声 (Noise) 则需要识别并剔除随机的像素级干扰。

现有的扩散模型虽有强大的生成能力，但往往采用“一刀切”的策略。MiM-DiT 则是通过“动态架构演变”来应对这种异质性。

核心动机：为何需要“专家中的专家”？

作者观察到，简单的 MoE（同质专家）只能增加参数量，却无法提供不同的归纳偏置（Inductive Bias）。例如，如果所有专家都是空间自注意力机制，它们就很难处理全局通道间的关系。

为此，MiM-DiT 提出了 MiM (MoE in MoE) 概念：

1. Inter-MoE (层间专家)：解决“选什么武器”的问题——是侧重空间关系的 Spatial Attention，还是侧重光照规律的 SE Attention？
2. Intra-MoE (层内专家)：解决“用多少力气”的问题——针对同一种退化，选择处理轻微或严重程度的特定子模块。

方法论：双层 MoE 与 DiT 的深度融合

1. 架构解析

MiM-DiT 的核心逻辑是：将退化后的 LQ 图像通过 MiM 模块提取“修复先验”，再通过 ZeroLinear 层注入到 DiT 的主干网络中。

• Inter-MoE 密集路由：采用了四种异构注意力机制：Spatial Self-attention (长程依赖)、Channel Self-attention (特征校准)、Swin Attention (局部/全局平衡)、SE Attention (全局光照)。
• Intra-MoE 稀疏路由：在每个路径内部，采用 Top-k 策略激活子专家。这种设计既保证了灵活性，又避免了全激活带来的巨大计算开销。

2. 注入机制：Zero-Linear Pathway

为了保证预训练 DiT 权重的稳定性，作者引入了零初始化（Zero-initialized）的线性层。这使得模型在训练初期表现得像原始扩散模型，随着训练进行，逐渐学会利用 MiM 提取的修复信号。

实验战绩

MiM-DiT 在 FoundIR 数据集上的表现堪称惊艳。

• 感知质量飞跃：在去噪、去雾任务上，LPIPS 指标大幅优于对比方法（如 DA-CLIP, DiffUIR）。
• 消融验证：实验证明，如果去掉 Intra-MoE，感知指标（如 MUSIQ）会显著下降，图像边缘会出现过度平滑。
• 路由自适应性：分析发现，当处理模糊时，模型会优先分配权重给 Spatial 和 Swin 专家；而处理低光时，则偏向于 Channel 和 SE 专家。这证明了模型确实学到了不同物理退化的特性。

深度洞察与局限性

总结：MiM-DiT 的成功在于它意识到“架构多样性”在多任务学习中的重要性。将宏观的架构路径选择与微观的参数调制结合，为通用底层视觉模型提供了新范式。

局限性：

• 虽然 Intra-MoE 是稀疏的，但 Inter-MoE 的密集融合在推理时仍会带来一定的延迟。
• 对于“复合退化”（如同时模糊且有噪声）的专家分配逻辑仍有待进一步深挖。

未来展望： 随着万亿级参数图像修复大模型的出现，这种“分层架构专家”的设计可能会成为处理无限退化类型的标配。