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

Source-linked AI summary

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

arXiv:2609.11156v1Sep 10, 2026cs.CV

TL;DR

All-in-one medical image restoration must handle diverse modalities and degradations, but existing methods largely overlook shared anatomical priors while focusing on task differences. UniH3 unifies hierarchical homogeneity modeling with hierarchical heterogeneity balancing, and achieves state-of-the-art performance on both all-in-one and single-task restoration. The authors also identify a scope boundary: the study covers only the primary restoration task within each modality.

Problem
Existing all-in-one MedIR methods emphasize inter-task heterogeneity while overlooking shared anatomical structures and intra-task variations, limiting their use of hierarchical medical-image knowledge.
Method
UniH3 combines H2M and HGA to retrieve and use intra- and inter-task homogeneity priors, while H2B mitigates optimization conflicts at both hierarchical heterogeneity levels.
Results
UniH3 achieves state-of-the-art performance on both all-in-one and single-task medical image restoration across the MedIR-2D-500K and MedIR-3D-3K benchmarks.
Takeaways & Limitations
The study presents UniH3 as a unified framework that leverages hierarchical homogeneity while balancing heterogeneous multi-task optimization for general-purpose MedIR.
Takeaways & Limitations
The study covers only the primary restoration task within each modality and does not address other tasks or degradation types in the same modality.

Abstract

from arXiv · show

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.

1 Introduction

UniH3 addresses the limitations of all-in-one medical image restoration by modeling hierarchical homogeneity alongside inter- and intra-task heterogeneity. Its H2M, HGA, and H2B components support guided restoration and balanced multi-task learning, achieving state-of-the-art results on all-in-one and single-task restoration.

1 Introduction: Existing all-in-one methods emphasize inter-task differences but largely overlook shared anatomical structures and intra-task variations across scanners, centers, and demographics.The paper argues that hierarchical modeling of both homogeneity and heterogeneity is needed for robust all-in-one restoration.
1 Introduction: H2M distills intra- and inter-task priors from high-quality images, retrieves input-relevant structural knowledge, and injects it through HGA to guide restoration.The retrieved priors are tailored to the input during inference.
1 Introduction: H2B dynamically resolves optimization conflicts at both inter- and intra-task levels, enabling balanced convergence across diverse data distributions.This extends beyond traditional strategies that balance losses only at the task level.
1 Introduction: The benchmark contains 509,200 2D image pairs across seven tasks and 3,522 3D volume pairs across three tasks.These datasets are named MedIR-2D-500K and MedIR-3D-3K, respectively.
1 Introduction: UniH3 achieves state-of-the-art performance in both all-in-one and single-task medical image restoration.The evaluation uses the MedIR-2D-500K and MedIR-3D-3K benchmarks.

2 Related Work

Medical image restoration has traditionally been studied as separate modality-specific tasks, while all-in-one methods seek to adapt one model across diverse restoration problems. In the medical domain, this unified direction remains nascent and includes prompt-based and mixture-of-experts approaches.

2 Related Work: Single-task MedIR methods separately address modality-specific degradations, including PET and CT denoising, MRI super-resolution, and other imaging tasks.The related work surveys denoising and super-resolution across PET, CT, MRI, X-ray, OCT, ultrasound, and pathology images.
2 Related Work: All-in-one restoration uses prompts to guide a shared network, with methods such as PromptIR and AdaIR learning discriminative task-adaptation mechanisms.This prompt-based paradigm is described as dominant in general all-in-one restoration.
2 Related Work: Medical all-in-one restoration remains nascent; AMIR pioneered unified modeling for three medical restoration tasks using a mixture-of-experts strategy.The passage identifies AMIR as an early medical-domain effort.

3 Method

UniH3 combines hierarchical homogeneity priors with hierarchical heterogeneity balancing in a U-shaped all-in-one restoration pipeline. H2M and HGA guide restoration using clean anatomical information, while H2B addresses optimization conflicts across tasks and samples.

3 Method: The U-shaped backbone uses four H2M modules and HGATBs to provide multi-resolution global and local restoration guidance.HGATBs use HGA for global interactions and convolution with Squeeze-and-Excitation for local interactions.
3.1 Hierarchical Homogeneity Memory: H2M distills intra-task and inter-task anatomical priors into memory and retrieves input-relevant priors through cross-attention.Task-specific and task-shared slots are updated with HQ features during training, then queried using LQ features for restoration guidance.
3.2 Homogeneity-Guided Attention: UniH3 injects retrieved homogeneity priors through HGA, which combines LQ and HQ-derived values while biasing attention toward the more reliable prior.Learnable channel-wise weights control the prior-guided terms, and setting both weights to zero recovers conventional self-attention.
3.3 Hierarchical Heterogeneity Balancer: H2B balances optimization hierarchically by modeling both task-level and sample-level uncertainty instead of using only one scalar per task.The task term addresses inter-task heterogeneity, while the sample-specific correction addresses intra-task variation and hard samples.

4 Experiments

UniH3 is evaluated across all-in-one and single-task 2D and 3D medical image restoration settings using large-scale benchmarks, with component and ablation analyses. It achieves strong benchmark performance while H2M, HGA, and H2B each contribute to restoration quality and task balancing.

2D All-in-One MedIR: 0.21 dB PSNR over AdaIR is UniH3’s average advantage across seven 2D all-in-one tasks on MedIR-2D-500K.UniH3 significantly outperforms all comparison methods across the seven tasks, while visual comparisons show higher structural fidelity and finer detail.
3D Single-Task MedIR: 1.48 dB average PSNR over Spach Transformer is achieved by UniH3-3D across three 3D single-task tasks.The model also beats all comparison methods in 3D single-task restoration.
2D Single-Task MedIR: 0.15 dB PSNR over MambaIR is UniH3’s average improvement across seven 2D single-task tasks.UniH3 significantly outperforms all compared general image-restoration methods in this setting.
Component Analysis: Both H2M and H2B improve performance with minimal computational increase when selectively disabled in component analysis.The analysis removes H2M or replaces its HGA mechanism, and substitutes H2B’s loss with L1 loss.
H2M Retrieval: H2M retrieval shows 8/10 top-score overlap for two PET spine tokens, 2/10 PET–CT spine overlap, and 0/10 PET spine–lesion overlap.These retrieval patterns provide visual evidence that H2M maps, stores, and retrieves hierarchical anatomical priors according to similarity.
H2B Ablation: Jointly mitigating inter-task and intra-task heterogeneity produces the best H2B ablation results.The proposed uncertainty formulation estimates task- and sample-level uncertainty, capturing finer-grained relationships than task-only balancing.

5 Discussion and Limitation

UniH3 is presented as practical for all-in-one MedIR, matching SOTA single-task performance across seven tasks and transferring learned knowledge to improve individual tasks. The authors note that the study covers only primary restoration tasks within each modality.

Discussion: Comparable performance to SOTA single-task MambaIR models across seven MedIR tasks supports the practicality of one all-in-one UniH3 model.The comparison is reported in Fig. 8.
Discussion: Fine-tuning UniH3 for individual tasks yields further improvements, indicating its potential as a transferable pretrained backbone.These improvements are also shown in Fig. 8.
Limitation: The study covers only the primary restoration task within each modality, excluding other tasks and degradation types that may occur in the same modality.

6 Conclusion

UniH3 unifies hierarchical homogeneity and heterogeneity for all-in-one medical image restoration through memory-guided priors and balanced multi-task learning. Evaluations on MedIR-2D-500K and MedIR-3D-3K report state-of-the-art performance in all-in-one and single-task settings.

Conclusion: UniH3 achieves state-of-the-art performance in both all-in-one and single-task medical image restoration on MedIR-2D-500K and MedIR-3D-3K.
Conclusion: H2M and HGA distill and use inter- and intra-task homogeneity priors to guide restoration across medical imaging modalities.
Conclusion: H2B mitigates inter- and intra-task heterogeneity to support stable and balanced multi-task learning.
Conclusion: Future work will expand the task spectrum with additional modalities and degradation types.

A Availability of Code and Data

The paper releases its code and data, while the appendix describes HGA formulations and points to tabulated dataset information and low- to high-quality image-pair simulations.

Availability of Code and Data: Code and data are released through the UniH3 GitHub repository.The stated repository is https://github.com/Yaziwel/UniH3.
HGA: HGA uses query, key, and value tensors Q, K, and V to derive transposed self-attention.The tensors are specified as Q, K, V ∈ R^(H′W′×C′).
HGA: The transposed self-attention-based HGA augments attention with addition and subtraction operations on the value.The formulation is given in Eq. 16 and illustrated in Fig. 3.
Datasets: Dataset information is provided in Tab. I, followed by descriptions of private data and simulated LQ–HQ image pairs.

C.1 Private Data Source

The study collected private PET and CT datasets for its experiments, including three private datasets described in Tab. I.

Private Data Source: Private1 and Private2 contain PET data, while Private3 contains CT data collected for this study.The experiments involving all three private datasets received ethics approval, and informed consent was obtained from participating patients.

C. ADDITIONAL DATASET INFORMATION 21

The additional dataset information describes privately collected 3D PET and CT data, with 2D slices extracted after excluding slices without anatomical content.

The dataset information includes 90 3D whole-body PET images collected with a PolarStar Flight PET system.The average administered dose was 301 MBq of 18F-FDG.
2D images are extracted from 3D scans after excluding slices without anatomical content, such as air-only regions.
The dataset information includes 128 3D CT images collected with a Sinovision CT system.The images cover spine, lung, and soft-tissue examinations.

C.2 Methods for Generating LQ-HQ image Pairs

The paper generates low-quality/high-quality image pairs using degradation simulations tailored to the dominant imaging artifacts of each modality.

PET denoising pairs use randomly subsampled listmode data to simulate a dose reduction factor of 10 before standard OSEM reconstruction.Both high-quality and low-quality PET images undergo OSEM reconstruction.
The simulations focus primarily on the key degradation affecting each imaging modality.
OCT denoising pairs use five repeated scans averaged as the high-quality image and one randomly selected original scan as the low-quality image.

D. ADDITIONAL EXPERIMENTS. 23

Additional experiments evaluate UniH3 on 2D and 3D all-in-one medical image restoration across multiple tasks and datasets. UniH3 consistently outperforms comparison methods on the AMIR benchmark and best preserves structures and details in the reported visual comparisons.

UniH3 consistently outperforms all comparison methods across PET denoising, CT denoising, and MRI super-resolution on the AMIR dataset.The comparison results are reported in Table II.
UniH3 best preserves structures and details across seven 2D all-in-one restoration tasks.The visual comparisons are provided in Fig. I and Fig. II.
UniH3-3D best preserves structures and details across three 3D all-in-one restoration tasks.The visual comparison spans coronal, sagittal, and transverse planes in Fig. III.