Gradient Alignment Improves Test-Time Adaptation for Medical Image Segmentation

2024-08-14

arXiv.org

AI Summary

In medical image segmentation, cross-center domain shift severely degrades the generalization of pre-trained models on unseen hospital data; existing test-time adaptation (TTA) methods suffer from suboptimal fine-tuning due to gradient direction bias and fixed learning rates. To address this, we propose GraTa, a gradient-aligned TTA method. GraTa introduces a novel pseudo-gradient–auxiliary-gradient alignment mechanism to jointly rectify optimization directions, and designs a cosine-similarity-based dynamic learning rate strategy to enhance convergence and robustness. Operating entirely without additional annotations, GraTa integrates self-supervised pseudo-labeling into the TTA framework. Extensive experiments across multiple medical image segmentation benchmarks demonstrate that GraTa consistently outperforms state-of-the-art TTA methods. The source code and pre-trained weights are publicly available.

Technology Category

• Machine Learning: Transfer, Domain Adaptation, Multi-Task Learning
• Computer Vision: Segmentation
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Systems and Infrastructure for Web, Mobile and WoT: Applied ML and AI for Web-based mobile applications

Abstract

Although recent years have witnessed significant advancements in medical image segmentation, the pervasive issue of domain shift among medical images from diverse centres hinders the effective deployment of pre-trained models. Many Test-time Adaptation (TTA) methods have been proposed to address this issue by fine-tuning pre-trained models with test data during inference. These methods, however, often suffer from less-satisfactory optimization due to suboptimal optimization direction (dictated by the gradient) and fixed step-size (predicated on the learning rate). In this paper, we propose the Gradient alignment-based Test-time adaptation (GraTa) method to improve both the gradient direction and learning rate in the optimization procedure. Unlike conventional TTA methods, which primarily optimize the pseudo gradient derived from a self-supervised objective, our method incorporates an auxiliary gradient with the pseudo one to facilitate gradient alignment. Such gradient alignment enables the model to excavate the similarities between different gradients and correct the gradient direction to approximate the empirical gradient related to the current segmentation task. Additionally, we design a dynamic learning rate based on the cosine similarity between the pseudo and auxiliary gradients, thereby empowering the adaptive fine-tuning of pre-trained models on diverse test data. Extensive experiments establish the effectiveness of the proposed gradient alignment and dynamic learning rate and substantiate the superiority of our GraTa method over other state-of-the-art TTA methods on a benchmark medical image segmentation task. The code and weights of pre-trained source models are available at https://github.com/Chen-Ziyang/GraTa.

Problem

• Medical Image Segmentation
• Domain Adaptation
• Fine-tuning

Innovation

• GraTa
• Dynamic Learning Rate
• Cross-Hospital Adaptation