# Gradient Alignment Improves Test-Time Adaptation for Medical Image Segmentation

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

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 31 references that are in the map.
- First posted on arXiv 14 Aug 2024.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Learning to Generalize: Meta-Learning for Domain Generalization](https://mlsanity.com/paper/1710.03463/) (2017)
- [Towards Stable Test-Time Adaptation in Dynamic Wild World](https://mlsanity.com/paper/2302.12400/) (2023)
- [MEMO: Test Time Robustness via Adaptation and Augmentation](https://mlsanity.com/paper/2110.09506/) (2021)
- [A Comprehensive Survey on Test-Time Adaptation under Distribution Shifts](https://mlsanity.com/paper/2303.15361/) (2023)
- [Embracing Imperfect Datasets: A Review of Deep Learning Solutions for Medical Image Segmentation](https://mlsanity.com/paper/1908.10454/) (2019)
- [Evaluating Prediction-Time Batch Normalization for Robustness under Covariate Shift](https://mlsanity.com/paper/2006.10963/) (2020)
- [REFUGE Challenge: A Unified Framework for Evaluating Automated Methods for Glaucoma Assessment from Fundus Photographs](https://mlsanity.com/paper/1910.03667/) (2019)
- [The Norm Must Go On: Dynamic Unsupervised Domain Adaptation by Normalization](https://mlsanity.com/paper/2112.00463/) (2021)
- [Sharpness-Aware Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2303.10353/) (2023)
- [Dynamic Weighted Learning for Unsupervised Domain Adaptation](https://mlsanity.com/paper/2103.13814/) (2021)
- [Learning Anchored Unsigned Distance Functions with Gradient Direction Alignment for Single-view Garment Reconstruction](https://mlsanity.com/paper/2108.08478/) (2021)
- [TeST: Test-time Self-Training under Distribution Shift](https://mlsanity.com/paper/2209.11459/) (2022)
- [DLTTA: Dynamic Learning Rate for Test-time Adaptation on Cross-domain Medical Images](https://mlsanity.com/paper/2205.13723/) (2022)
- [DomainAdaptor: A Novel Approach to Test-time Adaptation](https://mlsanity.com/paper/2308.10297/) (2023)
- [MedBN: Robust Test-Time Adaptation against Malicious Test Samples](https://mlsanity.com/paper/2403.19326/) (2024)
- [Towards More Robust Interpretation via Local Gradient Alignment](https://mlsanity.com/paper/2211.15900/) (2022)

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

- [Test-time Adaptation for Foundation Medical Segmentation Model without Parametric Updates](https://mlsanity.com/paper/2504.02008/) (2025)
- [A Large Scale Benchmark for Test Time Adaptation Methods in Medical Image Segmentation](https://mlsanity.com/paper/2512.02497/) (2025)
- [SAM-aware Test-time Adaptation for Universal Medical Image Segmentation](https://mlsanity.com/paper/2506.05221/) (2025)
- [A3-TTA: Adaptive Anchor Alignment Test-Time Adaptation for Image Segmentation](https://mlsanity.com/paper/2602.03292/) (2026)
- [SPEGC: Continual Test-Time Adaptation via Semantic-Prompt-Enhanced Graph Clustering for Medical Image Segmentation](https://mlsanity.com/paper/2603.11492/) (2026)
- [DP-TTA: Test-time Adaptation for Transient Electromagnetic Signal Denoising via Dictionary-driven Prior Regularization](https://mlsanity.com/paper/2510.13160/) (2025)
- [EviATTA: Evidential Active Test-Time Adaptation for Medical Segment Anything Models](https://mlsanity.com/paper/2603.14666/) (2026)
- [When Can Test-Time Adaptation Help Zero-Shot CT Vision-Language Models?](https://mlsanity.com/paper/2607.15556/) (2026)
- [Learning Prostate Anatomy at Test Time for Cancer Detection in Micro-Ultrasound](https://mlsanity.com/paper/2608.20557/) (2026)
- [On the Interaction Between Model Compression and Test-Time Adaptation](https://mlsanity.com/paper/2609.03604/) (2026)
- [Should This Case Be Adapted? Prediction Fragmentation Controls Test-Time Adaptation](https://mlsanity.com/paper/2609.20700/) (2026)

## Abstract (arXiv)

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.