UniH$^3$: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration
Zhiwen Yang, Jiayin Li, Chengyu Liu, Hui Zhang, Bingzheng Wei, Yan Xu
ECCV 2026
cs.CV
2026-09-10
北航提出UniH3，把解剖结构的层内/跨任务共性写进记忆库，再按样本调多任务权重。七项2D医学复原All-in-One平均PSNR 36.77，比AdaIR高0.21 dB。
医学图像复原长期按模态单独立项：PET 去噪、CT 去噪、MRI 超分，各训各的。多模态机房里要同时伺候几套退化，专模换来换去成本高，也挡着通用复原模型。
近年 All-in-One 路线（AMIR、PromptIR、AdaIR）主要在建模任务之间的差异：对比学习、退化分类、visual prompt、MoE，让网络先分清「这是哪一项」。北航、清华这条线指出两处漏掉的结构。医学图像比自然图像更同质，同一模态不同病人、甚至 CT 和 PET 之间，解剖结构是共享先验，可以拿来减学习负担。任务内部也有异质性：扫描仪、中心、人群不同，一个任务一个标量权重压不住难样本。
UniH3 是 U 形复原骨干，加上分层同质记忆 H2M 和分层异质平衡 H2B。输入先 3×3 卷积，四层不对称编解码，每层堆 Homogeneity-Guided Transformer Block：一块用 HGA 做全局交互，一块用卷积加 SE 做局部。输出残差加回低质量图。
H2M 维护记忆库 M 和可学习原型 P，都按 T 个任务槽加 1 个共享槽排。训练时用高质量图做蒸馏：原型以低质量特征为键、高质量特征为值做 cross-attention，再用 EMA（动量 0.99）写入共享槽（跨任务解剖）和当前任务槽（模态内解剖）。推理丢掉蒸馏支路，只用低质量特征去记忆库里取最相关的干净先验。四个尺度各放一个 H2M。
取回来的先验不走常见的 SFT 或普通 cross-attention。那两类都把低质量特征当底座、先验当添加剂。HGA 把学习起点锚在高质量先验上：注意力仍看混合 value，再加一项 (V^H − V) 把权重推向干净先验。λ1、λ2 为 0 时退回普通自注意力。实现跟 Restormer 一样用 transposed attention。
H2B 把 Kendall 式不确定性从「每任务一个 σ」拆开。总不确定度 σ{t,s} = σt + Δσs，Δσs 由轻量 UEB 看低质量图、预测（stop-gradient）和真值。难样本权重要降，任务之间也不让某一个拖死梯度。
2D 版 C=48，块数 2/3/3/4，Muon 优化 60 万步。3D 版把模块换成三维对应物，C=16，patch 64³。
两个自建基准：MedIR-2D-500K 含 509,200 对 2D 图、七项任务；MedIR-3D-3K 含 3,522 个 3D 体、三项任务。测试集约 51,200 张 2D。指标 PSNR / SSIM。
| 设置 | UniH3 | 次优 | 差距 |
| 2D All-in-One 平均 PSNR | 36.77 | AdaIR 36.56 | +0.21 dB |
| 2D 单任务平均 PSNR | 36.90 | MambaIR 36.75 | +0.15 dB |
| 3D All-in-One 平均 PSNR | 44.50 | Restore-RWKV-3D 43.95 | +0.55 dB |
| 3D 单任务平均 PSNR | 45.29 | Spach Transformer 43.81 | +1.48 dB |
2D All-in-One 七项全胜。PET 44.89 / CT 43.65 / MRI 39.55，对 AdaIR 的 44.55 / 43.49 / 39.17。超声和病理最难，27.80 和 28.63，仍各高一点。参数 28.96M，和 AdaIR 的 28.76M 接近，FLOPs 26.33G 低于 AdaIR 的 36.74G。一个 All-in-One UniH3（36.77）已经和单任务 SOTA MambaIR（36.75）持平。
消融（2D All-in-One）：两件都关 36.52；只开 H2M 36.66；只开 H2B 36.64；全开 36.77。HGA 36.77，SFT 36.74 但参数到 37.46M、FLOPs 36.99G，普通 cross-attention 36.67。H2M 里任务内槽（36.72）比跨任务槽（36.61）贡献更大，合在一起最好。检索注意力也支持这一点：PET 脊椎 token 主要查共享槽和 PET 槽，和 CT 脊椎还有 2/10 的 top 重叠，和病灶 token 是 0/10。
给「一个模型伺候一间影像科」多了一份可对照的大规模数字。0.21 dB 不是观感级跳变，作者自己也按「每项几千张测试图」来辩护这点增益。更有用的观察是：All-in-One 已经追平专模，再按任务微调还能再涨。H2M/H2B 接到 Uformer、Restormer、PromptIR、AdaIR 上同样涨点，模块不绑死在自家骨干。
临床落地仍取决于具体模态和退化。超声、病理的绝对 PSNR 仍低，论文也只覆盖每模态的主退化。
作者承认只做了各模态的主任务，同一模态里其他退化和下游诊断任务都没测。数据拼了公开集和私有集，复现边界取决于私有部分能否公开。2D 上对 AdaIR 的 0.21 dB、对 SFT 的 0.03 dB，都没有显著性检验，属于排行榜式领先。H2M 蒸馏要用高质量图，推理虽不需要，训练仍是配对 LQ–HQ 设定。3D 测试集 CT 只有 30 个体，那一项的 1.48 dB 单任务优势方差会偏大。