[論文レビュー] MRI Image Reconstruction via Learning Optimization Using Neural ODEs

Eric Z. Chen, Terrence Chen+1 more|arXiv (Cornell University)|Jun 24, 2020

Model Reduction and Neural Networks参考文献 39被引用数 6

ひとこと要約

この論文では、画像再構成の連続的最適化軌道をニューラル常微分方程式（Neural ODEs）を用いてモデル化する、新しい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ソルバ（例：Euler、Runge-Kutta）を用いて、初期のアンダーサンプリング画像から完全サンプリング画像へのODEの統合を行う。
• ネットワークがODEソルバの係数とステップサイズを学習可能にする「学習済みソルバ」（ReconODE-LT）を導入することで、ソルバを適応的かつトレーナブルにする。
• 注目メカニズムをニューラルODEアーキテクチャに組み込み、特徴表現を強化し、再構成精度を向上させる。
• 勾配チェックポイントを用いてメモリ使用量を削減しながら、バックプロパゲーションによるODEソルバの逆伝播（ReconODE-FT）または随伴法（ReconODE-FA）を用いてモデルを訓練する。
• 複数のODEソルバ（Euler、RK2、RK4）とデータセットを用いて、固定ソルバ（ReconODE-FT、ReconODE-FA）と学習済みソルバ（ReconODE-LT）を比較する。

実験結果

リサーチクエスチョン

• RQ1 Neural ODEを用いてMRI再構成を連続的最適化問題として定式化することで、離散的反復手法と比較して再構成品質が向上するか？
• RQ2 ODEソルバの係数とステップサイズを学習する（＝学習済みソルバ）ことで、固定で事前に定義されたソルバを使用する場合と比較して、再構成性能が顕著に向上するか？
• RQ3 ODEソルバの知識をネットワークアーキテクチャに統合することで、パラメータ効率と再構成忠実度にどのような影響を与えるか？
• RQ4 注目メカニズムは、Neural ODEに基づくMRI再構成モデルの性能にどのような影響を与えるか？
• RQ5 勾配チェックポイント技術は、Neural ODEに基づくMRI再構成におけるトレーニング時のメモリ使用量を効果的に削減できるか、かつ推論性能に悪影響を与えないか？

主な発見

• ReconODE-LT-RK4は、fastMRIバリデーションデータセットで最高のPSNR（4X時32.39、8X時30.27）とSSIM（4X時0.7333、8X時0.6617）を達成し、UNetを含むすべての他のモデルを上回った。
• ReconODE-LT-EulerはUNetのパラメータのたった0.9%（0.03M vs. 3.35M）で4X時にUNetと同等のSSIMを達成し、極めて高いパラメータ効率を示した。
• ReconODE-LT-RK2はD5C5のパラメータの64%で、4Xおよび8Xの両方でD5C5と同等のSSIMとPSNRを達成し、優れたパラメータ効率を示した。
• ReconODE-FT-RK4に注目メカニズムを組み込むことで、4X時のSSIMが0.726（注目なし）から0.733に向上し、再構成品質への肯定的影響が確認された。
• ReconODE-FTモデル（ソルバを介したバックプロパゲーション）はReconODE-FAモデル（随伴法）を上回り、トレーニングにおける勾配計算の正確性がより高いことを示唆した。
• 勾配チェックポイントにより、ReconODE-LT-RK4のGPUメモリ使用量が73%削減（13.6Gから3.7G）され、トレーニング時間は0.13秒増加するにとどまり、推論性能に影響を与えないことが確認された。