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

UniH^3. Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration. By Zhiwen Yang and five others, from Beihang University, Tsinghua University, and one other institution. Published on the tenth of September. Its code repository is worked out: about forty-five source files, with a readme and a licence.

All-in-one medical image restoration, MedIR, aims to handle many restoration tasks, across imaging modalities and types of degradation, with one universal model. Existing methods mostly focus on the differences between tasks: their distinct data distributions and degradation types. That is heterogeneity. What they largely neglect is the homogeneity that medical images carry, such as anatomical structures that are widely shared, both within one modality and across several. The authors argue that shared structure could make training easier and generalization better. Images from different modalities still share a great deal, and the argument is that a universal model should be able to lean on that.

UniH3 unifies the two views at more than one level, which is what hierarchical means in the title. It exploits homogeneity, by retrieving shared priors to guide restoration, and it manages heterogeneity, by balancing the conflicts between tasks during optimization. The name compresses the idea: unifying hierarchical homogeneity and heterogeneity.

For homogeneity, a Hierarchical Homogeneity Memory module, H2M, progressively distills priors from high-quality images during training. These priors capture what is shared within a task and across tasks. At restoration time, the module adaptively retrieves the priors most relevant to the input. They are then injected into the restoration pipeline through an efficient Homogeneity-Guided Attention mechanism, HGA. The effect is that restoration does not start from nothing. Each degraded input is paired with priors learned from clean images that resemble it, and those priors steer what the output should look like.

For heterogeneity, a Hierarchical Heterogeneity Balancer, H2B, mitigates conflicts during optimization, both between tasks and within a task. That balance is what lets the multi-task learning stay effective, instead of one task pulling training away from the others. Where the memory draws tasks together through what they share, the balancer keeps them from interfering through what they do not.

The experiments use two large-scale benchmarks, MedIR-2D-500K and MedIR-3D-3K. On both, UniH3 achieves state-of-the-art performance, and it does so both in the all-in-one setting and on single-task medical image restoration. The authors hope it will serve as a strong benchmark for future work. Code is available. The two benchmarks cover both two-dimensional and three-dimensional medical images, so the results span both kinds of data.

The argument is that a universal restoration model should use what medical images have in common as well as what separates them. Earlier work mainly handled the differences. UniH3 adds a memory of what is shared, retrieved on demand, and a balancer for what conflicts, and treats both as parts of one design aimed at general-purpose medical image restoration.