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 stores intra- and inter-task anatomical priors in a memory bank and balances sample uncertainty. Seven 2D MedIR tasks reach 36.77 dB, 0.21 dB above AdaIR.
Medical image restoration is still mostly a per-modality craft: denoise PET, denoise CT, super-resolve MRI, each with its own network. A clinic that runs several degradations at once then has to swap specialists, and the field stalls short of a general restorer.
All-in-One methods such as AMIR, PromptIR, and AdaIR mainly model how tasks differ. Contrastive prompts, degradation classifiers, and mixture-of-experts all try to tell tasks apart. This paper flags two missing pieces. Medical images are more homogeneous than natural photos: anatomy repeats across patients and even across CT and PET, which is unused shared knowledge. Tasks are also heterogeneous inside, across scanners, sites, and populations, so a single scalar weight per task cannot rebalance hard cases.
UniH3 is a U-shaped restorer plus a Hierarchical Homogeneity Memory (H2M) and a Hierarchical Heterogeneity Balancer (H2B). A 3×3 stem feeds a four-level encoder–decoder. Each Homogeneity-Guided Transformer Block pairs a global HGA layer with a local conv+SE layer. The network predicts a residual that is added back to the low-quality input.
H2M keeps a memory M and a learnable prototype P, both laid out as T task slots plus one shared slot. During training, the prototype attends into high-quality features with low-quality features as keys, then writes into the shared slot (cross-task anatomy) and the current task slot (within-modality anatomy) by EMA with momentum 0.99. At test time the distillation path is dropped; low-quality features retrieve the closest clean prior from M. Four H2M modules sit at four scales.
Retrieved priors are not injected by SFT or vanilla cross-attention, both of which treat the degraded feature as the base. Homogeneity-Guided Attention anchors on the clean prior: mixed values go through attention, plus a (V^H − V) term that pushes weight onto the memory. λ1 = λ2 = 0 recovers ordinary self-attention. The implementation uses Restormer-style transposed attention.
H2B splits Kendall-style uncertainty. Total uncertainty is σ{t,s} = σt + Δσs. A small UEB predicts the sample term from the low-quality image, the stop-gradient prediction, and the high-quality target, so hard samples and conflicting tasks are down-weighted together.
The 2D model uses C=48, block counts 2/3/3/4, and Muon for 600k steps. UniH3-3D swaps in 3D counterparts, C=16, 64³ patches.
Two assembled benchmarks: MedIR-2D-500K with 509,200 2D pairs across seven tasks, and MedIR-3D-3K with 3,522 volumes across three tasks. About 51,200 2D images are held out. Metrics are PSNR and SSIM.
| Setting | UniH3 | Runner-up | Gap |
| 2D all-in-one mean PSNR | 36.77 | AdaIR 36.56 | +0.21 dB |
| 2D single-task mean PSNR | 36.90 | MambaIR 36.75 | +0.15 dB |
| 3D all-in-one mean PSNR | 44.50 | Restore-RWKV-3D 43.95 | +0.55 dB |
| 3D single-task mean PSNR | 45.29 | Spach Transformer 43.81 | +1.48 dB |
The 2D all-in-one model wins all seven tasks. PET 44.89, CT 43.65, MRI 39.55, versus AdaIR at 44.55 / 43.49 / 39.17. Ultrasound and pathology remain the hard pair at 27.80 and 28.63. Parameter count is 28.96M, close to AdaIR's 28.76M, with 26.33G FLOPs against 36.74G. One all-in-one UniH3 (36.77) already matches specialist MambaIR (36.75).
Ablations on 2D all-in-one: both modules off 36.52; H2M only 36.66; H2B only 36.64; both 36.77. HGA scores 36.77 versus 36.74 for SFT at 37.46M / 36.99G, and 36.67 for cross-attention. Inside H2M, the intra-task slot (36.72) helps more than the shared slot (36.61). Retrieval maps agree: PET spine tokens query the shared slot and the PET slot, overlap 2/10 top keys with CT spine, and 0/10 with a lesion token.
This is a large public-scale number for "one network for a radiology stack." 0.21 dB is not a visual leap; the authors defend it with thousands of test images per task. The more useful fact is that the all-in-one model already matches a specialist, and task-wise finetuning still adds a bit. H2M and H2B also lift Uformer, Restormer, PromptIR, and AdaIR, so the modules are not locked to this backbone.
Clinical use still depends on the modality. Ultrasound and pathology PSNR stay low, and each modality is reduced to its primary degradation.
The authors only cover the main restoration task per modality. Other degradations and downstream diagnosis are untested. Training mixes public and private data, so exact reproduction depends on the private split. The 0.21 dB edge over AdaIR and the 0.03 dB edge over SFT have no significance test. H2M distillation needs paired high-quality images even though retrieval at test time does not. The 3D CT test split has 30 volumes, so the 1.48 dB single-task gap there is a small-N number.