[논문 리뷰] On-the-Fly Test-time Adaptation for Medical Image Segmentation

Jeya Maria Jose Valanarasu, Guo, Pengfei외 2명|arXiv (Cornell University)|2022. 03. 10.

Domain Adaptation and Few-Shot Learning인용 수 4

한 줄 요약

이 논문은 백프로파게이션 없이 단일 테스트 이미지에 대해 사전 학습된 모델을 적응시키는 제로샷, 에피소딕 프레임워크인 On-the-Fly Test-Time Adaptation를 제안한다. 도메인 코드 기반 적응형 배치 정규화를 갖춘 적응형 UNet을 사용하여 2D 및 3D 도메인 이동 상황에서 최신 기술 성능을 달성하며, 더 엄격한 임상적 제약 조건에도 불구하고 기존의 테스트 시간 적응 방법들을 능가한다.

ABSTRACT

One major problem in deep learning-based solutions for medical imaging is the drop in performance when a model is tested on a data distribution different from the one that it is trained on. Adapting the source model to target data distribution at test-time is an efficient solution for the data-shift problem. Previous methods solve this by adapting the model to target distribution by using techniques like entropy minimization or regularization. In these methods, the models are still updated by back-propagation using an unsupervised loss on complete test data distribution. In real-world clinical settings, it makes more sense to adapt a model to a new test image on-the-fly and avoid model update during inference due to privacy concerns and lack of computing resource at deployment. To this end, we propose a new setting - On-the-Fly Adaptation which is zero-shot and episodic (i.e., the model is adapted to a single image at a time and also does not perform any back-propagation during test-time). To achieve this, we propose a new framework called Adaptive UNet where each convolutional block is equipped with an adaptive batch normalization layer to adapt the features with respect to a domain code. The domain code is generated using a pre-trained encoder trained on a large corpus of medical images. During test-time, the model takes in just the new test image and generates a domain code to adapt the features of source model according to the test data. We validate the performance on both 2D and 3D data distribution shifts where we get a better performance compared to previous test-time adaptation methods. Code is available at https://github.com/jeya-maria-jose/On-The-Fly-Adaptation

연구 동기 및 목표

• 학습 데이터와 테스트 데이터 간의 도메인 이동으로 인해 의료 영상 분할 성능이 떨어지는 문제를 해결하기 위해.
• 백프로파게이션을 피하고 테스트 시 전체 타겟 분포에 대한 액세스가 필요 없는 임상적으로 구현 가능한 적응 방법을 개발하기 위해.
• 개별 이미지당 제로샷, 온더플라이 적응을 가능하게 하여 실제 임상 환경의 제약 조건(예: 개인정보 보호 및 제한된 컴퓨팅 자원)과 부합하기 위해.
• 사전 학습된 인코더에서 유도된 도메인 코드를 사용해 모델 특징을 동적으로 적응시키는 프레임워크를 설계하기 위해.
• 중대한 도메인 이동 상황에서 다양한 2D 및 3D 의료 영상 분할 벤치마크에서 방법을 검증하기 위해.

제안 방법

• 테스트 시간 적응의 새로운 설정을 제안: On-the-Fly Test-Time Adaptation로, 적응은 에피소딕이며 제로샷이며 추론 중 모델 업데이트가 없음.
• 표준 UNet에 도메인 코드를 사용해 특징을 적응시키는 적응형 배치 정규화 레이어를 추가한 Adaptive UNet을 도입.
• 대규모 의료 영상 코퍼스에 대해 사전 학습된 인코더를 사용해 각 테스트 이미지의 도메인 코드를 생성하여 모ality 및 스캐너 특성과 같은 특징을 캡처.
• 도메인 코드를 사용해 인코더 및 디코더의 배치 정규화 레이어의 누적 통계를 조절하여 인스턴스별 특징 적응을 가능하게 함.
• 원천 데이터나 테스트 시간 학습이 필요 없이도, 사전 학습된 모델에서 유도된 도메인 사전 지식을 활용해 적응을 안내함.
• 각각의 이미지에 대해 한 번씩 적응된 모델을 적용하고, 추론 후 적응 상태를 초기화함으로써 개인정보 보호 및 자원 소모를 최소화함.

실험 결과

연구 질문

• RQ1 추론 중 백프로파게이션 없이도 성능이 높은 테스트 시간 적응 방법이 가능할 수 있는가?
• RQ2 전체 타겟 분포에 액세스할 수 없는 상황에서, 이미지당 제로샷, 온더플라이 적응은 얼마나 효과적인가?
• RQ3 학습된 도메인 코드를 활용한 적응형 배치 정규화는 2D 및 3D 의료 영상에서의 도메인 이동을 효과적으로 다룰 수 있는가?
• RQ4 더 엄격한 임상 구현 조건 하에서, 제안된 방법은 기존의 테스트 시간 적응 기반 방법들을 능가하는가?
• RQ5 예를 들어, MRI 시퀀스 간 또는 펌터스 영상 장치 간의 다양한 모달리티 이동 상황에서 방법의 일반화 능력은 어떠한가?

주요 결과

• Adaptation UNet을 사용한 On-the-Fly Adaptation은 2D 펌터스 영상 이동에서 63.14의 Dice 스코어를 기록하며, TENT(57.10)과 RN+CR(59.15)를 능가한다.
• 3D MRI 분할에 대해선 T1→T1ce 이동에서 각각 60.66/58.73/39.30의 Dice 스코어를 기록하며, TENT(55.47/52.89/39.17)와 RN+CR(57.30/54.68/40.01)를 초월한다.
• 2D 및 3D 환경에서의 9개의 도메인 이동 상황에서 직접 테스트 및 모든 기반 방법보다 일관된 성능 향상을 보였다.
• 정성적 결과에서는 Adaptive UNet이 지식 기반 예측 및 진짜 마스크에 더 가까운 분할 마스크를 생성하며, 과다 분할 및 오분류를 감소시켰다.
• 2D 이동에 비해 3D 이동에서 성능가 다소 낮지만, 이는 교차 모달리티 MRI 시퀀스 이동의 더 큰 복잡성 때문일 것으로 보이며, 여전히 뛰어난 성능를 기록했다.
• 절단 분석 결과 도메인 코드 기반 적응이 핵심임을 확인했으며, 이를 제거할 경우 비적응 모델과 유사한 성능 저하가 발생했다.