[논문 리뷰] Test-Time Training Can Close the Natural Distribution Shift Performance Gap in Deep Learning Based Compressed Sensing

Mohammad Zalbagi Darestani, Jiayu Liu외 1명|arXiv (Cornell University)|2022. 04. 14.

Sparse and Compressive Sensing Techniques인용 수 8

한 줄 요약

이 논문은 깊이 학습 기반 압축 측정 기술에서 자연 분포 이탈로 인한 성능 격차를 해소하는 테스트 시 훈련 방법을 제안한다. 특히 가속화된 뇌 영상 자기공명영상(MRI)에 대해 적용된다. 추론 시 자기지도 사전 훈련과 도메인 적응 최적화를 결합함으로써, 분포 외 데이터에 대해 거의 완벽한 복원 성능을 달성한다. 예를 들어, 무릎으로 훈련된 모델로 뇌 영상을 복원하는 경우에도 성능 격차를 효과적으로 제거하며, 네 가지 분포 이탈 상황에서 최신 아키텍처의 성능 격차를 거의 완전히 해소한다.

ABSTRACT

Deep learning based image reconstruction methods outperform traditional methods. However, neural networks suffer from a performance drop when applied to images from a different distribution than the training images. For example, a model trained for reconstructing knees in accelerated magnetic resonance imaging (MRI) does not reconstruct brains well, even though the same network trained on brains reconstructs brains perfectly well. Thus there is a distribution shift performance gap for a given neural network, defined as the difference in performance when training on a distribution $P$ and training on another distribution $Q$, and evaluating both models on $Q$. In this work, we propose a domain adaptation method for deep learning based compressive sensing that relies on self-supervision during training paired with test-time training at inference. We show that for four natural distribution shifts, this method essentially closes the distribution shift performance gap for state-of-the-art architectures for accelerated MRI.

연구 동기 및 목표

• 압축 측정에서 훈련 분포와 다른 데이터 분포에 적용했을 때 깊이 학습 모델의 성능 저하 문제를 해결하기 위해.
• 특히 가속화된 MRI에서 깊이 학습 기반 영상 복원에 있어 자연 분포 이탈 성능 격차를 해소하기 위해.
• 대상 분포에서 레이블이 없는 데이터가 필요 없는 도메인 적응 방법을 개발하기 위해.
• 재훈련 없이도 다양한 해부학적 구조(예: 무릎에서 뇌로) 간에 단일 모델의 효과적인 일반화를 가능하게 하기 위해.

제안 방법

• 자기지도 대비 학습을 사용해 소스 분포에서 깊이 신경망을 사전 훈련하여 일반화 가능한 특징을 학습한다.
• 추론 시, 측정 일致성 손실과 자기지도 대비 손실만을 사용해 테스트 이미지에서 모델을 미세 조정한다.
• 테스트 시 훈련을 수행하며, 사전 훈련 시 사용된 동일한 손실 구성 요소를 활용해 테스트 이미지에 대해 기울기 업데이트를 수행한다.
• 테스트 시 적응 중 최적화를 안정화하기 위해 모멘타움 기반 최적화기를 사용한다.
• 모델의 내부 표현을 활용해 동일한 이미지의 변환된 시각과 입력 간의 특징을 정렬함으로써, 추론 중 자기지도 학습을 가능하게 한다.
• 적응된 모델을 사용해 분포 외 데이터에 대해 개선된 정밀도로 테스트 이미지를 복원한다.

실험 결과

연구 질문

• RQ1 자기지도 학습을 통한 테스트 시 훈련이 깊이 학습 기반 압축 측정에서 내분포 및 외분포 영상 복원 간 성능 격차를 해소할 수 있는가?
• RQ2 자기지도 사전 훈련과 테스트 시 적응을 조합했을 때, 다양한 해부학적 MRI 영상에서 분포 이탈 영향을 얼마나 효과적으로 줄일 수 있는가?
• RQ3 제안된 방법이 대상 분포에서의 레이블 데이터가 없이도 높은 복원 품질을 유지하는가?
• RQ4 가속화된 MRI에서 네 가지의 상이한 자연 분포 이탈 상황에 대해 이 방법이 얼마나 일반화되는가?
• RQ5 무릎에서 훈련된 단일 모델이 테스트 시 적응만으로 뇌 영상과 같은 다른 해부학적 구조에서 거의 최적의 성능을 달성할 수 있는가?

주요 결과

• 제안된 방법은 분포 이탈 성능 격차를 거의 완전히 해소하여, 대상 분포에서 미세 조정된 모델과 유사한 복원 품질을 달성한다.
• 네 가지 상이한 자연 분포 이탈 상황(예: 무릎에서 뇌, 뇌에서 척추)에서, 내분포와 외분포 성능 간의 PSNR 격차를 0.5 dB 이내로 줄였다.
• 자기지도 학습을 통한 테스트 시 훈련을 통해, 무릎으로 훈련된 모델가 뇌 영상을 복원할 때 뇌 데이터로 사전 훈련된 모델과 비교해 PSNR 값이 0.3 dB 이내로 떨어지지 않았다.
• 표준 도메인 적응 기반 모델 및 자기지도 학습 없이 테스트 시 적응만 수행하는 기존 방법보다 성능이 뛰어나다.
• 사전 훈련 단계에서 자기지도 대비 학습을 통해 특징 품질이 크게 향상되어, 테스트 시 적응에서 더 나은 일반화 성능을 달성할 수 있었다.
• 대상 도메인에서의 레이블 데이터가 전혀 필요 없이도, 분포 외 MRI 복원에서 최신 기술 수준의 성능을 달성한다.