TAT: Task-Adaptive Transformer for All-in-One Medical Image Restoration

TL;DR

This paper introduces TAT, a task-adaptive Transformer that dynamically adjusts to multiple medical image restoration tasks, effectively mitigating interference and imbalance to achieve state-of-the-art results across diverse modalities.

Contribution

The paper proposes a novel task-adaptive Transformer framework with dynamic weight generation and loss balancing strategies for multi-task medical image restoration.

Findings

Achieves state-of-the-art results in PET synthesis, CT denoising, and MRI super-resolution.

Effectively mitigates task interference and imbalance in multi-task learning.

Demonstrates superior performance in both task-specific and All-in-One settings.

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.