UniH3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration
Abstract
All-in-One medical image restoration (MedIR) aims to ad-dress diverse tasks across modalities and degradation types using a singleuniversal 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 med-ical images, such as widely shared anatomical structures within andacross modalities, which can be leveraged to ease model training andimprove generalization. To this end, we propose UniH3 , a novel frame-work that Unifies Hierarchical Homogeneity and Heterogeneity for all-in-one medical image restoration. Specifically, to comprehensively exploithomogeneity, we introduce a Hierarchical Homogeneity Memory (H2 M)module that progressively distills intra- and inter-task homogeneity pri-ors from high-quality images during training, and adaptively retrieves themost relevant priors tailored to the input for guided restoration. Theseretrieved priors are then injected into the restoration pipeline via an ef-ficient Homogeneity-Guided Attention (HGA) mechanism. Furthermore,to comprehensively address heterogeneity, we design a Hierarchical Het-erogeneity Balancer (H2 B) that mitigates both inter- and intra-task con-flicts during optimization, facilitating balanced and effective multi-tasklearning. 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 imagerestoration. We hope this work establishes a strong benchmark and ad-vances the development of general-purpose medical image restorationmodels. Code is available at https://github.com/Yaziwel/UniH3.