[論文レビュー] Unsupervised Deep-Learning Based Deformable Image Registration: A Bayesian Framework

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

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）を採用する。
• 変形場予測のためのバックボーンとして、2次元のVoxelMorphに類似した畳み込みニューラルネットワークを用いる。
• バーニングインフェーズを経て事後分布のサンプルを収集し、最終的な変形場推定値はこれらのサンプルの平均を取ることで得られる。
• 事後分布にパラメトリックな形を仮定しないため、複雑で非ガウス型の重み分布の柔軟なモデリングが可能である。
• 事後分布サンプリング中の収束性と一般化性能の向上を図るため、勾配更新に適応的ノイズ分散を用いる。
• 2次元畳み込みを3次元畳み込みに置き換えることで3次元登録に拡張可能であり、統合層を用いることで微分同相登録にも対応可能である。

実験結果

リサーチクエスチョン

• 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）を有意に上回った。
• 全テストノイズレベルにおいて、性能向上は統計的に有意であった（MNISTではp < 0.05、MGH10ではp < 10^-5）。
• 事後分布サンプルの平均化により、特にノイズの多い条件下でも耐障害性が向上し、最終モデル重みのみを用いる場合と比較して優れた性能を示した。
• 正規性の仮定をせず、真の事後分布を特徴付けることで、変形場における不確実性の定量化に成功した。
• 非パラメトリックな事後分布サンプリングのおかげで、特に小規模データセットにおいて一般化性能が向上し、過学習が軽減された。