TAT: Task-Adaptive Transformer for All-in-One Medical Image Restoration
Zhiwen Yang, Jiaju Zhang, Yang Yi, Jian Liang, Bingzheng Wei, Yan Xu
2025-12-17
Summary
This paper focuses on improving medical image restoration, which is the process of taking blurry or low-quality medical images and making them clearer and more useful for doctors. It specifically looks at a new way to handle multiple restoration tasks at once using a single computer program.
What's the problem?
When you try to use one program to fix different types of medical images (like PET scans, CT scans, and MRIs) and different problems within those images (like noise or low resolution), things get tricky. The program can get confused because the different tasks pull it in different directions during learning, causing conflicts. Also, some tasks are easier to learn than others, leading to the program focusing too much on the easy ones and neglecting the harder ones.
What's the solution?
The researchers developed a new system called a task-adaptive Transformer, or TAT. This system cleverly adjusts itself for each specific task. It does this in two main ways: first, it creates unique 'weights' for each task, preventing the learning process from getting pulled in conflicting directions. Second, it automatically adjusts how much attention the program pays to each task based on how difficult it is, ensuring that all tasks get adequate learning and none dominate the others.
Why it matters?
This research is important because it allows for a single, more efficient program to handle a variety of medical image restoration tasks. This could lead to faster and more accurate diagnoses, as doctors would have access to higher-quality images without needing separate programs for each type of scan or problem. The new system achieves better results than existing methods on several common medical imaging tasks.
Abstract
Medical image restoration (MedIR) aims to recover high-quality medical images from their low-quality counterparts. Recent advancements in MedIR have focused on All-in-One models capable of simultaneously addressing multiple different MedIR tasks. However, due to significant differences in both modality and degradation types, using a shared model for these diverse tasks requires careful consideration of two critical inter-task relationships: task interference, which occurs when conflicting gradient update directions arise across tasks on the same parameter, and task imbalance, which refers to uneven optimization caused by varying learning difficulties inherent to each task. To address these challenges, we propose a task-adaptive Transformer (TAT), a novel framework that dynamically adapts to different tasks through two key innovations. First, a task-adaptive weight generation strategy is introduced to mitigate task interference by generating task-specific weight parameters for each task, thereby eliminating potential gradient conflicts on shared weight parameters. Second, a task-adaptive loss balancing strategy is introduced to dynamically adjust loss weights based on task-specific learning difficulties, preventing task domination or undertraining. Extensive experiments demonstrate that our proposed TAT achieves state-of-the-art performance in three MedIR tasks--PET synthesis, CT denoising, and MRI super-resolution--both in task-specific and All-in-One settings. Code is available at https://github.com/Yaziwel/TAT.