[논문 리뷰] Learning to solve inverse problems using Wasserstein loss

Jonas Adler, Axel Ringh외 2명|arXiv (Cornell University)|2017. 10. 30.

Medical Imaging Techniques and Applications참고 문헌 33인용 수 25

한 줄 요약

이 논문은 복소단층촬영(CT)을 포함한 역문제에서 깊이 학습 모델을 훈련시키기 위해 평균제곱오차(MSE) 대신 워샤르스타인 손실을 사용하는 것을 제안한다. 최적 운반 이론을 활용함으로써, 훈련 데이터의 공간적 비일치에 대해 더 강건한 성능을 발휘하며, 이러한 왜곡 상황에서 MSE 기반 훈련이 흐릿한 결과를 낳는 것에 비해 더 선명한 재구성을 얻을 수 있다.

ABSTRACT

We propose using the Wasserstein loss for training in inverse problems. In particular, we consider a learned primal-dual reconstruction scheme for ill-posed inverse problems using the Wasserstein distance as loss function in the learning. This is motivated by miss-alignments in training data, which when using standard mean squared error loss could severely degrade reconstruction quality. We prove that training with the Wasserstein loss gives a reconstruction operator that correctly compensates for miss-alignments in certain cases, whereas training with the mean squared error gives a smeared reconstruction. Moreover, we demonstrate these effects by training a reconstruction algorithm using both mean squared error and optimal transport loss for a problem in computerized tomography.

연구 동기 및 목표

• 역문제의 재구성 품질 저하를 해결하기 위해 훈련 데이터의 공간적 비일치에 기인한 영향을 다루는 것.
• 훈련 데이터의 기하학적 왜곡에 강건한 학습된 프리미벌-듀얼 재구성 프레임워크를 개발하는 것.
• 비일치 조건 하에서 워샤르스타인 손실이 평균제곱오차(MSE)보다 더 나은 재구성 품질을 제공함을 보여주는 것.
• 워샤르스타인 손실이 데이터 비일치를 보완하는 데서 우월한 이유를 이론적으로 설명하는 것.

제안 방법

• 역문제를 위한 학습된 프리미벌-듀얼 재구성 네트워크에서 워샤르스타인 거리(최적 운반 비용)를 손실 함수로 사용한다.
• 노이즈가 섞이고 간접적인 측정값에서 재구성된 이미지를 도출하는 딥 네트워크를 워샤르스타인 손실로 훈련시킨다.
• 전방 연산자와 그 수반 연산자가 네트워크 아키텍처 내부에 통합된 학습된 반복적 방법을 적용한다.
• 훈련 중 워샤르스타인 손실을 계산하기 위해 효율적인 최적 운반 해법을 사용한다.
• 복소단층촬영 재구성 작업에서 워샤르스타인 손실 훈련을 표준 평균제곱오차(MSE) 손실과 비교한다.
• 이론적 분석을 통해 워샤르스타인 손실은 특정 경우에 비일치를 정확히 보완하지만, MSE는 흐릿한 결과를 낳는다는 것을 증명한다.

실험 결과

연구 질문

• RQ1 훈련 데이터에 공간적 비일치가 존재할 경우 워샤르스타인 손실을 사용하면 재구성 품질이 향상되는가?
• RQ2 역문제 훈련에서 기하학적 왜곡을 다룰 때 워샤르스타인 손실은 평균제곱오차(MSE)와 어떻게 비교되는가?
• RQ3 워샤르스타인 손실은 데이터 비일치를 정확히 보완하는 재구성 연산자를 도출할 수 있는가?
• RQ4 불안정한 역문제의 맥락에서 워샤르스타인 손실은 어떤 이론적 성질을 제공하는가?

주요 결과

• 워샤르스타인 손실로 훈련하면 훈련 데이터의 비일치를 정확히 보완하는 재구성 연산자가 도출되며, 반면 MSE 훈련은 비일치로 인해 흐릿한 재구성을 낳는다.
• 이론적 분석을 통해 워샤르스타인 손실은 기대 운반 비용을 최소화하여, 비일치 상황에서 진정한 기저 분포와 더 잘 일치하는 결과를 이끈다.
• 복소단층촬영 실험에서 워샤르스타인로 훈련된 모델은 비일치된 데이터 하에서 MSE로 훈련된 모델보다 훨씬 선명한 재구성을 생성했다.
• 워샤르스타인 손실은 MRI에서 자주 발생하는 자기장 불균일성이나 CT에서의 빔 아티팩트와 같은 기하학적 정확도 저하에 더 강건하다.
• 진정한 이미지가 최대 4 mm까지 왜곡되어도 워샤르스타인 손실 기반 방법은 여전히 뛰어난 성능을 보였다. 이는 MRI에서 알려진 문제이다.
• 이론적 결과는 워샤르스타인 손실 하에서 최적 해가 비일치가 대칭일 경우 유일하고 안정적이며, 반면 MSE는 이러한 보장이 없음을 확인했다.