Progressive Test Time Energy Adaptation for Medical Image Segmentation

2025-03-20

AI Summary

Medical image segmentation suffers from performance degradation due to domain shifts across imaging devices and clinical centers; conventional domain adaptation methods require multiple accesses to target data, violating the clinical constraint of single-pass inference. To address this, we propose a test-time progressive energy-based adaptation method that operates without target-domain labels. Our approach introduces, for the first time, a single-forward-pass adaptation mechanism grounded in a shape energy model: it quantifies the plausibility of predicted segmentation shapes at the patch level via an energy score, jointly enabling distribution alignment and erroneous prediction correction. Leveraging gradient-driven lightweight fine-tuning, the model is optimized dynamically during inference. Evaluated on eight public MRI and X-ray datasets covering cardiac, spinal cord, and pulmonary anatomies, our method consistently outperforms state-of-the-art approaches, achieving significant improvements in key metrics—while strictly preserving single-forward-pass inference efficiency.

Technology Category

• Computer Vision: Segmentation
• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Cognitive Modeling & Cognitive Systems: Adaptive Behavior

Application Category

• Search and Retrieval-Augmented AI: Vertical and domain-specific search
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Systems and Infrastructure for Web, Mobile and WoT: Energy management for devices in mobile Web and WoT environments

Abstract

We propose a model-agnostic, progressive test-time energy adaptation approach for medical image segmentation. Maintaining model performance across diverse medical datasets is challenging, as distribution shifts arise from inconsistent imaging protocols and patient variations. Unlike domain adaptation methods that require multiple passes through target data - impractical in clinical settings - our approach adapts pretrained models progressively as they process test data. Our method leverages a shape energy model trained on source data, which assigns an energy score at the patch level to segmentation maps: low energy represents in-distribution (accurate) shapes, while high energy signals out-of-distribution (erroneous) predictions. By minimizing this energy score at test time, we refine the segmentation model to align with the target distribution. To validate the effectiveness and adaptability, we evaluated our framework on eight public MRI (bSSFP, T1- and T2-weighted) and X-ray datasets spanning cardiac, spinal cord, and lung segmentation. We consistently outperform baselines both quantitatively and qualitatively.

Problem

• Adapts pretrained models for medical image segmentation during test time
• Addresses distribution shifts from inconsistent imaging protocols and patient variations
• Improves segmentation accuracy by minimizing shape energy scores

Innovation

• Progressive test-time energy adaptation for segmentation
• Shape energy model minimizes patch-level errors
• Model-agnostic approach refines during test processing