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

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.

• The paper introduces a transformer-based architecture that generates task-specific weights and employs dynamic loss balancing for all-in-one medical image restoration.
• It demonstrates state-of-the-art performance on PET synthesis, CT denoising, and MRI super-resolution, achieving higher PSNR, SSIM, and lower RMSE compared to baselines.
• It effectively addresses task interference and imbalance through dedicated task representation extraction and sample-level adaptive loss scaling.

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

Introduction

The proliferation of modality-specific models has advanced medical image restoration (MedIR), but these models are hindered by poor cross-task generalization, redundancy in resource usage, and susceptibility to data scarcity. All-in-One approaches promise to overcome these limitations but must resolve fundamental challenges: task interference (gradient conflicts in shared parameters across tasks) and task imbalance (unequal optimization trajectories due to varying task learning difficulties). The "TAT: Task-Adaptive Transformer for All-in-One Medical Image Restoration" (2512.14550) directly addresses these challenges by introducing a transformer-based architecture with explicit task-adaptivity in both parameterization and loss balancing.

Figure 1: The TAT architecture integrates a Transformer encoder-decoder with task-adaptive specialization modules for end-to-end multi-task MedIR.

Methodology

Task-Adaptive Transformer Architecture

TAT employs a multi-level, U-shaped transformer network comprising a three-stage transformer encoder and a four-stage decoder built around Weight-Adaptive Transformer Blocks (WATBs). A dedicated Task Representation Extraction Network (TREN) processes latent features (with stopped gradients) to yield a low-dimensional but task-discriminative representation vector $Z \in \mathbb{R}^d$. This vector conditions task adaptation mechanisms throughout the decoder, inducing task-specific behavior without sacrificing architectural amortization.

Task-Adaptive Weight Generation

TAT innovates upon standard transformer sharing by generating depthwise convolutional weights specific to each task, circumventing the gradient contention that arises from multitask learning on a shared parameter set. The task-adaptive weights are computed as:

$$W^G = \operatorname{Reshape}(\operatorname{MLP}(Z)),$$

with $W = W^S + \lambda W^G$ for each depthwise convolution, where $W^S$ is the shared backbone weight and $\lambda$ is a learnable scaling factor. By leveraging depthwise convolutions, the parameter generation remains computationally tractable ($\mathcal{O}(C)$), and the integration with transformer blocks ensures effective local-global feature fusion, which is critical in diverse MedIR tasks.

Task-Adaptive Loss Balancing

Beyond standard task-level uncertainty weighting, TAT introduces a sample-level adaptive loss balancing strategy. For an input $(I^{LQ}, I^{HQ}, \hat{I}^{HQ})$, three L1 losses encoding input-target, input-prediction, and prediction-target discrepancies are pooled and passed through an MLP (with gradients detached from the main model) to generate a sample-specific scaling parameter $\sigma$:

$$\sigma = \mathrm{MLP}([L_1(I^{LQ},I^{HQ}), L_1(I^{LQ}, \hat{I}^{HQ}), L_1(\hat{I}^{HQ}, I^{HQ})]).$$

The loss is then scaled for each sample as

$$\text{Loss} = \frac{1}{2\sigma^2} L_1(\hat{I}^{HQ}, I^{HQ}) + \log \sigma,$$

yielding dynamic attenuation for hard/easy samples and allowing the network to remain responsive to sample-level imbalance within each task. This contrasts with prior work that only modulates loss terms globally per task.

Empirical Results

TAT is evaluated on three canonical MedIR tasks: PET synthesis, CT denoising, and MRI super-resolution, using standardized splits and metrics (PSNR, SSIM, RMSE). The method is benchmarked both as a task-specific model and as an All-in-One model against state-of-the-art task and multitask baselines.

Numerical Outcomes: In both task-specific and multitask regimens, TAT achieves the highest PSNR, SSIM, and lowest RMSE across all restoration tasks. Notably, in All-in-One settings, it exceeds the current SOTA AMIR [yang2024amir] by consistent and statistically significant margins, demonstrating effective mitigation of task interference and imbalance.

Figure 2: Qualitative comparison: TAT yields superior perceptual restoration across PET, CT, and MRI tasks, preserving fine structural details and suppressing artifacts robustly.

Ablation studies validate both core innovations: removing task-adaptive weight generation or using non-sample-adaptive loss balancing leads to measurable performance degradation. Crucially, attempting to generate all weights (rather than restricting to depthwise convs) introduces computational overhead without further gains, underlining the efficacy of TAT's design constraints.

Theoretical and Practical Implications

Addressing Task Interference

By dynamically generating task-specific parameters conditioned on latent task representations, TAT avoids destructive gradient interference in shared multitask settings, a challenge identified but not resolved by prior approaches (e.g., routing networks or naive mixture-of-experts). This paradigm can be extended to arbitrarily diverse MedIR tasks, given the architecture's modular and scalable adaptation strategy.

Handling Task Imbalance

Sample-level adaptive loss reweighting—conditioned not on task identity but on input-output and reconstruction statistics—enables fine-grained, automated regularization even in highly imbalanced imaging datasets. This directly addresses a common limitation in joint-modality clinical workflows, where different modalities and degradations coexist in the same dataset.

Future Directions

TAT's design is architecture-agnostic; both the weight generation and sample-adaptive balancing modules can be integrated with more recent transformer variants (e.g., sparse/memory-efficient models [yang2024rat, yang2024restore_rwkv, guo2025mambair]), and could facilitate development of unified, real-world-deployable MedIR toolkits. Moreover, insights from adaptive parameterization and loss scaling may inform general multi-task learning outside the MedIR context, including multimodal and multitask scene understanding.

Conclusion

TAT presents a principled, empirically validated contribution to unified medical image restoration by incorporating explicit dynamic specialization in both parameter generation and optimization. The architecture achieves robust SOTA performance across diverse MedIR tasks, and its components—task-adaptive weight generation and loss balancing—are generalizable for future adaptation to novel and more sophisticated network backbones. The work substantiates the value of dynamic modularity in multitask transformers and introduces scalable directions for clinical and research deployment in complex multi-modal imaging environments.