【论文精度】Mattias P.H & Lasse Hansen：PDD-net 3D，2.5D 两篇论文笔记

简介

PDD-Net 是德国吕贝克大学的 Mattias P. Heinrich 和 Lasse Hansen 在深度神经网络应用于医学图像配准（MICCAI 2020 - Learn2Reg chanllenge）的 SOTA 成果，他们在2021年3月发表了总结论文，参见Mattias & Lasse：PDD-net in MICCAI。其中，在他们的 PDD-Net 在所有挑战选手中排名前三，四个子任务分别获得两个第一名和三个第二名。以下是 PDD-Net 前后两个版本的相应文章以及他们在 MICCAI 赛后发表的解释性文章。

参考 DL 医学图像处理大佬 棉花糖 的文章 【论文笔记】PDD-Net（概率密度位移网络）：弱监督的3D医学图像配准网络。

一. Closing the Gap between Deep and Conventional Image Registration using Probabilistic Dense Displacement Networks 👉

• Mattias P. Heinrich
• 3D abdominal CT scans of the VIS-CERAL3
• pdd-net
• weakly-supervised
• non-local loss

1. Abstract：

We address this shortcoming by leverag-ing ideas from probabilistic dense displacement optimisation that hasexcelled in many registration tasks with large deformations. We proposeto design a network with approximate min-convolutions and mean fieldinference for differentiable displacement regularisation within a discrete weakly-supervised registration setting.

• Question： 深度学习方法在大变形配准下表现很差。
• Method： 引入概率密集位移优化（probabilistic dense displacement optimisation，pdd），设计一个具有近似最小卷积和平均场推断的网络，用于在离散弱监督配准环境下的可微位移正则化。
• Answer： 参数量少、训练快；SOTA；能处理大变形配准。

2. Introduction：

• 为什么研究这个课题？

  • Deep Learning based Image Registration(DLIR) 主要应用于不太复杂（变形小，标签多）的配准；
  • 对于腹部 CT 这种大变形配准，SOTA（label-reg、niftyReg）表现很差；
  • justifies further research，值得进一步研究。

• 目前这个课题的研究进行到了哪一阶段？

  • 如果不借助复杂的多阶段扭曲（warp）管道（pipeline），就很难用深度连续回归网络对不同患者之间的大而高的变形进行建模。相反，离散配准（discrete registration）的使用，探索了大量的定量位移空间，已被证明更有效地捕获腹部和胸部这种大变形，并且可以用几个或一个扭曲步骤来实现。
  • FlowNet-C，离散位移场在 2D 图像配准上应用，但是优化参数太大，用不到3D上；
  • 概率和不确定性建模已经在 DLIR 中研究过，但没有在离散配准中研究过。

• 作者使用的理论是基于哪些假设？（看不明白也没关系，先抄下来）

  • 我们提出了一个新的 DLIR 学习模型，通过引入带有可微约束的强正则化项，该模型可以更好地利用概率密集位移采样（pdd sampling）的优点，并明确考虑到解决 3D 配准问题。
  • 我们的特征提取器使用 3D 可变形卷积（OBELISK）¹，非常轻量级。据我们所知，这是第一个将离散 DLIR 与均场正则化的可微分应用相结合的方法。
  • 与以前的工作相比，我们的模型需要较少的可训练权重，能捕捉较大的变形，并且可以用少量标签训练到高精度。
  • 引入了一种新的非局部标签相似性损失函数²，代替了目前广泛使用的基于空间变压器的损失函数，提高了求导性能。

3. Conclusion：

• 这篇文章存在哪些缺陷？
  • 对更大的数据集进行更详细的评估，并采用额外的评估指标（表面距离），可以更深入地了解该方法的优缺点。
• 作者关于这个课题的构思有哪几点？
  • 我们的新型pdd-net结合了 概率密集位移（probabilistic dense displacements，pdd） 和 微分平均场正则化（differentiable mean-field regularisation），为患者间腹部 CT 中的 7 个大的解剖结构的一对一配准实现了超过 70% 的 Dice 精度。
  • 未来的工作可以通过使用多个配准阶段和更自适应的控制点采样来获得更多的增益。

4. 研究方法：

• 研究的数据从哪里来？
  • 3D abdominal CT scans of the VIS-CERAL3
• 研究中用到的重要指标有哪些？
  • Dice
  • std(Jac)
  • runtime
• 模型分哪几步？每一步分别得出了什么结论？

  • Convolutional feature learning network：we employ the Obelisk approach¹

  • PDD-net 在 OBELISK-net 之前加上一层 4 通道的 5×5×5 的卷积核，以便提取到边缘特征；这个网络共有 64 个空间滤波偏移量（卷积核），总共120k 的可训练参数量。

    We extend the authors’ implementation by adding a normal 5 ×5×5 convolution kernel with 4 channels prior to the Obelisk layer to also learn edge-like features. The network has 64 spatial filter offsets and in total 120k trainable parameters, shared for fixed and moving scan to yield $f(I_F)$ and $f(I_M)$.

  • Correlation layer for dense displacement dissimilarity： 给出第一部分的特征表示，我们的目标是找到一个正则化的位移场，它为非线性变换的每个控制点分配一个向量

  • Regularisation using min-convolutions and mean-field inference： 我们将正则化约束建模为网络体系结构的一部分。

  • Probabilistic transform losses and label supervision： 我们进一步利用位移采样的概率性质，并基于非局部均值加权（non-local means weighting）计算有监督的标签损失函数。

  • 3D pdd-net： 下一篇：We build our efficient 2.5D discrete registration framework upon the 3D pdd-net (probabilistic dense displacement network) proposed in 这一篇。这篇 PDD-net 总结如下：

    • 使用可变性卷积（OBELISK-net）来提取特征，输出具有 3 个体素的步幅的 24 通道特征图（ 4D tensor）。
    • 定义一个粗网格（coarser grid）上的一组 $|K|\in {\mathbb{R}}^3$ 控制点的 B-样条曲线转换模型，以及一个线性间距为 $\mathcal{L}=q\cdot {\left\{-1,-\frac{6}{7},-\frac{5}{7},\dots ,+\frac{5}{7},+\frac{6}{7},+1\right\}}^3$ 的 3D 位移空间，这里 $q$ 是定义捕获范围的标量。
    • 使用不带任何可训练权值的相关层（参见³）来计算固定图像特征向量与在 $\mathcal{L}$ 所跨越的搜索空间内的所有离散位移的浮动图像特征向量的匹配代价（matching cost）。
    • 网络的第二部分包括若干 3D 最大和平均池化运算（步长=1），这些运算交替作用于三个空间或三个位移维度，并对 MRF 正则化的两个迭代进行建模（近似于⁴中的平均场推断）。

如上所述，这些密集的三维位移计算能够实现高精度的配准，但会带来大量的内存和计算成本。因此，下一篇我们用三个 2.5D 子空间来近似稠密的 3D 空间，并定义 L 2 D = { L x y , L x z , L y z } \mathcal{L_{2D}} = \left\{ \mathcal{L_{xy}}, \mathcal{L_{xz}}, \mathcal{L_{yz}} \right\} L2D​={Lxy​,Lxz​,Lyz​}，其中每个子空间包含二维位移平面网格：${\mathcal{L}}_{xy}=q\cdot {\left\{-1,-\frac{6}{7},-\frac{5}{7},\dots ,+\frac{5}{7},+\frac{6}{7},+1\right\}}^2$ 和第三维的常数值 $z_1$。这一步极大地减少了在可变形卷积部分学习更好的特征提取所需的内存。

5. 论文图表:

二. Highly accurate and memory efficient unsupervised learning-based discrete CT registration using 2.5D displacement search 👉

• Mattias P. Heinrich and Lasse Hansen
• 3D abdominal CT scans
• pdd-net 2.5
• unsupervised
• MIND loss

1. Abstract：

We propose a highly accurate unsupervised learning framework for 3D abdominal CT registration that uses a discrete displacement layer and a contrast-invariant metric (MIND descriptors) that is evaluated in a probabilistic fashion. We realise a substantial reduction in memory and computational demand by iteratively subdividing the 3D search space into orthogonal planes.

• Question： 深度学习方法在3D大变形配准下精度不够高、内存消耗大。
• Method： 提出无监督三维腹部CT配准框架，该框架采用一个离散位移层和以概率方式评估的对比不变（contrast-invariant）度量函数（MIND-loss）。
• Answer： 我们将神经网络搜索空间缩小了 5 倍，运行速度提高了 2 倍，并且在精度方面与之前的离散 3D 网络相当。

2. Introduction：

• 为什么研究这个课题？

  • Deep Learning based Image Registration（DLIR）对于大变形配准要减少计算用时和内存；
  • 无监督（Voxelmorph）方法⁵是值得关注的，因为它们可以减少推理运行时间，而无需依赖大量的手动注释的真实标签（ground truth）；
  • DLIR 配准精度没有得到提升，其应用仍然局限于较不复杂的任务；
  • U-Net 类型的（原本是应用于图像分割）多尺度（multi-scale）编解码配准网络结构似乎限制了配准精度的提升；
  • 采用稀疏空间采样可以减少参数量的爆炸性增长，但这种非常稀疏的3D空间采样会导致配准精度的显著降低。

• 目前这个课题的研究进行到了哪一阶段？

  • 在基于MRF的离散优化中，探索多个量化为标签的潜在位移，在二维图像和医疗卷的大变形图像配准中取得了巨大成功：FlowNet、PWC-Net、pdd-net；
  • FlowNet-C，离散位移场在2D图像配准上应用，但是优化参数太大，用不到3D上；
  • 上一篇 pdd-net 取得了很好的成果；
  • 我们是第一个提出一个无监督的离散深度学习框架的 3D 医学图像配准，我们利用概率预测，以改善损失函数的性能。

• 作者使用的理论是基于哪些假设？（看不明白也没关系，先抄下来）

  • Discrete displacements： 离散方法不是直接回归连续位移值，而是量化预期位移的范围并估计（空间正则化）标签分配的概率（或成本）。我们预计，全采样的3D位移场对于达到 SOTA 是不必要的，因此减少标签和内存的使用，将使离散 DLIR 作用更大。
  • Subdivision of label spaces： 在许多传统的基于 MRF（马尔可夫）的配准方法和策略中，参数数量的爆炸性增长是一个问题，但是，非常稀疏的三维空间采样会导致配准精度的显著降低。在这里，我们提出了一种中间策略，将三维空间分解成三个正交的二维平面来进行神经网络结构搜索（NAS），从而更准确地扫描整个空间，在保证精度不丢失的同时减少参数量。
  • 2.5D approaches in deep learning： 一些最近的医学 3D 图像分割和分类网络把 3D 的输入数据看成多个 2D 视图，这样得到的信息是片面的。不使用 3D 卷积核处理 3D 输入的方式主要有以下几种：
    • 通过提取 3D 面片的正交平面并将其表示为 2D CNN 的 RGB 颜色通道，创建了 2.5D 输入视图，从而实现了从 ImageNet 预训练模型的迁移学习。
    • ⁶ 将三个正交空间平面上的运动预测解耦（轴向、冠状面和矢状面），并将连续回归值融合。
    • ⁷ 提出了可分离 3D 卷积的 S3D 模型。每个 3D 卷积可以被两个连续的卷积层代替：一个 2D 卷积学习空间特征，一个 1D 卷积学习时间特征。本文在 3 个正交 2D 平面采样，有点类似第二种方式。

可分离 3D 卷积：

3. Conclusion：

• 这篇文章存在哪些缺陷？
  • 某个方法必须作用于6D张量
• 作者关于这个课题的构思有哪几点？
  • We demonstrate that the high accuracy of a full 3D search space can be matched using two iterations of sampling three orthogonal 2D displacement maps in combination with on-the-fly instance optimisation.
  • 提出了一种漂亮的非局部度量 MIND，利用无监督学习的概率预测，实现快速训练。
  • 基于离散深度学习的图像配准（DLIR）的 2.5D 位移空间细分，与快速实例优化相结合，为高度复杂的腹部主体间配准提供了 SOTA 技术。
  • 未来的工作可以通过使用更强大的正则化网络结合上下文损失项，来进一步获得内存高效的 2.5D 位移空间。

4. 研究方法：

• 研究的数据从哪里来？

  • 3D abdominal CT scans （MICCAI learn2reg 2020 task3）

• 研究中用到的重要指标有哪些？

  • Dice
  • memory
  • infer time (GPU、CPU)

• 模型分哪几步？每一步分别得出了什么结论？

  • Decoupled 2.5D subspaces： 相较于上一篇的 3D PDD-net，这次我们用三个 2.5D 子空间来近似稠密的 3D 空间，并定义 L 2 D = { L x y , L x z , L y z } \mathcal{L_{2D}} = \left\{ \mathcal{L_{xy}}, \mathcal{L_{xz}}, \mathcal{L_{yz}} \right\} L2D​={Lxy​,Lxz​,Lyz​}，其中每个子空间包含二维位移平面网格：${\mathcal{L}}_{xy}=q\cdot {\left\{-1,-\frac{6}{7},-\frac{5}{7},\dots ,+\frac{5}{7},+\frac{6}{7},+1\right\}}^2$ 和第三维的常数值 $z_1$。

    • 这一步极大地减少了在可变形卷积部分学习更好的特征提取所需的内存，相关层内的特征差异计算从 4 GFlops减少到 790 MFlops。
    • 空间平滑仍需在 3D 上操作（但通道数量更少了），而调整位移尺寸的操作现在在 2D 中。
    • 为了估计三维位移场 $\phi$，网络的输出被转换成每个控制点的三个二维伪概率映射（pseudo-probability maps，使用 softmax）
    • 将概率网格（mesh-grid ）与位移网格相乘，并在所有三个相交映射（maps）的两个非零元素之间求平均值，从而找到三维向量位移场 $\phi$。增加了扩散正则化 $\lambda \cdot \left({\left|\nabla {\phi }_1\right|}^2+{\left|\nabla {\phi }_2\right|}^2+{\left|\nabla {\phi }_3\right|}^3\right)$，以促进合理的变形。

  • Two-step instance optimisation with gradient descent： 全三维位移空间中三个二维平面的逼近精度取决于它们的交点与真实最优点的接近程度。在初始的 2.5D 子集上，网络的正则化部分的直接传播可能会导致训练和推理的不准确。因此，本文提出了一种两步方法，即在度量损失和扩散正则化之间寻找最佳折衷点，以减轻前馈网络的有限能力。

    • 与 ⁵ 和 ⁸ 类似，对中间 2.5D 位移概率映射进行了迭代的动态（on-the-fly）优化。
    • 考虑了一个连续的三维 B-样条变换模型，并使用 Adam 在每对配准图片之间的改进后的位移场 ${\phi }^{\ast }:=\left(x+\Delta x^{\ast },y+\Delta y^{\ast },z+\Delta z^{\ast }\right)$ 进行优化。
    • 离散代价张量：$\mathcal{C}\in H\times W\times D\times 15\times 15\times 3$，其中最后的维数 3 表示三维位移搜索区域的三个 2.5D 子平面。
    • 我们定义一个子维度 $L_{xy}=\mathcal{C}\left(x,y,z,\Delta x^{\ast },\Delta y^{\ast },0\right)$
    • 最小化损失函数：$L_{\text{instance }}=L_{xy}+L_{xz}+L_{yz}+L_{\text{diff.-reg }}$，三个 2D 位移度量值和扩散正则化损失的总和）。
    • 以连续的方式迭代更新变形并优化代价函数，将离散优化和连续优化的互补优势结合起来。一旦第一次变形配准优化完成，将执行离散网络的第二次传递，更精确地初始化放置正交的 2.5D 子平面（参见图1右侧），然后再次进行连续优化。（此为两步法）

  • Non-local metric loss： 上一篇中的损失函数仅限于带分割标签的弱监督学习，因此，我们引入了一种新的离散 DLIR 的无监督非局部度量损失：Modality Independent Neighbourhood Descriptors (MIND) ⁹

    • 充分利用 2.5D 概率位移预测：${\mathrm{M}\mathrm{I}\mathrm{N}\mathrm{D}}_{\text{warped }}=\frac{1}{3}{\mathrm{M}\mathrm{I}\mathrm{N}\mathrm{D}}_{\mathrm{x}\mathrm{y}}+\frac{1}{3}{\mathrm{M}\mathrm{I}\mathrm{N}\mathrm{D}}_{\mathrm{x}\mathrm{z}}+\frac{1}{3}{\mathrm{M}\mathrm{I}\mathrm{N}\mathrm{D}}_{\mathrm{y}\mathrm{z}}$
    • 概率位移定义为：${\mathcal{P}}_{xy}\left(x,y,z,\Delta x^{\prime },\Delta y^{\prime }\right)=\frac{\exp \left(-\alpha \mathcal{C}\left(x,y,z,\Delta x^{\prime },\Delta y^{\prime }\right)\right)}{{\sum }_{\Delta x^{\prime },\Delta y^{\prime }}\exp \left(-\alpha \mathcal{C}\left(x,y,z,\Delta x^{\prime },\Delta y^{\prime }\right)\right)}$
    • 离散扭曲的 MIND 特征图计算如下：${\mathrm{M}\mathrm{I}\mathrm{N}\mathrm{D}}_{\mathrm{x}\mathrm{y}}={\sum }_{\Delta x^{\prime },\Delta y^{\prime }}{\mathcal{P}}_{xy}\left(x,y,z,\Delta x^{\prime },\Delta y^{\prime }\right)\cdot \mathrm{MIND}(x+\Delta x^{\prime },y+\Delta y^{\prime },z)$

• 论文图表