UniH3：统一层级同质性与异质性以实现全合一医学图像复原

快速导读：本文提出 UniH3 框架，通过层级同质性记忆（H2M）和层级异质性平衡器（H2B）同时建模医学图像中的层级同质性与异质性，在 2D 和 3D 基准上实现了全合一及单任务医学图像复原的 SOTA 性能。

医学图像复原长期面临多模态、多退化类型下的效率与泛化瓶颈。传统单任务模型难以应对临床场景的多样性，而现有 All-in-One 方法虽尝试统一处理，却往往忽略医学图像固有的层级同质性，导致多任务学习中的优化冲突。UniH3 提出了一种新的视角，将同质性与异质性统一建模，旨在提升全合一复原框架的鲁棒性与性能。

本文的核心贡献在于提出 UniH3 框架，包含层级同质性记忆（H2M）、同质性引导注意力（HGA）和层级异质性平衡器（H2B）。通过在训练阶段蒸馏高质量先验，并在推理阶段动态检索与平衡，UniH3 在 MedIR-2D-500K 和 MedIR-3D-3K 数据集上均取得了显著的性能提升，验证了层级建模策略的有效性。

这篇论文解决什么问题？

医学图像复原任务涵盖去噪、超分辨率、去模糊等多种退化类型，且涉及 CT、MRI、超声等不同模态。传统方法通常针对单一任务和模态进行优化，导致模型数量庞大且难以维护。All-in-One 框架旨在用单一模型处理多种任务，但现有方法如 AMIR 和 AdaIR 主要关注任务间的异质性，通过 MoE 或 Prompt 机制区分任务，却忽略了跨模态共享的解剖结构等层级同质性。

这种忽略导致模型在处理不同任务时缺乏共享先验的引导，容易陷入局部最优。此外，任务内的异质性，如不同扫描仪或中心导致的分布差异，也未被充分处理。Kendall 等人的不确定性加权方法仅处理任务间差异，无法应对样本级的分布变化。因此，如何同时建模层级同质性与异质性，成为提升 All-in-One 医学图像复原性能的关键挑战。

核心创新

• 提出 H2M 模块，显式建模并蒸馏层级同质性先验（Inter-task 和 Intra-task），通过记忆库存储和检索机制指导复原。
• 设计 HGA 机制，改变传统以 LQ 特征为基础的引导范式，转而以 HQ 同质性先验为锚点，通过可学习权重 $\lambda_1, \lambda_2$ 调整注意力偏好，降低学习难度。
• 开发 H2B 策略，引入层级不确定性加权损失，不仅平衡任务间异质性，还通过样本级不确定性估计缓解任务内异质性引起的优化冲突。
• 构建大规模基准数据集 MedIR-2D-500K 和 MedIR-3D-3K，涵盖多种模态和退化类型，促进通用 MedIR 研究。

方法概览

UniH3 包含三个核心组件：1. U 型复原骨干网络，集成 Homogeneity-Guided Transformer Blocks (HGATBs)；2. Hierarchical Homogeneity Memory (H2M)，在训练时通过 EMA 从高质量图像中蒸馏任务共享（Inter-task）和任务特定（Intra-task）的同质性先验存入记忆库，推理时通过 Cross-Attention 检索相关先验；3. Homogeneity-Guided Attention (HGA)，将检索到的同质性先验注入注意力机制，以 HQ 先验为锚点引导复原；4. Hierarchical Heterogeneity Balancer (H2B)，通过层级不确定性模型（任务级 $\sigma_t$ 和样本级 $\Delta\sigma_s$）动态平衡损失，缓解任务间和任务内的优化冲突。

• U 型复原骨干网络：集成 Homogeneity-Guided Transformer Blocks (HGATBs)，作为特征提取与复原的基础架构，支持多尺度特征融合。
• 层级同质性记忆 (H2M)：在训练时通过 EMA 从高质量图像中蒸馏 Inter-task 和 Intra-task 同质性先验，存入记忆库；推理时通过 Cross-Attention 检索相关先验，为复原提供高质量锚点。
• 同质性引导注意力 (HGA)：改变传统以 LQ 特征为基础的引导范式，转而以 HQ 同质性先验为锚点，通过可学习权重调整注意力偏好，降低学习难度，提升复原质量。
• 层级异质性平衡器 (H2B)：引入层级不确定性模型，包括任务级 σ_t 和样本级 Δσ_s，动态平衡损失，缓解任务间和任务内的优化冲突，提升多任务学习的稳定性。
• 记忆库更新机制：仅在训练阶段进行，确保推理时的效率与稳定性，但需注意测试数据分布与训练差异对检索效果的影响。
• 通用性验证：将 H2M 和 H2B 模块应用于其他骨干网络（如 Uformer, Restormer），验证其在不同架构下的有效性与可迁移性。

逐图理解论文

失败案例

该表展示了 MedIR-2D-500K 数据集上 All-in-One 设置的对比结果。需关注 UniH3 与 AdaIR 等方法的 PSNR 差异，验证层级建模策略的有效性。

研究缺口

该图详细展示了 H2M 模块，包括记忆库的构建与检索过程。需关注 EMA 蒸馏机制与 Cross-Attention 检索，理解如何提取并应用层级同质性先验。

核心贡献

该图展示了 HGA 机制，包括 HQ 先验如何引导注意力。需关注可学习权重 λ1, λ2 的作用，理解如何降低学习难度并提升复原质量。

最强结论

该表展示了 MedIR-3D-3K 数据集上 3D All-in-One 设置的对比结果。需关注 UniH3 与 Restore-RWKV-3D 等方法的 PSNR 差异，验证其在 3D 任务中的优势。

意义与局限

该图展示了 H2M 的检索注意力图，Top-10 分数用红框标出。需关注检索到的先验是否与输入图像语义相关，验证 H2M 的有效性。

实验如何设计？

• 数据集：构建 MedIR-2D-500K 和 MedIR-3D-3K 数据集，涵盖多种模态和退化类型，促进通用 MedIR 研究。
• 对比方法：包括 SwinIR, Restormer, MambaIR, AMIR, AdaIR 等 SOTA 方法，在 All-in-One 和 Single-Task 设置下进行对比。
• 评估指标：使用 PSNR 和 SSIM 作为主要评估指标，全面衡量复原质量。
• 消融实验：验证 H2M, H2B, HGA 各组件的有效性，以及 Inter-task 和 Intra-task 同质性/异质性建模的贡献。
• 可视化分析：展示 H2M 的检索注意力图和 H2B 的不确定性分布，验证层级先验的语义相关性与平衡效果。

关键结果与论文证据

• 在 MedIR-2D-500K 数据集上，UniH3 在 All-in-One 设置下的平均 PSNR 达到 36.77 dB，优于 AdaIR 的 36.56 dB，提升 0.21 dB（第 10 页）。
• 在 MedIR-2D-500K 数据集上，UniH3 在 Single-Task 设置下的平均 PSNR 达到 36.90 dB，优于 MambaIR 的 36.75 dB，提升 0.15 dB（第 10 页）。
• 在 MedIR-3D-3K 数据集上，UniH3 在 All-in-One 设置下的平均 PSNR 达到 44.50 dB，优于 Restore-RWKV-3D 的 43.95 dB，提升 0.55 dB（第 11 页）。
• 在 MedIR-3D-3K 数据集上，UniH3 在 Single-Task 设置下的平均 PSNR 达到 45.29 dB，优于 Spach Transformer 的 43.81 dB，提升 1.48 dB（第 11 页）。
• 消融实验显示，H2M 和 H2B 组件的结合使 PSNR 从基线的 36.52 dB 提升至 36.77 dB，验证了层级建模策略的有效性（第 11 页）。
• 可视化分析表明，H2M 的检索注意力图能够准确捕捉语义相关的先验，H2B 的不确定性分布合理，有效缓解了优化冲突（第 13 页）。

阅读时需要注意

• 仅关注每种模态的主要复原任务，未覆盖同一模态下的其他退化类型或任务，限制了模型的通用性。
• All-in-One 模型的实际临床价值仍在讨论中，尽管实验显示其性能接近单任务 SOTA，但临床部署需进一步验证。
• 数据集包含部分私有数据，可能影响完全复现性，需公开更多数据以促进研究。
• H2M 的记忆库更新仅在训练阶段进行，若测试数据分布与训练差异极大，检索效果可能下降，需进一步研究动态更新机制。

关联工作

• AMIR：先驱工作，使用 MoE 策略处理三个医学复原任务，但主要关注任务间异质性，未充分利用层级同质性。
• AdaIR：SOTA All-in-One 方法，使用频率挖掘和调制，UniH3 在 2D All-in-One 中超越其性能，验证了层级建模的优势。
• MambaIR：基于 State-Space Model 的单任务复原 SOTA，UniH3 在单任务设置下超越其性能，展示了多任务框架的潜力。
• Restore-RWKV：基于 RWKV 的医学图像复原方法，在 3D 任务中作为主要对比基线，UniH3 在 3D 任务中表现更优。
• Kendall et al. (Uncertainty Weighting)：H2B 的理论基础，传统不确定性加权仅处理任务间异质性，H2B 扩展至样本级，提升了平衡效果。