Gradient Alignment Improves Test-Time Adaptation for Medical Image Segmentation

要約

近年、医療画像のセグメンテーションは大幅に進歩していますが、さまざまなセンターからの医療画像間のドメインシフトという蔓延する問題が、事前トレーニングされたモデルの効果的な展開を妨げています。
推論中にテストデータを使用して事前トレーニングされたモデルを微調整することで、この問題に対処するために、多くのテスト時適応 (TTA) 手法が提案されています。
ただし、これらの方法では、次善の最適化方向 (勾配によって決定される) と固定ステップ サイズ (学習率に基づく) により、満足のいく最適化が得られないことがよくあります。
この論文では、最適化手順における勾配の方向と学習率の両方を改善する勾配アライメントベースのテスト時間適応 (GraTa) 方法を提案します。
自己教師付き対物レンズから得られる擬似勾配を主に最適化する従来の TTA 法とは異なり、私たちの方法では、勾配の調整を容易にするために擬似勾配に補助勾配を組み込んでいます。
このような勾配の調整により、モデルは異なる勾配間の類似点を発見し、現在のセグメンテーション タスクに関連する経験的な勾配に近似するように勾配の方向を修正できます。
さらに、擬似勾配と補助勾配の間のコサイン類似性に基づいて動的学習率を設計することで、さまざまなテスト データに対する事前トレーニング済みモデルの適応微調整が可能になります。
広範な実験により、提案された勾配アライメントと動的学習率の有効性が確立され、ベンチマーク医療画像セグメンテーション タスクにおける他の最先端の TTA 手法に対する GraTa 手法の優位性が実証されました。
事前トレーニングされたソース モデルのコードと重みが利用可能になります。

要約(オリジナル)

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 will be available.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google