UniH^3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration

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.

한 줄 요약

UniH3는 의료 이미지 복원의 다중 태스크와 모달리티를 통합하는 새로운 프레임워크로, 계층적 동질성과 이질성을 동시에 모델링한다.

핵심 기여도

Hierarchical Homogeneity Memory (H2M) 모듈을 도입하여 훈련 중 동질성 정보를 추출하고 추론 시 입력에 맞는 구조적 힌트를 주입.
Homogeneity-Guided Attention (HGA) 메커니즘을 통해 복원 과정에서 효율적으로 동질성 정보를 활용.
Hierarchical Heterogeneity Balancer (H2B)를 설계하여 태스크 간 및 내부의 이질성을 동적으로 조정.
MedIR-2D-500K (509,200개)와 MedIR-3D-3K (3,522개) 데이터셋에서 SOTA 성능 달성.

핵심 아이디어

의료 이미지 복원에서 기존 연구는 주로 태스크 간 이질성을 모델링하는 데 집중했으나, UniH3는 의료 이미지가 공유하는 해부학적 구조와 같은 **동질성**을 활용함으로써 모델 학습을 용이하게 하고 일반화 능력을 향상시킨다.
이를 위해 H2M 모듈은 훈련 중 고질량 이미지에서 **intra-task 및 inter-task 동질성 prior**를 점진적으로 추출하고, 추론 시 입력에 맞는 prior를 적응적으로 검색하여 복원 과정에 주입한다.
또한, H2B는 기존 방법이 무시하는 **intra-task 이질성**(예: 스캐너, 환자 인구 통계 등에 따른 분포 변화)을 포함한 이질성을 계층적으로 조정하여, **다양한 데이터 분포에서 균형 잡힌 학습**을 가능하게 한다.

기술적 접근법

**모델 구조**: U-형 복원 백본에 H2M과 H2B를 결합.
**H2M 모듈**:
훈련 시 HQ 이미지에서 intra- 및 inter-task prior를 추출.
메모리 길이 128, EMA 계수 0.99.
**HGA 메커니즘**:
H2M에서 검색된 prior를 복원 파이프라인에 주입.
트랜스포머 블록 내에서 global 및 local 정보를 결합.
**H2B 전략**:
loss 함수 조정 대신, **optimization conflict**를 inter- 및 intra-task 수준에서 동적으로 해결.
**하이퍼파라미터**:
2D 모델: HGATB 수 N1=2, N2=N3=3, N4=4, 입력 채널 48.
3D 모델: HGATB 수 N1=N2=1, N3=N4=5, 입력 채널 16.

주요 결과

**MedIR-2D-500K 데이터셋**에서 PSNR 36.8 dB, SSIM 0.946 (기존 최고 기준 대비 +2.1 dB, +0.03 SSIM).
**MedIR-3D-3K 데이터셋**에서 PSNR 34.2 dB, SSIM 0.931 (기존 최고 기준 대비 +1.8 dB, +0.025 SSIM).
H2M과 H2B 모듈을 제거했을 때 성능이 각각 1.2 dB, 0.8 dB 하락함을 확인 (Table 8 참조).

의의 및 한계

UniH3는 의료 이미지 복원 분야에서 **일반 목적 모델**(All-in-One) 개발을 촉진하며, 다양한 태스크와 모달리티를 통합하는 데 기여한다. 특히, H2M과 H2B를 통해 **동질성과 이질성을 동시에 모델링**함으로써 기존 방법보다 더 높은 일반화 능력을 보인다.
하지만, 현재는 2D 및 3D 의료 이미지만 다루며, **다양한 병리학적 조건**(예: 종양, 신경 퇴행 등)에 대한 훈련 데이터가 제한적이라는 한계가 있다. 또한, H2M의 prior 저장 및 검색 메커니즘은 **메모리 사용량 증가**를 동반할 수 있다.

실용적 활용

UniH3는 병원 및 연구소에서 **다양한 의료 이미지 복원**(예: PET/CT, MRI, X-ray)을 위한 단일 모델로 활용 가능하며, **의료 AI 플랫폼**에 통합되어 복잡한 복원 작업을 자동화할 수 있다. 특히, **다중 태스크 학습**을 기반으로 유지보수 비용을 줄이고, 의료 이미지 품질을 일관되게 향상시킬 수 있다.