2020 MICCAI Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains

Shape-aware Meta-learningShape-aware Meta-learning : 形状感知元学习

通过写文章总结，进一步理解作者的思路，如何发现问题，针对问题，作者又是怎么解决的？解决的方法是来自哪些文章？这些文章什么时间发表在什么期刊或会议上？

Shape-aware Meta-learning ：形状感知元学习

Generalizing Prostate MRI Segmentation to Unseen Domains : 泛化前列腺MRI分割到未见的域

中文题目：形状感知元学习用于泛化前列腺MRI分割到未见过的域

作者信息：中国香港中文大学，计算机科学与工程系，Quande Liu, Di Dou

摘要

问题：这篇文章主要解决域泛化（Domain Generalization）问题，例如：从多个源域数据学习一个模型直接泛化到未见过的目标域中。

作者提出一个基于形状感知元学习方法来改善模型关于MRI前列腺分割的泛化能力。该方法主要梯度的元学习，通过在训练过程中利用虚拟的元训练（meta-train）与元测试（meta-test）来显示地模拟域迁移（Domain shift）。同时作者还发现一个问题，将分割模型运用到未见过的域（unseen domain）数据会存在一些问题，例如（模型预测的mask存在不完整的mask，以及模糊的边界），作者针对该问题，引入两个互补的损失函数来增强元优化（meta optimization），即在模拟的域迁移情况下，鼓励分割的形状紧凑性（shape compactness）和形状平滑性（shape smoothness）。作者在6个不同机构所提供的前列腺MRI数据上评估方法的性能。

方法

形状感知元学习方法的流程图（Overview）如下所示。

源域数据被随机划分成元训练（meta-train）和元测试（meta-test）来模仿域迁移。（1）通过在元测试中约束形状紧凑性实现具有完整形状的分割。（2）通过改善轮廓（contour）嵌入和轮廓附近背景区域的嵌入之间的类内凝聚力（intra-class cohesion）和类间分离（inter-class separation），并且与数据域无关，进而增强域不变性（domain-invariant）以实现稳定的边界分割。

$(X,Y)$ 表示分割任务中的联合输入和标签空间。 $\left \{ D_{1},D_{2},...,D_{K} \right \}$ 表示K个域。每个域 $D_{K}$ 中都包含 $N_{K}$ 个配对的图像标注数据 $\left\{ (x_{n}^{k},y_{n}^{k}) \right \}_{n=1}^{N_{k}}$ 。目的是从所有的源域数据D中学习一个分割模型 $F:X \rightarrow Y$ ，可以泛化到没有见过的目标域 $D_{tg}$ 。

2.1 基于梯度的元学习方法 Gradient-based Meta-learning Scheme

这篇文章的基础是基于以下这篇文章。通过在训练过程中模拟真实世界中的域迁移（domain shift）来改善鲁棒优化（robust optimization）。

Li D, Yang Y, Song Y Z, et al. Learning to generalize: Meta-learning for domain generalization[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2018, 32(1).

具体来说，在模型的每次迭代过程中。源域数据被随机划分成元训练（meta-train） $D_{tr}$ ，元测试数据 $D_{te}$ 。元学习（meta learning）可以分为两步。(1) 在Dice分割损失函数约束下，使用元训练数据 $D_{tr}$ 更新分割模型Segementor 的参数 $\theta$ ： $\theta^{'}= \theta - \alpha \bigtriangledown_{\theta} L_{seg}(D_{tr};\theta)$ ， $\alpha$ 是内循环（inner-loop）的学习率参数。（2）

第二步需要结合代码进一步理解。总的来说，模型除了在元训练上学习分割任务，还需要学习如何在元训练和元测试数据模拟的域迁移（Domain Shift）来进行泛化。换句话说，对模型进行优化，使在虚拟源域训练数据上所学习到的参数更新，也能提高虚拟目标域的性能。

在分割问题中，作者希望模型在未见过的域数据上能够很好地保留分割结果的完整的形状（complete shape）和平滑的边界（smooth boundary）。作者通过引入两个互补的损失项到元目标函数， $L_{meta} = L_{seg}+\lambda_{1} L_{compact} + \lambda_{2}L_{smooth}$ ，在域迁移的情况下，显示地增加分割图的形状紧凑性与形状平滑，进一步提高泛化性能。

2.2 元形状紧凑性约束 Meta shape compactness constraint

传统的分割损失函数（Dice, 交叉熵）评估像素对应的准确率，没有对分割图进行全局的约束。这些损失函数会导致模型在分布迁移（Distribution shift）的情况下不能产生完整的分割。以下文献研究表明：

Fan R, Jin X, Wang C C L. Multiregion segmentation based on compact shape prior[J]. IEEE Transactions on Automation Science and Engineering, 2014, 12(3): 1047-1058.

对于紧凑性目标，约束形状紧凑性有助于改善完整形状的分割，因为不规则形状的不完整分割往往对应一个糟糕的紧凑性属性。而前列腺区域通常表现一个紧凑性形状，这种形状先验信息是独立于所观察到的域，（换一种说法，给定一只猫的图像，保留猫的轮廓，再使用大象的纹理去取代猫的皮毛纹理，人类倾向于认为图像的类别是猫，但是网络却会判定为大象。这种不一致性，会导致网络强行学习到的规律和人类不一致，很难完成对人类视觉系统的模拟。具体到深度估计领域，图像的纹理变化，例如不同的光照、天气、季节造成的影响都会对模型产生较大的影响[1]。不同域中前列腺的区别包括纹理变化，例如不同的成像设备，成像机制，不同患者造成的影响都会对模型产生较大的影响。）。因此，作者显示地整合紧凑性形状约束到元目标函数（meta-objective） $L_{meta}$ ，进而鼓励模型在域迁移的情况下，很好地保留形状完整性（shape completeness）。具体地来说，作者采用已经建立的Iso-perimetric Quotient 测量[2]来量化形状紧凑度，具体定义为 $C_{IPQ} = 4\pi A/P^{2}$ ，A是形状面积，P是周长。作者采用 $C_{IPQ}$ 的倒数形式定义形状紧凑型损失函数 $L_{compact} = \frac{P^{2}}{4 \pi A}=\frac{\sum_{i\in \Omega }\sqrt{(\bigtriangledown p_{u_{i}})^{2}+ (\bigtriangledown p_{vi})^{2}) + \epsilon }}{4 \pi (\sum_{i \in \Omega }\left | p_{i} \right | + \epsilon)}$ ，p是预测的概率图， $\Omega$ 是概率图中所有像素的集合， $\bigtriangledown p_{ui},\bigtriangledown p_{vi}$ 分别是像素在水平方向和垂直方向上的概率梯度， $\epsilon$ 是用于计算稳定性的超参数。P是所有像素的梯度大小之和。A是概率率图P绝对值absolute value之和。因此，最小化 $L_{compact}$ 可以得到完整形状的分割图（segmentation map），因为不规则形状（irregular shape）往往表现出相对较小区域面积A，相对较大的长度P，导致 $L_{compact}$ 损失值变大。通过约束 $L_{compact}$ ，有利于模型保留未见过域数据中完整的形状，而不是过拟合源域数据。

2.3 元形状平滑增强 Meta shape smoothness Enhancement

作者观察到在不同域上模型的性能下降主要来自模糊的边界区域。因此，作者提出一个目标函数 $L_{smooth}$ 来增强边界轮廓（boundary deineation），进而显示地促进轮廓相关嵌入（contour-relevant embedding ）与背景相关嵌入（background-relevant embedding）之间的类内一致性（intra-class cohesion）和类间分离性（iter-class separation)，这两个嵌入从所有域的每个样本中提取的到。

轮廓相关嵌入（contour-relevant embedding） $E_{m}^{con} \in \mathbb{R}^{C_{l}}$ ，背景相关嵌入（background-relevant embedding） $E_{m}^{bg} \in \mathbb{R} ^{C_{l}}$ 计算方式如下： $E_{m}^{con}=\frac{\sum_{i \in \Omega} (T_{m}^{l})_{i} \cdot (c_{m})_{i}}{ \sum_{i \in \Omega} (c_{m})_{i} }$ ， $E_{m}^{bg}=\frac{\sum_{i \in \Omega} (T_{m}^{l})_{i} \cdot (b_{m})_{i}}{ \sum_{i \in \Omega} (b_{m})_{i} }$ ，其中， $T_{m}^{l} \in \mathbb{R}^{H \times W \times C_{l}}$ 是使用双线性上采样采样第l层激活图（activation map） $M_{m}^{l} \in \mathbb{R}^{H_{l} \times W_{l} \times C_{l}}$ 得到， $c_{m} \in \mathbb{R}^{H \times W \times 1},b_{m} \in \mathbb{R}^{H \times W \times 1}$ 分别是二值边界轮廓mask，二值背景mask，使用形态学操作从one-hot 标注 $y_{m}$ 计算的到。其中 $b_{m}$ 仅仅采样边界周围的背景像素，进一步增加对边界区域像素的判别性。在实现过程中，作者对解码器器的最后两层的输出进行上采样，然后进行特征拼接得到 $T_{m}^{l} \in \mathbb{R}^{H \times W \times C_{l}}$ ，具体还要结合代码进行理解。

考虑到讲这种正则化直接加入到网络的嵌入中，可能过于严格导致 $L_{seg},L_{compact}$ 不能收敛。作者采用对比学习（Constrast Learning）[3] 来实现这种约束，利用一个嵌入网络（embedding network） $H_{\phi}$ 投影特征 $E^{con},E^{bg}$ 到低纬空间，然后计算 $H_{\phi}$ 输出特征向量之间的距离 $d_{\phi}(E_{m}, E_{n})=\left \| H_{\phi}(E_{m})-H_{\phi} (E_{n}) \right \|_{2}$ ，其中样本m,n是从所有域中随机采样得到的。作者想通过这种约束协调（harmonize） $D_{tr}, D_{te}$ 嵌入空间来捕获边界区域周围的域不变表征。因此，对比损失定义如下：

$\mathit{l}_{constrastive} (m,n)=\left\{\begin{matrix} d_{\phi}(E_{m},E_{n}) ,& if \tau (E_{m}) = \tau (E_{n})\\ (max(0,\zeta -d_{\phi}(E_{m},E_{n})))^{2},& if \tau(E_{m})\neq \tau(E_{n}) \end{matrix}\right.$ ， $\tau(E)$ 表示类别，如果 $E$ 是 $E^{con}$ ，则 $\tau(E)=1$ ；如果E是 $E^{bg}$ ，则 $\tau(E)=0$ 。 $\zeta$ 是根据度量学习的做法预先确定的距离差值。

$C(q,2)$ 是组合的数量。通过约束 $L_{seg}(D_{tr};\theta)$ ， $L_{meta}(D_{tr},D_{te};\theta^{'})$ 优化原始参数 $\theta$ 。通过约束 $L_{smooth}$ 优化 $H_{\phi}$ 。

数据集

作者使用6不同机构提供的公开MRI前列腺数据集，其中包括NCI-ISBI13数据集[4]（Site A、B），I2CVB数据集（Site C）[5]，PROMISE12 (Site D,E,F)[6]。各个不同数据的基本信息比较：

数据预处理：横截面调整到384 * 384，并归一化到零均值和单位方差，并裁减每个样本仅仅保留前列腺区域。作者使用Dice系数以及ASD（Average Surface Distance）来评估分割结果。

训练集，测试集的划分方法：

对比方法：

Zhang L, Wang X, Yang D, et al. Generalizing deep learning for medical image segmentation to unseen domains via deep stacked transformation[J]. IEEE transactions on medical imaging, 2020, 39(7): 2531-2540. BIGAug
Li D, Zhang J, Yang Y, et al. Episodic training for domain generalization[C]//Proceedings of the IEEE/CVF International Conference on Computer Vision. 2019: 1446-1455. Epi-FCR
Aslani S, Murino V, Dayan M, et al. Scanner invariant multiple sclerosis lesion segmentation from MRI[C]//2020 IEEE 17th International Symposium on Biomedical Imaging (ISBI). IEEE, 2020: 781-785. LatReg
Dou Q, Castro D C, Kamnitsas K, et al. Domain generalization via model-agnostic learning of semantic features[J]. arXiv preprint arXiv:1910.13580, 2019. MASF

定性比较：

定量比较：

DeepALL 使用所有域的数据训练一个深度模型；Intra-site :训练和测试仅仅使用一个域上的数据。

作者还比较了增加不同类别（域）的数据作为训练集训练模型（即当测试数据是SiteF, 情况1:使用训练数据SiteA，情况2:使用训练数据SiteA、B;情况3:使用训练数据SiteA、B、C....情况5:使用训练数据SiteA、B、C、D、E），并在未见过域数据上测试模型的性能。如下图所示。

参考文献：

https://mp.weixin.qq.com/s/xdV3VHbSrgQa3G51KMHamA
Li W, Goodchild M F, Church R. An efficient measure of compactness for two-dimensional shapes and its application in regionalization problems[J]. International Journal of Geographical Information Science, 2013, 27(6): 1227-1250. 二维形状紧凑性度量以及在区域化问题中的应用。
Chen T, Kornblith S, Norouzi M, et al. A simple framework for contrastive learning of visual representations[C]//International conference on machine learning. PMLR, 2020: 1597-1607.
Bloch, N., Madabhushi, A., Huisman, H., Freymann, J., Kirby, J., Grauer, M.,Enquobahrie, A., Jaffe, C., Clarke, L., Farahani, K.: Nci-isbi 2013 challenge: automated segmentation of prostate structures. The Cancer Imaging Archive 370 (2015).
Lema^ıtre, G., Mart´ı, R., Freixenet, J., Vilanova, J.C., Walker, P.M., Meriaudeau, F.: Computer-aided detection and diagnosis for prostate cancer based on mono and multi-parametric mri: a review. CBM 60, 8{31 (2015).
Litjens, G., Toth, R., van de Ven, W., Hoeks, C., Kerkstra, S., van Ginneken, B.,Vincent, G., Guillard, G., Birbeck, N., Zhang, J., et al.: Evaluation of prostate segmentation algorithms for mri: the promise12 challenge. MIA 18(2), 359{373 (2014).