UniH^3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration
UniH3 unifies hierarchical homogeneity and heterogeneity modeling for all-in-one medical image restoration across modalities and degradation types.
UniH3 introduces a Hierarchical Homogeneity Memory module that distills shared anatomical priors from high-quality images, injected via a Homogeneity-Guided Attention mechanism. A Hierarchical Heterogeneity Balancer mitigates inter- and intra-task conflicts during multi-task optimization. It achieves state-of-the-art on MedIR-2D-500K and MedIR-3D-3D benchmarks for both all-in-one and single-task restoration, with code released on GitHub.
- H2M module retrieves homogeneity priors to guide restoration
- HGA mechanism injects retrieved priors into the pipeline
- H2B balancer addresses multi-task optimization conflicts
- SOTA on MedIR-2D-500K and MedIR-3D-3K benchmarks
- Code released at github.com/Yaziwel/UniH3
Full article216 words · extracted from huggingface.co · click to collapse
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.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.11156