Sample-Aware Test-Time Adaptation for Medical Image-to-Image Translation

2025-08-01

AI Summary

Medical image cross-domain translation suffers significant performance degradation on out-of-distribution (OOD) samples. To address this, we propose a test-time adaptation (TTA) framework that quantifies sample-level domain shift via reconstruction error and introduces a sample-aware dynamic adaptation module. This module selectively adjusts intermediate features of a pre-trained model only for OOD samples, thereby avoiding redundant interventions on in-distribution (ID) samples. Unlike conventional TTA methods that uniformly adapt all test samples, our approach enables fine-grained, lightweight online optimization. Evaluated on low-dose CT denoising and T1-to-T2 MRI synthesis, the method consistently outperforms both non-TTA baselines and state-of-the-art TTA approaches, achieving PSNR/SSIM improvements of 1.2–2.8 dB and 0.015–0.032, respectively. Moreover, it enhances model robustness and inference stability under distribution shifts.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Diffusion Models for Vision
• Application Domains: Transportation

Application Category

• User Modeling, Personalization and Recommendation: User modeling for targeted and personalized online advertising
• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Image-to-image translation has emerged as a powerful technique in medical imaging, enabling tasks such as image denoising and cross-modality conversion. However, it suffers from limitations in handling out-of-distribution samples without causing performance degradation. To address this limitation, we propose a novel Test-Time Adaptation (TTA) framework that dynamically adjusts the translation process based on the characteristics of each test sample. Our method introduces a Reconstruction Module to quantify the domain shift and a Dynamic Adaptation Block that selectively modifies the internal features of a pretrained translation model to mitigate the shift without compromising the performance on in-distribution samples that do not require adaptation. We evaluate our approach on two medical image-to-image translation tasks: low-dose CT denoising and T1 to T2 MRI translation, showing consistent improvements over both the baseline translation model without TTA and prior TTA methods. Our analysis highlights the limitations of the state-of-the-art that uniformly apply the adaptation to both out-of-distribution and in-distribution samples, demonstrating that dynamic, sample-specific adjustment offers a promising path to improve model resilience in real-world scenarios. The code is available at: https://github.com/cosbidev/Sample-Aware_TTA.

Problem

• Handles out-of-distribution samples in medical image translation
• Dynamically adjusts translation based on test sample characteristics
• Improves model resilience without compromising in-distribution performance

Innovation

• Dynamic Test-Time Adaptation for medical images
• Reconstruction Module quantifies domain shift
• Selective feature modification in pretrained model