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

2025-12-16

AI Summary

In medical image restoration (MedIR), “All-in-One” models must jointly handle diverse multimodal and multi-degradation tasks—including PET synthesis, CT denoising, and MRI super-resolution—yet suffer from gradient interference (task conflict) and task imbalance (heterogeneous learning difficulty). To address these challenges, we propose the first full-task adaptive framework tailored for MedIR. Our method introduces a task-conditioned lightweight adapter module that dynamically generates parameter weights to mitigate gradient conflicts, coupled with a differentiable loss-weighting mechanism guided by estimated task difficulty to ensure balanced cross-task optimization. Built upon a Transformer architecture, the framework achieves both strong representational capacity and computational efficiency. Evaluated under an All-in-One setting across three task categories, our approach achieves significant PSNR and SSIM improvements over state-of-the-art methods. The source code is publicly available.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Multi-modal Vision
• Intelligent Robots: Multimodal Perception & Sensor Fusion

Application Category

• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• User Modeling, Personalization and Recommendation: Federated recommendation systems and personalization
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Addresses task interference in multi-task medical image restoration
• Mitigates task imbalance across different medical imaging modalities
• Enhances All-in-One model performance for diverse restoration tasks

Innovation

• Task-adaptive weight generation to mitigate gradient conflicts
• Task-adaptive loss balancing to adjust for learning difficulties
• Transformer framework dynamically adapting to multiple medical restoration tasks