Shape-Aware Meta-Learning 在模型泛化中引入形状约束

论文来源：Liu, Quande, Qi Dou, and Pheng-Ann Heng. “Shape-aware Meta-learning for Generalizing Prostate MRI Segmentation to Unseen Domains.” In International Conference on Medical Image Computing and Computer-Assisted Intervention, pp. 475-485. Springer, Cham, 2020. (Code from https://github.com/liuquande/SAML)

Motivation

域泛化问题

例如不同中心采集的MRI前列腺数据，存在明显的域差异问题。因此采用在已知中心 (如 Site A, B, C, D, E) 上训练的模型在未知中心 (如 Site F) 上测试，会导致显著的分割错误 —— 即域泛化问题。

元学习 Meta-learning

Meta learning (元学习) 可以用来处理上述模型方法，通过在训练过程中切分数据 meta-train 和 meta-test 来显式模拟已知和未知域之间的域差异。然而此类方法大多数用来处理 image-level 的自然图像分类问题，不适合用来处理逐像素预测的分割问题——其中重要的问题在于如何引入 Shape-based regularization，导致不完整的形状以及含糊的边界。

Shape-Aware Meta-Learning

shape-aware loss function

基于上述对元学习缺点的分析，我们希望模型能够保持形状完整性 shape compactness 和 形状平滑性 shape smoothness。因此，将两个互补的形状约束项引入传统的 meta-learning 损失函数中：
$${\mathcal{L}}_{\text{meta}}={\mathcal{L}}_{\text{seg}}+{\lambda }_1{\mathcal{L}}_{\text{compact}}+{\lambda }_2{\mathcal{L}}_{\text{smooth}}$$
其中 ${\lambda }_1$ 和 ${\lambda }_2$ 分别表示对形状完整性和形状平滑性的均衡。

shape complementation

考虑到前列腺呈现 compact shape，因此采用 Iso-Perimetric Quotient 度量 $C_{IPQ}=4\pi A/P^2$，其中 $A$ 和 $P$ 分别代表形状面积和边缘长度。将上述度量转化到分割任务中，形成 shape compactness constraint：
$${\mathcal{L}}_{\text{compact}}=\frac{P^2}{4\pi A}=\frac{{\sum }_{i\in \Omega }\sqrt{{\left(\nabla p_{u_i}\right)}^2+{\left(\nabla p_{v_i}\right)}^2+ϵ}}{4\pi \left({\sum }_{i\in \Omega }\left|p_i\right|+ϵ\right)}$$
直观来说，最小化上述 ${\mathcal{L}}_{compact}$ 鼓励分割结果具有完整 compact shape —— 因为不完整的形状常具有较小的区域 $A$ 然而具有较大的 $P$，即较大的 ${\mathcal{L}}_{compact}$。

shape smoothness

要求分割结果具有平滑的边缘，通过正则化 domain-invariant contour-relevant background-relevant embedding，提升类内一致性和类间差异性。具体来说，通过 mask-average pooling 方法得到边缘和背景 embedding 的结果：
$$E_m^{con}=\frac{{\sum }_{i\in \Omega }{\left(T_m^l\right)}_i\cdot {\left(c_m\right)}_i}{{\sum }_{i\in \Omega }{\left(c_m\right)}_i},E_m^{bg}=\frac{{\sum }_{i\in \Omega }{\left(T_m^l\right)}_i\cdot {\left(b_m\right)}_i}{{\sum }_{i\in \Omega }{\left(b_m\right)}_i}$$
直接约束上述 embedding 结果是过于严格的，因此采用对比学习方法 —— 即再过一个 embedding network，将其再映射到低维空间，然后在此低维空间上计算距离 $d_{\varphi }\left(E_m,E_n\right)={\Vert H_{\varphi }\left(E_m\right)-H_{\varphi }\left(E_n\right)\Vert }_2$。最后形成 shape smoothness constraint：
ℓ contrastive ( m , n ) = { d ϕ ( E m , E n ) ,  if  τ ( E m ) = τ ( E n ) ( max ⁡ { 0 , ζ − d ϕ ( E m , E n } ) 2 ,  if  τ ( E m ) ≠ τ ( E n ) \ell_{\text {contrastive}}(m, n)=\left\{\begin{array}{ll} d_{\phi}\left(E_{m}, E_{n}\right), & \text { if } \tau\left(E_{m}\right)=\tau\left(E_{n}\right) \\ \left(\max \left\{0, \zeta-d_{\phi}\left(E_{m}, E_{n}\right\}\right)^{2},\right. & \text { if } \tau\left(E_{m}\right) \neq \tau\left(E_{n}\right) \end{array}\right. ℓcontrastive​(m,n)={dϕ​(Em​,En​),(max{0,ζ−dϕ​(Em​,En​})2,​ if τ(Em​)=τ(En​) if τ(Em​)​=τ(En​)​
直观来说，上述约束确保同样属于边缘的像素具有更加相似的特征，然而边缘和背景的像素具有更加具有区分度的特征 —— 使分割边缘不再 ambiguous。

总结

个人认为上述论文的关键在于将形状约束做了两个角度的拆分 —— 形状和边缘。通常而言形状更加关注内部特性和整体拓扑结构 (如上文中的 $A$ 和 $P$)，而边缘通常可以添加平滑约束等来确保相对外部特征 (如上文的 $E_m^{con}$)。总体来说，这篇文章还是具有相当的启发性，其形状约束是可以在其他工作中重复使用的。