# UniH3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration

arXiv: https://arxiv.org/abs/2609.11156 · HTML: https://mlsanity.com/paper/2609.11156/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 46 of its 76 references that are in the map.
- First posted on arXiv 10 Sep 2026.

## Builds on (most cited first, 30 of 46)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Squeeze-and-Excitation Networks](https://mlsanity.com/paper/1709.01507/) (2017)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [SwinIR: Image Restoration Using Swin Transformer](https://mlsanity.com/paper/2108.10257/) (2021)
- [Restormer: Efficient Transformer for High-Resolution Image Restoration](https://mlsanity.com/paper/2111.09881/) (2021)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [Uformer: A General U-Shaped Transformer for Image Restoration](https://mlsanity.com/paper/2106.03106/) (2021)
- [RWKV: Reinventing RNNs for the Transformer Era](https://mlsanity.com/paper/2305.13048/) (2023)
- [Simple Baselines for Image Restoration](https://mlsanity.com/paper/2204.04676/) (2022)
- [Muon is Scalable for LLM Training](https://mlsanity.com/paper/2502.16982/) (2025)
- [Recovering Realistic Texture in Image Super-resolution by Deep Spatial Feature Transform](https://mlsanity.com/paper/1804.02815/) (2018)
- [MambaIR: A Simple Baseline for Image Restoration with State-Space Model](https://mlsanity.com/paper/2402.15648/) (2024)
- [Deep Learning for Segmentation using an Open Large-Scale Dataset in 2D Echocardiography](https://mlsanity.com/paper/1908.06948/) (2019)
- [TransWeather: Transformer-based Restoration of Images Degraded by Adverse Weather Conditions](https://mlsanity.com/paper/2111.14813/) (2021)
- [Gland Segmentation in Colon Histology Images: The GlaS Challenge Contest](https://mlsanity.com/paper/1603.00275/) (2016)
- [Can AI help in screening Viral and COVID-19 pneumonia?](https://mlsanity.com/paper/2003.13145/) (2020)
- [AdaIR: Adaptive All-in-One Image Restoration via Frequency Mining and Modulation](https://mlsanity.com/paper/2403.14614/) (2024)
- [Complexity Experts are Task-Discriminative Learners for Any Image Restoration](https://mlsanity.com/paper/2411.18466/) (2024)
- [Efficient and Accurate MRI Super-Resolution using a Generative Adversarial Network and 3D Multi-Level Densely Connected Network](https://mlsanity.com/paper/1803.01417/) (2018)
- [Swin Transformer for Fast MRI](https://mlsanity.com/paper/2201.03230/) (2022)
- [OCTA-500: A Retinal Dataset for Optical Coherence Tomography Angiography Study](https://mlsanity.com/paper/2012.07261/) (2020)
- [CTformer: Convolution-free Token2Token Dilated Vision Transformer for Low-dose CT Denoising](https://mlsanity.com/paper/2202.13517/) (2022)
- [MedFMC: A Real-world Dataset and Benchmark For Foundation Model Adaptation in Medical Image Classification](https://mlsanity.com/paper/2306.09579/) (2023)
- [Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV](https://mlsanity.com/paper/2407.11087/) (2024)
- [Universal Image Restoration Pre-training via Degradation Classification](https://mlsanity.com/paper/2501.15510/) (2025)
- [All-In-One Medical Image Restoration via Task-Adaptive Routing](https://mlsanity.com/paper/2405.19769/) (2024)
- [Enhanced generative adversarial network for 3D brain MRI super-resolution](https://mlsanity.com/paper/1907.04835/) (2019)
- [GraphEcho: Graph-Driven Unsupervised Domain Adaptation for Echocardiogram Video Segmentation](https://mlsanity.com/paper/2309.11145/) (2023)

## Abstract (arXiv)

All-in-One medical image restoration (MedIR) aims to address diverse tasks across modalities and degradation types using a single universal model. Existing methods typically prioritize modeling inter-task heterogeneity (e.g., distinct data distributions and degradation types). However, they largely neglect the inherent homogeneity present in medical images, such as widely shared anatomical structures within and across modalities, which can be leveraged to ease model training and improve generalization. To this end, we propose UniH3, a novel framework that Unifies Hierarchical Homogeneity and Heterogeneity for all-in-one medical image restoration. Specifically, to comprehensively exploit homogeneity, we introduce a Hierarchical Homogeneity Memory (H2M) module that progressively distills intra- and inter-task homogeneity priors from high-quality images during training, and adaptively retrieves the most relevant priors tailored to the input for guided restoration. These retrieved priors are then injected into the restoration pipeline via an efficient Homogeneity-Guided Attention (HGA) mechanism. Furthermore, to comprehensively address heterogeneity, we design a Hierarchical Heterogeneity Balancer (H2B) that mitigates both inter- and intra-task conflicts during optimization, facilitating balanced and effective multi-task learning. Extensive experiments on two large-scale benchmarks, MedIR-2D-500K and MedIR-3D-3K, demonstrate that UniH3 achieves state-of-the-art performance on both all-in-one and single-task medical image restoration. We hope this work establishes a strong benchmark and advances the development of general-purpose medical image restoration models. Code is available at https://github.com/Yaziwel/UniH3.