UniH3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration¶
会议: ECCV 2026
论文: ECCV 原文
代码: https://github.com/Yaziwel/UniH3
领域: 医学图像
关键词: 全能医学图像恢复, 跨模态同质性, 分层记忆先验, 同质性引导注意力, 分层异质性平衡器
一句话总结¶
UniH3 首次在全能医学图像恢复中系统性联合建模多任务间的分层同质性与分层异质性，通过分层同质性记忆（H2M）与同质性引导注意力（HGA）蒸馏检索解剖共享先验，并利用分层异质性平衡器（H2B）消解任务间与样本级梯度冲突，在 MedIR-2D-500K 和 MedIR-3D-3K 两个大规模基准上均刷新 SOTA。
研究背景与动机¶
医学图像恢复（MedIR）长期依赖单任务专用模型，分别针对 PET、CT 去噪或 MRI 超分辨率等特定物理退化进行独立优化。随着多模态成像（如 PET/CT、PET/MRI）临床普及，单任务模型在面对复合退化场景时泛化能力骤降，且多模型冗余部署带来了沉重维护开销。近年来计算机视觉中兴起全能恢复（All-in-One）范式，医疗领域也涌现出如 AMIR 等基于混合专家架构（MoE）的探索。然而现有全能方法几乎完全将精力倾斜于任务异质性（Heterogeneity）的刻画，依靠对比学习、退化分类标签或视觉提示（Prompt）强行隔离不同任务的特征空间。
这种单向聚焦异质性的思路在医学成像中遭遇了严重瓶颈。首先，现有方法严重忽视了医学图像内在的固有同质性（Homogeneity）。与分布开放、内容杂乱的自然图像相比，不同模态乃至不同退化下的医学影像始终承载着人体高度一致的解剖拓扑结构与强生理先验。割裂这种跨任务/跨模态的同质关联，使得多任务学习演变为相互孤立的参数竞争，随着任务与模态数扩展，模型优化难度急剧攀升。其次，在异质性建模上，现有方法仅停留在粗粒度的任务级别（如区分 PET 还是 MRI），却完全忽略了由采集设备、扫描中心、采集协议及患者个体差异带来的任务内（Intra-task）精细异质性，导致优化时难以化解样本层面的隐式梯度冲突。
针对上述两大约束，本文提出必须在多任务体系中自顶向下同时兼顾同质性协同与异质性解耦。核心 idea：构建分层同质性与异质性协同框架 UniH3，在表征侧利用分层同质性记忆（H2M）与偏置注意力（HGA）从高质量数据中蒸馏并检索跨模态与模态内的解剖结构先验作为重构锚点；在优化侧利用分层异质性平衡器（H2B）联合动态解耦任务级与样本级不确定性，实现全能医学图像恢复的均衡收敛。
方法详解¶
整体框架¶
UniH3 采用基于 U 型非对称编码器-解码器的图像恢复主干网络，整体流程如下图所示。针对输入的低质医学图像 \(I_{LQ} \in \mathbb{R}^{H \times W \times 1}\)，网络首先经由 \(3 \times 3\) 卷积投影提取浅层特征 \(I_S \in \mathbb{R}^{H \times W \times C}\)（基准通道 \(C=48\)）。特征经由 4 级 U 型编码器-解码器逐级抽取深层特征 \(I_D\)，各阶段分别堆叠 \(N_1=2, N_2=3, N_3=3, N_4=4\) 个同质性引导 Transformer 块（HGATB）。在 U 型结构的各尺度跳跃连接处，挂载分层同质性记忆模块（H2M），自适应检索多分辨率解剖先验并注入 HGATB。最后深层特征经由输出卷积预测残差 \(I_R\)，通过全局残差连接输出重建的高质图像 \(\hat{I}_{HQ} = I_{LQ} + I_R\)。在训练阶段，优化目标由分层异质性平衡器（H2B）动态调控。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
Input["输入退化医学图像 ILQ"] --> Stem["浅层特征投影<br/>3x3 Conv 映射至 C 维空间"]
Stem --> H2M["分层同质性记忆模块 H2M<br/>自适应检索跨任务与任务内解剖先验 VH"]
H2M --> HGA["同质性引导注意力 HGA<br/>以高质量先验为基准的偏置重构"]
HGA --> HGATB["局部交互增强<br/>Conv + SE 块补充精细局部纹理"]
HGATB --> Recon["残差重构投影<br/>预测残差 IR 并叠加得到恢复图像"]
Recon --> H2B["分层异质性平衡器 H2B<br/>联合任务级与样本级不确定性动态调控"]
H2B --> Output["输出高质量医学图像 IHQ-hat"]关键设计¶
1. 分层同质性记忆模块：解耦蒸馏与自适应检索解剖先验
全能医学恢复模型随着支持模态和任务增加，学习负荷急剧加重。为激活不同任务间潜在的知识迁移，H2M 模块显式构建了两套结构对称的组件：参数化记忆库 \(M \in \mathbb{R}^{(T+1)L \times C'}\) 与可学习原型矩阵 \(P \in \mathbb{R}^{(T+1)L \times C'}\)。其中槽位被严格划分为 \(T\) 个任务特定槽位（存储任务内同质先验）与 1 个任务共享槽位（存储跨任务跨模态同质先验），每个槽位长度为 \(L=128\)。模块分为训练期的“同质性蒸馏”和全周期的“同质性检索”两个阶段：
- 同质性蒸馏（仅训练期）：将配对图像 \(I_{LQ}\) 与 \(I_{HQ}\) 经由 pixel-unshuffle 下采样与 \(3 \times 3\) 卷积映射为特征 \(F^{LQ}, F^{HQ} \in \mathbb{R}^{H'W' \times C'}\)。利用原型矩阵 \(P\) 作为查询向量，通过交叉注意力从真实高质量特征 \(F^{HQ}\) 中萃取代表性解剖基元特征 \(V^{HQ} = \text{Softmax}(P (F^{LQ})^\mathsf{T} / \sqrt{C'}) F^{HQ}\)。依据当前输入的任务标签，抽取对应的共享特征 \(V^{Sh}\) 与特定任务特征 \(V^{Sp}\)，并通过动量更新（EMA，动量系数 \(\alpha=0.99\)）写入记忆库： $\(M_{Sh/Sp} \leftarrow \alpha M_{Sh/Sp} + (1 - \alpha) V^{Sh/Sp}\)$ 测试期直接关闭蒸馏分支，避免任何对真实高质量图像的依赖。
- 同质性检索（训练与推理期）：以退化特征 \(F^{LQ}\) 作为查询，原型 \(P\) 作为键，记忆库 \(M\) 作为值，通过交叉注意力自适应检索出匹配当前解剖位置的无损结构先验 \(V^H\)： $\(V^H = \text{CrossAttention}(F^{LQ}, P, M)\)$ 通过在 U 型网络的 4 个分辨率尺度挂载独立的 H2M 模块，网络能够在从宏观解剖轮廓到细粒度组织纹理的多尺度上持续获得先验补偿。
2. 同质性引导注意力：以高质量先验为锚点的偏置重构机制
传统引导机制（如 SFT 或常规交叉注意力）通常以严重退化的低质特征为学习基准，引导信号仅作为辅助外挂，模型仍需艰难地从受污染的低质表征中拟合高质量分布。HGA 颠覆了这一范式，将学习锚点直接建立在 H2M 检索出的无损高质量同质先验 \(V^H\) 上。
在计算自注意力映射矩阵 \(A = \text{Softmax}(Q K^\mathsf{T} / \sqrt{C'})\) 后，若仅采用对称相加融合 \(A(V + V^H)\)，模型对受损特征 \(V\) 与纯净先验 \(V^H\) 一视同仁，无法充分发挥高质量先验的主导作用。为此，HGA 在注意力图上引入恒等偏置项，构造不对称倾向性 \((A - I)V + (A + I)V^H = A(V + V^H) + (V^H - V)\)，其中 \(\pm I\) 项抑制了低质特征自身的直接映射，强制强化了先验特征的自贡献，而残差项 \((V^H - V)\) 则形成强烈的向优质先验倾斜的归纳偏置。为保证训练稳定性与自适应表达，最终的 HGA 引入通道级可学习缩放参数 \(\lambda_1, \lambda_2 \in \mathbb{R}^{C'}\)： $\(V^O = A[(1 - \lambda_1)V + \lambda_1 V^H] + \lambda_2 (V^H - V)\)$ 当 \(\lambda_1 = \lambda_2 = 0\) 时，HGA 平滑退化为标准注意力。在实际工程实现中，UniH3 遵循 Restormer 采用沿通道维度的转置自注意力（Transposed Self-Attention），在维持线性复杂度的同时实现高保真度引导。
3. 分层异质性平衡器：任务级与样本级双层不确定性协同优化
不同医学任务退化机理差异显著，多任务联合训练极易发生梯度冲突，导致部分易学任务主导梯度方向，而困难任务欠拟合。传统基于同方差不确定性的损失加权（如 Kendall 等人方法）仅为每个任务学习一个全局任务标量 \(\sigma_t\)。然而该方案忽略了同一任务内部因设备厂家差异、病灶形态特异性所带来的任务内异质性。
H2B 建立了双层分层不确定性模型，将第 \(t\) 个任务中第 \(s\) 个样本的综合不确定性参数分解为任务级全局项与样本级精细校正项： $\(\sigma_{t,s} = \sigma_t + \Delta \sigma_s\)$ 其中 \(\sigma_t\) 为任务级可学习参数，负责应对任务间的大尺度分布差异；\(\Delta \sigma_s\) 则由一个轻量的不确定性估计块（UEB）实时预测： $\(\Delta \sigma_s = \text{UEB}\left( \text{Concat}[I_s^{LQ}, \text{sg}(\hat{I}_s^{HQ}), I_s^{HQ}] \right)\)$ 其中 \(\text{sg}(\cdot)\) 表示阻断梯度反向传播（stop-gradient），严格解耦不确定性估计与恢复主干的梯度流，防止平衡机制投机取巧。最终的 H2B 损失函数在整个批次（大小为 \(S\)）和多任务间联合求和： $\(\mathcal{L}_{H2B} = \frac{1}{TS}\sum_{t=1}^T \sum_{s=1}^S \left( \frac{1}{2\sigma_{t,s}^2} \mathcal{L}_{rec}^{(t,s)} + \log \sigma_{t,s} \right)\)$ 其中重构损失 \(\mathcal{L}_{rec}\) 采用基础 \(L_1\) 损失。当某一样本由于极端伪影或罕见病变导致重构误差畸高时，\(\sigma_{t,s}\) 自适应增大以平抑异常梯度冲击；而在平稳区域则降低 \(\sigma_{t,s}\) 放大监督信号，保证多任务全样本的稳健收敛。
损失函数 / 训练策略¶
模型采用端到端联合优化，训练损失完全由分层异质性平衡损失 \(\mathcal{L}_{H2B}\) 主导。在 MedIR-2D-500K 上，批大小设置为 14，图像切块大小为 \(128 \times 128\)。优化器采用前沿的 Muon 优化器，总计训练 \(6 \times 10^5\) 次迭代，初始学习率设为 \(3 \times 10^{-4}\)，并采用余弦退火策略逐步衰减至 \(1 \times 10^{-7}\)。针对 3D 医疗体数据，UniH3-3D 将全部卷积和注意力算子替换为三维算子，HGATB 阶段配置调整为 \(N_1=N_2=1, N_3=N_4=5\)，基准通道缩减至 \(C=16\)，体块尺寸为 \(64 \times 64 \times 64\)，批大小为 6，初始学习率调整为 \(5 \times 10^{-5}\)。
实验关键数据¶
主实验¶
实验在构建的大规模 2D 基准 MedIR-2D-500K（覆盖 7 项任务，509,200 个图像对）与 3D 基准 MedIR-3D-3K（覆盖 3 项任务，3,522 个体数据对）上全面展开。下表汇报了在 2D 全能医学恢复设置下的综合性能指标与计算开销：
| 方法 | 参数量 (M) | 计算量 (G) | PET (PSNR/SSIM) | CT (PSNR/SSIM) | MRI (PSNR/SSIM) | X-ray (PSNR/SSIM) | OCT (PSNR/SSIM) | 超声 (PSNR/SSIM) | 病理 (PSNR/SSIM) | 平均 (PSNR/SSIM) |
|---|---|---|---|---|---|---|---|---|---|---|
| SwinIR | 11.50 | 187.93 | 44.24 / 0.9866 | 43.17 / 0.9338 | 38.44 / 0.9453 | 35.87 / 0.9275 | 35.70 / 0.8892 | 27.52 / 0.8089 | 28.35 / 0.7941 | 36.18 / 0.8979 |
| Uformer | 50.47 | 21.42 | 44.51 / 0.9874 | 43.45 / 0.9356 | 39.01 / 0.9507 | 36.59 / 0.9353 | 35.84 / 0.8908 | 27.61 / 0.8135 | 28.55 / 0.8011 | 36.51 / 0.9020 |
| Restormer | 26.12 | 35.21 | 44.47 / 0.9873 | 43.44 / 0.9353 | 39.05 / 0.9515 | 36.46 / 0.9335 | 35.84 / 0.8909 | 27.66 / 0.8138 | 28.53 / 0.8006 | 36.49 / 0.9018 |
| NAFNet | 67.89 | 15.74 | 44.40 / 0.9871 | 43.32 / 0.9346 | 38.90 / 0.9502 | 36.33 / 0.9326 | 35.82 / 0.8905 | 27.59 / 0.8131 | 28.49 / 0.7995 | 36.41 / 0.9011 |
| AMIR | 23.54 | 31.76 | 44.49 / 0.9873 | 43.47 / 0.9356 | 39.09 / 0.9519 | 36.47 / 0.9333 | 35.89 / 0.8914 | 27.69 / 0.8150 | 28.57 / 0.8019 | 36.52 / 0.9023 |
| PromptIR | 35.59 | 39.49 | 44.52 / 0.9874 | 43.48 / 0.9355 | 39.13 / 0.9524 | 36.57 / 0.9341 | 35.84 / 0.8909 | 27.69 / 0.8152 | 28.54 / 0.8006 | 36.54 / 0.9023 |
| AdaIR | 28.76 | 36.74 | 44.55 / 0.9875 | 43.49 / 0.9356 | 39.17 / 0.9527 | 36.60 / 0.9344 | 35.86 / 0.8915 | 27.69 / 0.8150 | 28.56 / 0.8015 | 36.56 / 0.9026 |
| UniH3 (本文) | 28.96 | 26.33 | 44.89 / 0.9883 | 43.65 / 0.9368 | 39.55 / 0.9564 | 36.88 / 0.9368 | 35.96 / 0.8921 | 27.80 / 0.8179 | 28.63 / 0.8035 | 36.77 / 0.9045 |
同时，在 3D 全能医学恢复（MedIR-3D-3K）任务中，UniH3-3D 取得了 44.50 dB 平均 PSNR 与 0.9674 平均 SSIM，全面碾压前序最优 3D 方案 Restore-RWKV-3D（43.95 dB / 0.9644），平均增益高达 +0.55 dB。
消融实验¶
在 MedIR-2D-500K 基准上针对核心组件、HGA 引导形式以及 H2M/H2B 的分层有效性进行了系统性消融：
| 实验组别 | 配置说明 | 参数量 (M) | 计算量 (G) | 平均 PSNR (dB) | 平均 SSIM | 表现与说明 |
|---|---|---|---|---|---|---|
| 核心组件分析 | Baseline (去除 H2M 与 H2B) | 27.17 | 25.50 | 36.52 | 0.9023 | 退化为基础转置注意力与标准 L1 损失 |
| + H2M (仅引入同质性记忆与 HGA) | 28.96 | 26.33 | 36.66 | 0.9034 | 显著补充结构先验，增益 +0.14 dB | |
| + H2B (仅引入分层异质性平衡器) | 27.17 | 25.50 | 36.64 | 0.9033 | 有效解耦梯度冲突，增益 +0.12 dB | |
| UniH3 完整模型 (+ H2M + H2B) | 28.96 | 26.33 | 36.77 | 0.9045 | 两者协同达到最优，累计增益 +0.25 dB | |
| HGA 机制消融 | w/o HGA (标准自注意力) | 27.17 | 25.50 | 36.52 | 0.9023 | 缺乏先验引导基准 |
| SFT 特征变换机制 [55] | 37.46 | 36.99 | 36.74 | 0.9041 | 参数与计算量骤增，性能仍不及 HGA | |
| 经典交叉注意力引导 [11] | 29.60 | 27.43 | 36.67 | 0.9035 | 对称交互未建立先验优先偏置 | |
| HGA (本文偏置注意力) | 28.96 | 26.33 | 36.77 | 0.9045 | 轻量高效且引导精准 | |
| H2M 分层同质性 | 仅保留跨任务同质性 (Inter-task) | 28.96 | 26.33 | 36.61 | 0.9031 | 仅捕获全局模态共享解剖 |
| 仅保留任务内同质性 (Intra-task) | 28.96 | 26.33 | 36.72 | 0.9038 | 捕获单模态多患者一致性 | |
| 双层同质性联合 (Inter + Intra) | 28.96 | 26.33 | 36.77 | 0.9045 | 全面激活跨模态与模态内同质先验 | |
| H2B 分层异质性 | 基础同方差平衡 (仅 Inter-task) | 28.96 | 26.33 | 36.70 | 0.9036 | 仅能粗粒度平衡模态间损失 |
| 仅样本级不确定性 (仅 Intra-task) | 28.96 | 26.33 | 36.73 | 0.9041 | 缺少宏观任务尺度约束 | |
| 双层异质性平衡 (H2B 完整) | 28.96 | 26.33 | 36.77 | 0.9045 | 联合任务与样本级解耦效果最优 |
关键发现¶
- 同质性与异质性协同是突破全能瓶颈的钥匙：从消融数据可以看出，单独建模同质先验（H2M 提升 0.14 dB）或单独调控异质梯度（H2B 提升 0.12 dB）均能带来稳健增益，而两者结合时能产生显著的叠加效应（达到 36.77 dB），证明共享先验的学习与无冲突的多任务优化相互促进。
- H2M 内部呈现高度符合解剖学规律的检索模式：可视化检索注意力分布发现，解剖位置相似的 token（例如两枚 PET 脊柱 token）其 Top-10 检索槽位重合度高达 8/10；跨模态相似解剖（PET 脊柱与 CT 脊柱）也具有 2/10 的共享重合度；而同一模态内不同病灶解剖（PET 脊柱与 PET 肿瘤）的重合度为 0/10，证明 H2M 确实在隐空间精确解耦并存储了可迁移的解剖拓扑基元。
- 全能模型具备超越单任务模型的泛化潜能：UniH3 单一全能模型平均得分 36.77 dB，不仅以大幅优势碾压过往全能模型，甚至直接战胜了为每个任务单独训练的专用单任务模型（如单任务 MambaIR 平均 36.75 dB）；将全能 UniH3 作为预训练底座进一步微调时，平均 PSNR 可跃升至 36.97 dB，确立了全能预训练在医学领域的巨大迁移价值。
亮点与洞察¶
- 颠覆先验注入逻辑的 HGA 偏置设计：传统恢复模型将先验视作从属特征（辅助输入），而 HGA 通过在自注意力中植入代数恒等偏置 \((A-I)V + (A+I)V^H\)，从数学上直接迫使模型将干净解剖先验作为重构锚点，这一巧思以极低计算代价化解了深层特征严重退化导致的漂移难题。
- 解剖同质先验的双层解耦记忆机制：首次将医学影像中“跨模态共有的人体解剖结构（Inter-task）”与“同模态内部不同病人的组织拓扑（Intra-task）”拆解为共享槽位与专用槽位，在训练期利用动量蒸馏提炼、在测试期无损检索，彻底打破了以往医疗全能模型“任务越多、干扰越大”的魔咒。
- 可平替的即插即用泛化能力：论文实验证明，将 H2M 与 H2B 模块直接移植到 Uformer、Restormer、PromptIR 和 AdaIR 等主流恢复架构上，所有模型均无一例外获得 0.15~0.20 dB 的显著性能跃升，展现出卓越的模块通用性。
局限与展望¶
- 退化类型的覆盖广度仍有局限：当前基准与实验在每种模态下仅聚焦于最主要的一种物理退化（如 PET 仅做低剂量去噪、MRI 仅做超分辨率），尚未建模单模态内复合退化（如 MRI 同时伴随运动伪影与欠采样噪声）及非配对复杂临床场景。
- 样本级不确定性估计的额外计算前传：H2B 在训练期需要引入轻量级 UEB 网络并输入预测结果 \(\hat{I}^{HQ}\)，虽然测试期完全剥离该分支，但在极大规模训练集上的整体训练显存占用有所增加。
- 未来演进方向：未来可探索将大语言模型（LLM）的解剖语义先验与 H2M 记忆矩阵进行多模态对齐，并在下游精准分割、病灶检测等临床诊断任务中验证恢复后图像的诊断增强价值。
相关工作与启发¶
- vs AdaIR / PromptIR (视觉提示全能恢复)：PromptIR 与 AdaIR 依赖输入特征自适应生成连续 Prompt 向量，主要用于引导解码器区分任务类型（建模异质性），对跨模态共有解剖结构缺乏显式存储与复用机制。UniH3 创新性地提出分层同质性记忆，主动将解剖同质先验与退化异质性解耦，检索效率更高且结构保真度更强。
- vs AMIR (医学多任务 MoE 架构)：AMIR 采用稀疏门控混合专家网络（MoE）根据不同任务动态分发特征流，计算开销较大且各专家分支易发生局部过拟合。UniH3 保持主干网络参数紧凑统一（仅 28.96M 参数、26.33G FLOPs），通过 H2B 动态损失平衡从根源上缓解梯度冲突，在计算开销更低的同时大幅提升多任务恢复质量。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 首次系统性提出在全能医学恢复中联合建模分层同质性与分层异质性，H2M 记忆槽与 HGA 偏置注意力设计极具洞见。
- 实验充分度: ⭐⭐⭐⭐⭐ 构建了 MedIR-2D-500K 与 MedIR-3D-3K 两个大规模多任务基准，涵盖 2D/3D、全能/单任务、丰富消融及跨主干验证。
- 写作质量: ⭐⭐⭐⭐⭐ 概念清晰，公式推导严谨自洽，动机与机制阐述环环相扣，图表详实规范。
- 价值: ⭐⭐⭐⭐⭐ 为通用医学图像大模型与通用预训练底座提供了开创性范式，代码开源且模块具备强大的即插即用迁移潜力。