[论文解读] 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）的全贝叶斯无监督深度学习框架，用于可变形图像配准，通过采样网络权重的真实后验分布来避免过拟合并实现不确定性量化。该方法在MNIST和MGH10数据集上的配准精度优于VoxelMorph，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）对网络权重进行后验采样，在训练过程中向损失梯度中注入自适应高斯噪声。
• 该方法使用类似VoxelMorph的2D卷积神经网络作为形变场预测的主干网络。
• 在预 burn-in 阶段后收集后验样本，最终通过平均这些样本获得形变场估计，以提升鲁棒性。
• 该方法不假设后验分布的参数形式，从而能够灵活建模复杂且非高斯的权重分布。
• 在梯度更新中使用自适应噪声方差，以增强后验采样过程中的收敛性与泛化能力。
• 通过将2D卷积替换为3D卷积，该方法可扩展至3D配准；通过引入积分层，亦可实现微分同胚配准。

实验结果

研究问题

• RQ1 与点估计模型相比，全贝叶斯深度学习框架是否能提升无监督可变形图像配准的配准精度？
• RQ2 通过SGLD结合自适应噪声注入的后验采样，是否能减少在小样本医学影像数据集上的过拟合并提升鲁棒性？
• RQ3 所提出的方法是否能在不假设参数化后验分布的前提下，实现对形变场预测的不确定性量化？
• RQ4 在不同输入图像噪声水平下，该贝叶斯方法与VoxelMorph的性能相比如何？

主要发现

• 所提出的ASGLD Reg方法在MNIST数据集上实现了0.0063的均方误差（MSE），优于VoxelMorph的0.0065。
• 在MGH10脑部MRI数据集上，该方法实现了0.736（±0.017）的Dice分数，显著高于VoxelMorph的0.7109（±0.027）。
• 在所有测试的噪声水平下，性能提升均具有统计学显著性（MNIST的p < 0.05，MGH10的p < 10^-5）。
• 与仅使用最终模型权重相比，平均后验样本显著提升了在噪声条件下的鲁棒性。
• 该方法通过表征真实后验分布成功实现了对形变场不确定性的量化，且无需假设后验分布为正态分布。
• 由于采用非参数化后验采样，该框架在小样本数据集上表现出更强的泛化能力并有效减少了过拟合。