All-In-One Medical Image Restoration via Task-Adaptive Routing

TL;DR

The paper tackles the challenge of all‑in‑one medical image restoration by introducing AMIR, a universal model that handles MRI super‑resolution, CT denoising, and PET synthesis through task‑adaptive routing. It combines a Unet‑style Restormer backbone with a Routing Instruction Network, Spatial Routing Modules, and Channel Routing Modules to dynamically route task‑specific information and mitigate gradient interference. The method demonstrates state‑of‑the‑art performance in single‑task and all‑in‑one settings across three MedIR tasks, with extensive ablations confirming the importance of the routing components and instruction learning. This approach enables a scalable, interpretable, and efficient framework for multi‑task MedIR, potentially reducing model maintenance and improving generalization across diverse imaging modalities. The results suggest strong practical impact for deploying a single universal model in clinical imaging workflows while preserving task specificity and image quality, with further work extending to more tasks.

Abstract

Although single-task medical image restoration (MedIR) has witnessed remarkable success, the limited generalizability of these methods poses a substantial obstacle to wider application. In this paper, we focus on the task of all-in-one medical image restoration, aiming to address multiple distinct MedIR tasks with a single universal model. Nonetheless, due to significant differences between different MedIR tasks, training a universal model often encounters task interference issues, where different tasks with shared parameters may conflict with each other in the gradient update direction. This task interference leads to deviation of the model update direction from the optimal path, thereby affecting the model's performance. To tackle this issue, we propose a task-adaptive routing strategy, allowing conflicting tasks to select different network paths in spatial and channel dimensions, thereby mitigating task interference. Experimental results demonstrate that our proposed \textbf{A}ll-in-one \textbf{M}edical \textbf{I}mage \textbf{R}estoration (\textbf{AMIR}) network achieves state-of-the-art performance in three MedIR tasks: MRI super-resolution, CT denoising, and PET synthesis, both in single-task and all-in-one settings. The code and data will be available at \href{https://github.com/Yaziwel/All-In-One-Medical-Image-Restoration-via-Task-Adaptive-Routing.git}{https://github.com/Yaziwel/AMIR}.

Figures (4)

• Figure 1: (a) Examples of LQ/HQ pairs for three different MedIR tasks. (b) The interference metric zhu2022uniperceiver-moe of task $j$ on task $i$ at the second and last blocks in Restormer zamir2022restormer. Red values indicate that task $j$ negatively impacts task $i$, while green values indicate a positive impact.
• Figure 2: Overview of the proposed all-in-one medical image restoration (AMIR) network.
• Figure 3: Visual comparison of different methods on all-in-one medical image restoration.
• Figure 4: (a) t-SNE visualization of $I^{IR}$ from different tasks, indicating a clear clustering of different task inputs. (b) Top-1 selected expert in each spatial routing module (SRM). In our AMIR network setting, there are 3 SRMs, each incorporating a mixture of experts (MOE) with 4 experts. Within these SRMs, the top-1 selected expert is identified across the 3 SRMs for each task. Remarkably, the top experts selected across SRMs form distinct paths, with variations observed across different tasks.