[논문 리뷰] Unsupervised Deep-Learning Based Deformable Image Registration: A Bayesian Framework

Samah Khawaled, Moti Freiman|arXiv (Cornell University)|2020. 08. 10.

Generative Adversarial Networks and Image Synthesis참고 문헌 28인용 수 5

한 줄 요약

이 논문은 변형 가능한 영상 정렬을 위한 완전 베이지안 비지도 딥 러닝 프레임워크를 제안하며, 네트워크 가중치의 진정한 사후 분포를 샘플링하기 위해 확률적 경량화 랭그비안 동역학(SGLD)을 사용하여 과적합을 방지하고 불확실성 정량화를 가능하게 한다. 이 방법은 VoxelMorph보다 정확도를 향상시켜, MNIST 및 MGH10 데이터셋에서 각각 더 낮은 MSE(0.0063 vs. 0.0065)와 더 높은 Dice 점수(0.73 vs. 0.71)를 달성한다.

ABSTRACT

Unsupervised deep-learning (DL) models were recently proposed for deformable image registration tasks. In such models, a neural-network is trained to predict the best deformation field by minimizing some dissimilarity function between the moving and the target images. After training on a dataset without reference deformation fields available, such a model can be used to rapidly predict the deformation field between newly seen moving and target images. Currently, the training process effectively provides a point-estimate of the network weights rather than characterizing their entire posterior distribution. This may result in a potential over-fitting which may yield sub-optimal results at inference phase, especially for small-size datasets, frequently present in the medical imaging domain. We introduce a fully Bayesian framework for unsupervised DL-based deformable image registration. Our method provides a principled way to characterize the true posterior distribution, thus, avoiding potential over-fitting. We used stochastic gradient Langevin dynamics (SGLD) to conduct the posterior sampling, which is both theoretically well-founded and computationally efficient. We demonstrated the added-value of our Basyesian unsupervised DL-based registration framework on the MNIST and brain MRI (MGH10) datasets in comparison to the VoxelMorph unsupervised DL-based image registration framework. Our experiments show that our approach provided better estimates of the deformation field by means of improved mean-squared-error ($0.0063$ vs. $0.0065$) and Dice coefficient ($0.73$ vs. $0.71$) for the MNIST and the MGH10 datasets respectively. Further, our approach provides an estimate of the uncertainty in the deformation-field by characterizing the true posterior distribution.

연구 동기 및 목표

• 작은 의료 영상 데이터셋에서 비지도 딥 러닝 기반 변형 가능한 영상 정렬의 과적합 문제를 해결하기 위해.
• 이전의 베이지안 방법에서 사용하는 정규성 가정을 피하기 위해 네트워크 가중치의 사후 분포에 대한 원칙적이고 비모수적 특성화를 제공하기 위해.
• 진정한 사후 분포를 샘플링하여 변형 필드 예측의 불확실성 정량화를 가능하게 하기 위해.
• 적응형 노이즈 주입을 통한 사후 샘플링을 통해 노이즈가 있는 입력 조건에서도 정확도와 강인성을 향상시키기 위해.

제안 방법

• 프레임워크는 네트워크 가중치의 사후 샘플링을 수행하기 위해 확률적 경량화 랭그비안 동역학(SGLD)을 사용하며, 훈련 중 손실 기울기에 적응형 가우시안 노이즈를 주입한다.
• 변형 필드 예측을 위한 배경으로 2D VoxelMorph 유사 컨볼루션 신경망을 사용한다.
• 버너-인 단계 이후 사후 샘플을 수집하고, 이들의 평균을 취해 최종 변형 필드 추정치를 도출함으로써 강인성을 향상시킨다.
• 이 방법은 사후 분포에 대해 파rametric 형태를 가정하지 않아 복잡한 비정규 분포를 탄력적으로 모델링할 수 있다.
• 사후 샘플링 동안 수렴성과 일반화 성능을 향상시키기 위해 기울기 업데이트에 적응형 노이즈 분산을 사용한다.
• 2D 컨볼루션을 3D 컨볼루션으로 교체함으로써 3D 정렬로 확장 가능하며, 통합 레이어를 통해 디퍼모르픽 정렬로도 적용 가능하다.

실험 결과

연구 질문

• RQ1 완전 베이지안 딥 러닝 프레임워크는 점 추정 모델에 비해 비지도 변형 가능한 영상 정렬의 정확도를 향상시킬 수 있는가?
• RQ2 적응형 노이즈 주입을 통한 SGLD를 통한 사후 샘플링은 작은 의료 영상 데이터셋에서 과적합을 줄이고 강인성을 향상시키는가?
• RQ3 제안된 방법은 사후 분포의 비모수적 형태를 가정하지 않고 변형 필드 예측의 불확실성 정량화를 수행할 수 있는가?
• RQ4 노이즈가 다양한 수준으로 포함된 입력 영상에서 베이지안 방법의 성능은 VoxelMorph에 비해 어떻게 비교되는가?

주요 결과

• 제안된 ASGLD Reg 방법은 MNIST 데이터셋에서 평균 제곱 오차(MSE) 0.0063을 기록하여 VoxelMorph의 MSE 0.0065를 초월했다.
• MGH10 뇌 MRI 데이터셋에서 이 방법은 Dice 점수 0.736(±0.017)를 달성했으며, VoxelMorph의 0.7109(±0.027)보다 유의미하게 높았다.
• 모든 테스트 노이즈 수준에서 성능 향상은 통계적으로 유의미했다(p < 0.05 for MNIST, p < 10^-5 for MGH10).
• 사후 샘플의 평균을 취하는 것이 최종 모델 가중치만을 사용하는 것보다 노이즈가 많은 조건에서 강인성을 향상시켰다.
• 정규성 가정 없이 진정한 사후 분포를 특성화함으로써 변형 필드의 불확실성을 성공적으로 정량화했다.
• 비모수적 사후 샘플링 덕분에 특히 작은 데이터셋에서 일반화 능력 향상과 과적합 감소를 보였다.