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

arXiv: https://arxiv.org/abs/2512.14550 · HTML: https://mlsanity.com/paper/2512.14550/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 27 references that are in the map.
- First posted on arXiv 16 Dec 2025.

## Builds on (most cited first, 18 of 18)

- [Restormer: Efficient Transformer for High-Resolution Image Restoration](https://mlsanity.com/paper/2111.09881/) (2021)
- [Multi-Task Learning Using Uncertainty to Weigh Losses for Scene Geometry and Semantics](https://mlsanity.com/paper/1705.07115/) (2017)
- [CMT: Convolutional Neural Networks Meet Vision Transformers](https://mlsanity.com/paper/2107.06263/) (2021)
- [Low-Dose CT with a Residual Encoder-Decoder Convolutional Neural Network (RED-CNN)](https://mlsanity.com/paper/1702.00288/) (2017)
- [PromptIR: Prompting for All-in-One Blind Image Restoration](https://mlsanity.com/paper/2306.13090/) (2023)
- [PET image denoising based on denoising diffusion probabilistic models](https://mlsanity.com/paper/2209.06167/) (2022)
- [Swin Transformer for Fast MRI](https://mlsanity.com/paper/2201.03230/) (2022)
- [CTformer: Convolution-free Token2Token Dilated Vision Transformer for Low-dose CT Denoising](https://mlsanity.com/paper/2202.13517/) (2022)
- [Restore-RWKV: Efficient and Effective Medical Image Restoration with RWKV](https://mlsanity.com/paper/2407.11087/) (2024)
- [Eformer: Edge Enhancement based Transformer for Medical Image Denoising](https://mlsanity.com/paper/2109.08044/) (2021)
- [All-In-One Medical Image Restoration via Task-Adaptive Routing](https://mlsanity.com/paper/2405.19769/) (2024)
- [EDCNN: Edge enhancement-based Densely Connected Network with Compound Loss for Low-Dose CT Denoising](https://mlsanity.com/paper/2011.00139/) (2020)
- [Swin Deformable Attention U-Net Transformer (SDAUT) for Explainable Fast MRI](https://mlsanity.com/paper/2207.02390/) (2022)
- [3D Segmentation Guided Style-based Generative Adversarial Networks for PET Synthesis](https://mlsanity.com/paper/2205.08887/) (2022)
- [DenoMamba: A fused state-space model for low-dose CT denoising](https://mlsanity.com/paper/2409.13094/) (2024)
- [DRMC: A Generalist Model with Dynamic Routing for Multi-Center PET Image Synthesis](https://mlsanity.com/paper/2307.05249/) (2023)
- [Region Attention Transformer for Medical Image Restoration](https://mlsanity.com/paper/2407.09268/) (2024)
- [Prompted Contextual Transformer for Incomplete-View CT Reconstruction](https://mlsanity.com/paper/2312.07846/) (2023)

## Built on by (most cited first, 3 of 3)

- [Vision-Language Controlled Deep Unfolding for Joint Medical Image Restoration and Segmentation](https://mlsanity.com/paper/2601.23103/) (2026)
- [MARR: Decoupling Policy, Execution, and Calibration for All-in-One Medical Image Restoration](https://mlsanity.com/paper/2609.06645/) (2026)
- [UniH3: Unifying Hierarchical Homogeneity and Heterogeneity for All-in-One Medical Image Restoration](https://mlsanity.com/paper/2609.11156/) (2026)

## Abstract (arXiv)

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.