[논문 리뷰] MRI Image Reconstruction via Learning Optimization Using Neural ODEs

Eric Z. Chen, Terrence Chen외 1명|arXiv (Cornell University)|2020. 06. 24.

Model Reduction and Neural Networks참고 문헌 39인용 수 6

한 줄 요약

이 논문은 신경 미분방정식(신경 ODE)을 사용하여 이미지 복원의 연속 최적화 궤적을 모델링하는 새로운 MRI 복원 방법인 ReconODE를 제안한다. 최적화 과정의 동역학을 학습하고 ODE 해법기법의 지식을 네트워크 아키텍처에 통합함으로써, UNet 및 캐스케이드 CNN과 같은 기존 딥러닝 모델보다 훨씬 적은 파라미터로 최신 기술 수준의 복원 품질을 달성한다.

ABSTRACT

We propose to formulate MRI image reconstruction as an optimization problem and model the optimization trajectory as a dynamic process using ordinary differential equations (ODEs). We model the dynamics in ODE with a neural network and solve the desired ODE with the off-the-shelf (fixed) solver to obtain reconstructed images. We extend this model and incorporate the knowledge of off-the-shelf ODE solvers into the network design (learned solvers). We investigate several models based on three ODE solvers and compare models with fixed solvers and learned solvers. Our models achieve better reconstruction results and are more parameter efficient than other popular methods such as UNet and cascaded CNN. We introduce a new way of tackling the MRI reconstruction problem by modeling the continuous optimization dynamics using neural ODEs.

연구 동기 및 목표

• 이미지 복원을 연속 최적화 문제로 재정의하여 가속화된 MRI 복원의 과제를 해결한다.
• 기존의 이산 반복 방법과 고정된 해법기법의 한계를 극복하기 위해, 일반 미분방정식(ODE)을 사용해 최적화 과정을 연속 동역학 시스템으로 모델링한다.
• 고정된 해법기법 대신 신경망 내에서 ODE 해법기법의 계수와 스텝 크기를 학습함으로써(학습된 해법기법) 파라미터 효율성과 복원 품질을 향상시킨다.
• 해법기법 지식을 네트워크 설계에 통합할 경우, UNet 및 캐스케이드 CNN과 같은 표준 딥러닝 아키텍처보다 성능 향상과 더 낮은 모델 복잡도를 달성할 수 있음을 입증한다.
• 다양한 ODE 해법기법과 주목적 메커니즘의 영향을 분석하여 복원 정밀도 및 일반화 능력 향상에 기여하는 바를 조사한다.

제안 방법

• 이미지 복원을 ODE의 초기값 문제로 재정의한다: dx(t)/dt = f(x(t), t, y, θ), 여기서 f는 연속 최적화 역학을 모델링하는 신경망이다.
• 사전 구현된 ODE 해법기법(예: 오일러, 룬게-쿠타)을 사용해 초기 과표본화된 이미지에서 복원된 전면 표본화된 이미지로 ODE를 통합한다.
• 네트워크가 ODE 해법기법의 계수와 스텝 크기를 학습하도록 하여 '학습된 해법기법'(ReconODE-LT)을 도입함으로써, 해법기법을 적응형이고 학습 가능한 것으로 만든다.
• 주목적 메커니즘을 신경 ODE 아키텍처에 통합하여 특징 표현을 향상시키고 복원 정확도를 높인다.
• 역전파를 ODE 해법기법을 통해(ReconODE-FT) 또는 고유한 방법을 통해(ReconODE-FA) 수행하고, 메모리 사용량을 줄이기 위해 기울기 체크포인팅을 사용한다.
• 다양한 ODE 해법기법(오일러, RK2, RK4)과 데이터셋에서 고정된 해법기법(ReconODE-FT, ReconODE-FA)과 학습된 해법기법(ReconODE-LT) 간의 성능을 비교한다.

실험 결과

연구 질문

• RQ1 신경 ODE를 사용해 MRI 복원을 연속 최적화 문제로 모델링할 경우, 이산 반복 방법에 비해 복원 품질이 향상되는가?
• RQ2 ODE 해법기법의 계수와 스텝 크기를 학습하는(학습된 해법기법) 방식이 고정된 사전 정의된 해법기법을 사용하는 것보다 복원 성능을 크게 향상시키는가?
• RQ3 ODE 해법기법 지식을 네트워크 아키텍처에 통합할 경우, 파라미터 효율성과 복원 정밀도에 어떤 영향을 미치는가?
• RQ4 주목적 메커니즘이 신경 ODE 기반 MRI 복원 모델의 성능에 어떤 영향을 미치는가?
• RQ5 기울기 체크포인팅 기법이 신경 ODE 기반 MRI 복원에서 학습 시 메모리 사용량을 효과적으로 줄일 수 있으며, 추론 성능에 영향을 주지 않는가?

주요 결과

• ReconODE-LT-RK4는 fastMRI 검증 데이터셋에서 4배 가속화 시 32.39의 PSNR, 8배 가속화 시 30.27의 PSNR, 4배 가속화 시 0.7333의 SSIM, 8배 가속화 시 0.6617의 SSIM을 기록하여 UNet을 포함한 모든 다른 모델보다 뛰어난 성능을 보였다.
• ReconODE-LT-Euler는 UNet의 파라미터의 0.9%인 0.03M(30,000)에 불과한 파라미터로 4배 가속화 시 UNet과 유사한 SSIM 성능을 달성하여 뛰어난 파라미터 효율성을 입증했다.
• ReconODE-LT-RK2는 D5C5 모델의 64%의 파라미터로 4배 및 8배 가속화 시 D5C5와 유사한 SSIM과 PSNR 성능을 기록하여 뛰어난 파라미터 효율성을 보였다.
• ReconODE-FT-RK4에 주목적 메커니즘을 통합하면 4배 가속화 시 SSIM이 주목적 없이 0.726에서 0.733으로 향상되어 복원 품질 향상에 긍정적 영향을 미친다.
• 역전파를 해법기법을 통해 수행하는 ReconODE-FT 모델이 고유한 방법을 사용하는 ReconODE-FA 모델보다 성능이 뛰어나, 학습 시 더 정확한 기울기 계산이 가능함을 시사한다.
• 기울기 체크포인팅은 ReconODE-LT-RK4의 GPU 메모리 사용량을 73% 감소시켰다(13.6G에서 3.7G로), 학습 시간은 0.13초 증가에 그치며 추론 성능에 영향을 주지 않았다.