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

Conference: ECCV 2026
Paper: ECCV Official
Code: https://github.com/Yaziwel/UniH3
Area: Medical Imaging
Keywords: All-in-One Medical Image Restoration, Cross-Modality Homogeneity, Hierarchical Memory Prior, Homogeneity-Guided Attention, Hierarchical Heterogeneity Balancer

TL;DR¶

UniH3 establishes the first all-in-one medical image restoration framework that systematically unifies hierarchical homogeneity and heterogeneity, distilling and retrieving cross-task anatomical priors via Hierarchical Homogeneity Memory (H2M) and Homogeneity-Guided Attention (HGA) while resolving task- and sample-level gradient conflicts via Hierarchical Heterogeneity Balancer (H2B), achieving new SOTA performance on MedIR-2D-500K and MedIR-3D-3K benchmarks.

Background & Motivation¶

Medical image restoration (MedIR) has traditionally operated within single-task paradigms where specialized models are developed separately for distinct imaging physics, such as PET denoising, low-dose CT artifact removal, or MRI super-resolution. With the widespread adoption of multi-modal scanners in modern clinical diagnosis (e.g., PET/CT and PET/MRI), these isolated single-task models falter when handling multifaceted degradation profiles, while maintaining separate models introduces substantial deployment and maintenance footprints. While all-in-one restoration architectures have recently gained significant traction in computer vision, preliminary explorations in the medical domain (such as mixture-of-experts in AMIR) remain predominantly fixated on task heterogeneity—striving to isolate distinct tasks and data distributions using contrastive learning, degradation classification, or visual prompts.

However, this conventional heterogeneity-first philosophy overlooks two fundamental realities in medical imaging. First, existing methods neglect the inherent anatomical homogeneity present across medical images. Unlike open-world natural images with chaotic semantic distributions, medical images across modalities and degradation levels consistently reflect shared human anatomical structures and strong biological topologies. Ignoring this rich homogeneous prior prevents models from leveraging cross-task synergies, rendering multi-task optimization increasingly difficult as the number of tasks expands. Second, prior heterogeneity modeling remains restricted to a coarse inter-task level (e.g., distinguishing PET from MRI), while entirely overlooking fine-grained intra-task variations caused by diverse scanner models, acquisition protocols, institutional centers, and patient demographics. This coarse balancing fails to alleviate subtle sample-level optimization conflicts.

To resolve this dilemma, all-in-one restoration requires a unified paradigm that simultaneously models homogeneity and heterogeneity at both inter- and intra-task levels. Core idea: propose UniH3 to unify hierarchical homogeneity and heterogeneity, utilizing a Hierarchical Homogeneity Memory (H2M) and Homogeneity-Guided Attention (HGA) to distill and retrieve anatomical priors as reconstruction anchors, while employing a Hierarchical Heterogeneity Balancer (H2B) to dynamically resolve inter- and intra-task optimization conflicts.

Method¶

Overall Architecture¶

UniH3 is built upon a 4-level asymmetric U-shaped encoder-decoder backbone. For an input degraded low-quality medical image $I_{LQ} \in \mathbb{R}^{H \times W \times 1}$, a $3 \times 3$ convolutional stem first extracts shallow features $I_S \in \mathbb{R}^{H \times W \times C}$ with base channel dimension $C=48$. The encoder-decoder hierarchy stacks $N_1=2, N_2=3, N_3=3, N_4=4$ Homogeneity-Guided Transformer Blocks (HGATBs) across four resolution stages. At each resolution level, Hierarchical Homogeneity Memory (H2M) modules are integrated to provide adaptive anatomical priors across scales. Each HGATB incorporates two consecutive units: a Homogeneity-Guided Attention (HGA) unit that captures global context anchored on retrieved priors, followed by a Convolution plus Squeeze-and-Excitation (SE) unit that models fine-grained local textures. Finally, deep features $I_D$ are projected to a residual map $I_R$, yielding the high-quality restoration $\hat{I}_{HQ} = I_{LQ} + I_R$. Optimization is dynamically regularized by the Hierarchical Heterogeneity Balancer (H2B).

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    Input["Input Degraded Image ILQ"] --> Stem["Shallow Feature Projection<br/>3x3 Conv maps input to C channels"]
    Stem --> H2M["Hierarchical Homogeneity Memory H2M<br/>Adaptive retrieval of cross- and intra-task prior VH"]
    H2M --> HGA["Homogeneity-Guided Attention HGA<br/>Prior-anchored global feature interaction"]
    HGA --> HGATB["Local Interaction Enhancement<br/>Conv + SE layer capturing local details"]
    HGATB --> Recon["Residual Reconstruction Projection<br/>Predicts IR and generates restored image"]
    Recon --> H2B["Hierarchical Heterogeneity Balancer H2B<br/>Dual task- and sample-level uncertainty balancing"]
    H2B --> Output["Restored High-Quality Output IHQ-hat"]

Key Designs¶

1. Hierarchical Homogeneity Memory: Disentangled Distillation and Adaptive Retrieval of Anatomical Priors

As the number of supported modalities and restoration tasks grows, training burden escalates without shared representations. H2M explicitly addresses this by establishing two structurally symmetric components: a memory bank $M \in \mathbb{R}^{(T+1)L \times C'}$ and a learnable prototype matrix $P \in \mathbb{R}^{(T+1)L \times C'}$. These matrices are partitioned into $T$ task-specific slots (capturing intra-task anatomical consistency across patients) and 1 task-shared slot (capturing inter-task structural invariants across modalities), each with slot length $L=128$. The module operates in two phases:

Homogeneity Distillation (training only): Paired images $I_{LQ}$ and $I_{HQ}$ are mapped to target feature dimensions $F^{LQ}, F^{HQ} \in \mathbb{R}^{H'W' \times C'}$ using pixel-unshuffle downsampling and a $3 \times 3$ convolution. Learnable prototype tokens in $P$ query ground-truth high-quality features $F^{HQ}$ via cross-attention $V^{HQ} = \text{Softmax}(P (F^{LQ})^\mathsf{T} / \sqrt{C'}) F^{HQ}$. Conditioned on the active task identifier, extracted shared features $V^{Sh}$ and task-specific features $V^{Sp}$ are archived into memory $M$ via an Exponential Moving Average (EMA, momentum $\alpha=0.99$): $$M_{Sh/Sp} \leftarrow \alpha M_{Sh/Sp} + (1 - \alpha) V^{Sh/Sp}$$ During testing, this distillation branch is discarded, guaranteeing zero dependence on ground-truth images.
Homogeneity Retrieval (training and testing): Given input degraded features $F^{LQ}$ as queries, prototype $P$ as keys, and distilled memory bank $M$ as values, cross-attention dynamically retrieves clean anatomical priors $V^H$: $$V^H = \text{CrossAttention}(F^{LQ}, P, M)$$ Four multi-scale H2M modules operate at distinct stages of the U-shaped hierarchy, supplying multi-resolution structural guidance from coarse organ geometry to fine-grained tissue textures.

2. Homogeneity-Guided Attention: Biased Reconstruction Anchored on High-Quality Priors

Conventional restoration guidance schemes (e.g., SFT or standard cross-attention) treat corrupted low-quality features as the primary anchor and external guidance as secondary modulation, forcing the network to infer clean structures from damaged baselines. HGA fundamentally shifts this paradigm by establishing the retrieved high-quality prior $V^H$ as the foundational anchor for representation learning.

In standard self-attention with attention map $A = \text{Softmax}(Q K^\mathsf{T} / \sqrt{C'})$, a naive linear fusion $A(V + V^H)$ treats degraded features $V$ and prior features $V^H$ symmetrically. Because $V^H$ carries reliable anatomical topology, attention should preferentially rely on $V^H$. HGA introduces identity terms into the attention operator $(A - I)V + (A + I)V^H = A(V + V^H) + (V^H - V)$. The $\pm I$ terms suppress the direct feedthrough of degraded features $V$ while amplifying the self-contribution of prior $V^H$, with $(V^H - V)$ serving as an explicit preference bias. Channel-wise learnable weights $\lambda_1, \lambda_2 \in \mathbb{R}^{C'}$ are incorporated to guarantee stability and expressive capacity: $$V^O = A[(1 - \lambda_1)V + \lambda_1 V^H] + \lambda_2 (V^H - V)$$ When $\lambda_1 = \lambda_2 = 0$, HGA gracefully reverts to standard self-attention. UniH3 implements HGA over transposed self-attention across channel dimensions following Restormer, maintaining linear computational complexity while delivering robust structural restoration.

3. Hierarchical Heterogeneity Balancer: Multi-Task Optimization via Dual-Level Uncertainty

Heterogeneous restoration objectives introduce conflicting gradient vectors that cause certain tasks to dominate the parameter updates while harder tasks remain undertrained. Traditional uncertainty-based multi-task balancing optimizes a single task-level homoscedastic uncertainty parameter $\sigma_t$, which remains too coarse to resolve intra-task variations stemming from diverse imaging centers, scanners, and abnormal pathology distributions.

H2B introduces a hierarchical uncertainty formulation that decomposes the total uncertainty $\sigma_{t,s}$ for sample $s$ of task $t$ into a macro task-level parameter $\sigma_t$ and an adaptive sample-specific correction $\Delta \sigma_s$: $$\sigma_{t,s} = \sigma_t + \Delta \sigma_s$$ While $\sigma_t$ is a learnable task scalar, $\Delta \sigma_s$ is dynamically inferred by a lightweight Uncertainty Estimation Block (UEB): $$\Delta \sigma_s = \text{UEB}\left( \text{Concat}[I_s^{LQ}, \text{sg}(\hat{I}_s^{HQ}), I_s^{HQ}] \right)$$ where $\text{sg}(\cdot)$ denotes a stop-gradient operator that prevents uncertainty estimation from leaking gradients into the restoration backbone. The overall H2B objective aggregates loss contributions across tasks $T$ and batch size $S$: $$\mathcal{L}_{H2B} = \frac{1}{TS}\sum_{t=1}^T \sum_{s=1}^S \left( \frac{1}{2\sigma_{t,s}^2} \mathcal{L}_{rec}^{(t,s)} + \log \sigma_{t,s} \right)$$ using standard $L_1$ reconstruction loss $\mathcal{L}_{rec}$. For abnormally degraded samples, $\sigma_{t,s}$ increases to attenuate destabilizing gradient surges, while well-behaved samples yield lower uncertainty to strengthen supervisory signals, ensuring stable and balanced multi-task convergence.

Loss & Training¶

The entire framework is optimized end-to-end using the H2B loss $\mathcal{L}_{H2B}$. On the 2D MedIR-2D-500K benchmark, the model is trained with patch size $128 \times 128$ and batch size 14 using the Muon optimizer. Training runs for $6 \times 10^5$ iterations with an initial learning rate of $3 \times 10^{-4}$ decayed to $1 \times 10^{-7}$ via cosine annealing. For 3D volumetric restoration, UniH3-3D adopts 3D convolutional and attention operators, configuring HGATB depths to $N_1=N_2=1, N_3=N_4=5$, channel width $C=16$, patch size $64 \times 64 \times 64$, batch size 6, and initial learning rate $5 \times 10^{-5}$.

Key Experimental Results¶

Main Results¶

Quantitative evaluations are conducted on MedIR-2D-500K (509,200 image pairs across 7 tasks) and MedIR-3D-3K (3,522 volume pairs across 3 tasks). The table below summarizes comparative performance under the 2D All-in-One setting:

Method	#Params (M)	FLOPs (G)	PET (PSNR/SSIM)	CT (PSNR/SSIM)	MRI (PSNR/SSIM)	X-ray (PSNR/SSIM)	OCT (PSNR/SSIM)	Ultrasound (PSNR/SSIM)	Pathology (PSNR/SSIM)	Average (PSNR/SSIM)
SwinIR	11.50	187.93	44.24 / 0.9866	43.17 / 0.9338	38.44 / 0.9453	35.87 / 0.9275	35.70 / 0.8892	27.52 / 0.8089	28.35 / 0.7941	36.18 / 0.8979
Uformer	50.47	21.42	44.51 / 0.9874	43.45 / 0.9356	39.01 / 0.9507	36.59 / 0.9353	35.84 / 0.8908	27.61 / 0.8135	28.55 / 0.8011	36.51 / 0.9020
Restormer	26.12	35.21	44.47 / 0.9873	43.44 / 0.9353	39.05 / 0.9515	36.46 / 0.9335	35.84 / 0.8909	27.66 / 0.8138	28.53 / 0.8006	36.49 / 0.9018
NAFNet	67.89	15.74	44.40 / 0.9871	43.32 / 0.9346	38.90 / 0.9502	36.33 / 0.9326	35.82 / 0.8905	27.59 / 0.8131	28.49 / 0.7995	36.41 / 0.9011
AMIR	23.54	31.76	44.49 / 0.9873	43.47 / 0.9356	39.09 / 0.9519	36.47 / 0.9333	35.89 / 0.8914	27.69 / 0.8150	28.57 / 0.8019	36.52 / 0.9023
PromptIR	35.59	39.49	44.52 / 0.9874	43.48 / 0.9355	39.13 / 0.9524	36.57 / 0.9341	35.84 / 0.8909	27.69 / 0.8152	28.54 / 0.8006	36.54 / 0.9023
AdaIR	28.76	36.74	44.55 / 0.9875	43.49 / 0.9356	39.17 / 0.9527	36.60 / 0.9344	35.86 / 0.8915	27.69 / 0.8150	28.56 / 0.8015	36.56 / 0.9026
UniH3 (Ours)	28.96	26.33	44.89 / 0.9883	43.65 / 0.9368	39.55 / 0.9564	36.88 / 0.9368	35.96 / 0.8921	27.80 / 0.8179	28.63 / 0.8035	36.77 / 0.9045

In the 3D All-in-One setting (MedIR-3D-3K), UniH3-3D attains 44.50 dB average PSNR and 0.9674 SSIM, surpassing the previous state-of-the-art Restore-RWKV-3D (43.95 dB / 0.9644) by a remarkable margin of +0.55 dB.

Ablation Study¶

Ablation studies on MedIR-2D-500K validate the efficacy of core architectural designs, attention mechanisms, and hierarchical components:

Category	Configuration	#Params (M)	FLOPs (G)	Avg PSNR (dB)	Avg SSIM	Observation & Note
Component Analysis	Baseline (w/o H2M & H2B)	27.17	25.50	36.52	0.9023	Standard transposed attention + L1 loss
	+ H2M (Memory + HGA only)	28.96	26.33	36.66	0.9034	Structural prior injection gains +0.14 dB
	+ H2B (Balancer only)	27.17	25.50	36.64	0.9033	Uncertainty balancing gains +0.12 dB
	UniH3 Full Model (+ H2M + H2B)	28.96	26.33	36.77	0.9045	Joint synergy achieves cumulative +0.25 dB gain
Attention Variant	w/o HGA (Standard self-attention)	27.17	25.50	36.52	0.9023	Lacks anatomical prior anchoring
	SFT Mechanism [55]	37.46	36.99	36.74	0.9041	Substantially higher parameters and FLOPs
	Cross Attention [11]	29.60	27.43	36.67	0.9035	Symmetric interaction lacks prior preference
	HGA (Ours)	28.96	26.33	36.77	0.9045	Superior quality with minimal compute overhead
H2M Homogeneity	Inter-task Homogeneity only	28.96	26.33	36.61	0.9031	Captures global anatomical invariants
	Intra-task Homogeneity only	28.96	26.33	36.72	0.9038	Captures modality-specific consistency
	Joint Inter + Intra Homogeneity	28.96	26.33	36.77	0.9045	Both hierarchical priors are indispensable
H2B Heterogeneity	Inter-task Heterogeneity only ($\sigma_t$)	28.96	26.33	36.70	0.9036	Coarse task-level uncertainty balancing
	Intra-task Heterogeneity only ($\Delta \sigma_s$)	28.96	26.33	36.73	0.9041	Sample-level correction without task anchor
	Joint Inter + Intra Heterogeneity	28.96	26.33	36.77	0.9045	Resolves macro and fine-grained conflicts

Key Findings¶

Synergy of Homogeneity and Heterogeneity: Modeling anatomical homogeneity (H2M +0.14 dB) and resolving heterogeneity conflicts (H2B +0.12 dB) provide mutually reinforcing benefits, cumulatively lifting PSNR by +0.25 dB over the baseline.
Biologically Meaningful Retrieval Patterns: Visual inspection of retrieval attention reveals that tokens corresponding to the same anatomy (e.g., two PET spine tokens) share an 8/10 Top-10 slot overlap; cross-modality spine tokens (PET and CT) share 2/10 overlap; while anatomically divergent regions (PET spine vs. PET lesion) have 0/10 overlap, demonstrating accurate structural clustering in memory space.
Universal Pretrained Backbone Potential: A single all-in-one UniH3 model (36.77 dB) outperforms dedicated single-task MambaIR models (36.75 dB). Fine-tuning pretrained UniH3 on individual tasks further elevates average PSNR to 36.97 dB, confirming its efficacy as a general-purpose pretrained foundation for medical vision.

Highlights & Insights¶

Prior-Anchored Attention Bias: By embedding identity-based biases $(A-I)V + (A+I)V^H$ into the attention operator, HGA fundamentally anchors feature reconstruction onto clean anatomical memory, preventing feature degradation without incurring excessive parameter bloat.
Hierarchical Memory Disentanglement: Decoupling anatomical priors into task-shared and task-specific memory slots enables the network to absorb universal human anatomy while preserving modality-specific contrast physics.
Plug-and-Play Generalizability: Integrating H2M and H2B into existing architectures (Uformer, Restormer, PromptIR, AdaIR) yields consistent 0.15~0.20 dB PSNR improvements across all backbones, verifying the universal compatibility of the proposed paradigm.

Limitations & Future Work¶

Single Degradation Assumption per Modality: Current benchmarks pair each modality with one primary physical degradation (e.g., denoising for CT/PET, super-resolution for MRI), leaving compound degradations and unpaired clinical settings for future exploration.
Training-Time Memory Overhead: The sample-level uncertainty estimator UEB requires processing generated predictions $\hat{I}^{HQ}$, which slightly increases peak GPU memory footprint during multi-task training (though test-time inference remains unaffected).
Future Directions: Integrating multimodal anatomical text descriptions via vision-language alignment into H2M memory tokens, and verifying restored diagnostic utility in downstream pathology segmentation and detection tasks.

vs AdaIR / PromptIR (Prompt-based All-in-One): Prompt-driven methods primarily focus on task discrimination (heterogeneity), neglecting explicit storage and reuse of shared anatomical topologies. UniH3 introduces explicit memory banks to distill and retrieve cross-modality structural invariants, achieving superior fidelity.
vs AMIR (Medical Multi-Task MoE): AMIR relies on routing networks and task-specific experts, incurring high parameter counts and risking expert over-specialization. UniH3 maintains a unified and compact backbone (28.96M parameters, 26.33G FLOPs) and resolves task competition via dynamic uncertainty balancing, outperforming AMIR by +0.25 dB.

Rating¶

Novelty: ⭐⭐⭐⭐⭐ Pioneering unification of hierarchical homogeneity and heterogeneity in all-in-one medical image restoration.
Experimental Thoroughness: ⭐⭐⭐⭐⭐ Comprehensive validation across MedIR-2D-500K and MedIR-3D-3K spanning 2D/3D, multi-task/single-task, and cross-backbone studies.
Writing Quality: ⭐⭐⭐⭐⭐ Rigorous mathematical formulations, clear intuitive motivations, and well-structured empirical analyses.
Value: ⭐⭐⭐⭐⭐ Provides an open-source, highly effective benchmark and foundation architecture for universal medical image restoration.