各种 Dice Loss 变体

各种 Dice Loss 变体

语雀文档：https://www.yuque.com/lart/idh721/gpix1i

Dice Loss也是图像分割任务中非常常见的一个损失函数。本文基于 Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks 中的内容进行了整理。

hard dice score for binary segmentation

dice score 被广泛使用的针对二值分割图 S 和 G 之间成对比较的重叠度量方式。

其可以表示为集合操作或统计性度量的形式：

$$D_{hard}=\frac{2|S\cap G|}{|S|+|G|}=\frac{2{\Theta }_{TP}}{2{\Theta }_{TP}+{\Theta }_{FP}+{\Theta }_{FN}}=\frac{2{\Theta }_{TP}}{2{\Theta }_{TP}+{\Theta }_{AE}}$$

这里涉及到几项，具体含义如下：

• $S\&G$：待评估图像和参考图像
• ${\Theta }_{TP}$：正阳性样本的数量，即 $S$ 和 $G$ 都为真的位置的数量。
• ${\Theta }_{FP}$：$S$ 中真而 $G$ 中为假的位置的数量。
• ${\Theta }_{FN}$：$S$ 中假而 $G$ 中为真的位置的数量。
• ${\Theta }_{AE}={\Theta }_{FP}+{\Theta }_{FN}$：$S$ 和 $G$ 不一致的位置的数量。

soft dice score for binary segmentation

对于软二值分割的扩展依赖于概率分类对的不一致概念。

对于 $S$ 和 $G$ 中的位置 $i\in \mathbf{X}$ 对应的类别 $S_i$ 和 $G_i$ 可以被定义为标签空间 L = { 0 , 1 } \mathbf{L}=\{0,1\} L={0,1} 上的随机变量。

概率分割可以被表示为标签概率图，其中 $P(\mathbf{L})$ 表示标签概率向量的集合：

• p = { p i : = P ( S i = 1 ) } i ∈ X p=\{p^i:=P(S_i=1)\}_{i \in \mathbf{X}} p={pi:=P(Si​=1)}i∈X​
• g = { g i : = P ( G i = 1 ) } i ∈ X g=\{g^i:=P(G_i=1)\}_{i \in \mathbf{X}} g={gi:=P(Gi​=1)}i∈X​

由此可以将前面的关于数据的统计量 ${\Theta }_{TP}\&{\Theta }_{AE}$ 扩展到软分割情形：

• ${\Theta }_{AE}={\sum }_{i\in \mathbf{X}}|p^i-g^i|$
• ${\Theta }_{TP}={\sum }_{i\in \mathbf{X}}{g}^i(1-|p^i-g^i|)$

对于一般情形中的 $g$，即 ∀ i ∈ X , g i ∈ { 0 , 1 } \forall i \in \mathbf{X}, g^i \in \{0, 1\} ∀i∈X,gi∈{0,1}，此时有：

• ${\Theta }_{AE}={\sum }_{i\in \mathbf{X}}{g}^i(1-p^i)+(1-g^i)p^i={\sum }_{i\in \mathbf{X}}{g}^i+p^i-2g^i{p}^i$
• ${\Theta }_{TP}={\sum }_{i\in \mathbf{X}}{g}^i{p}^i$

对应的 soft dice score 可以表示为：

$$D_{soft}(p,g)=\frac{2{\sum }_i{g}^i{p}^i}{{\sum }_i(g^i+p^i)}$$

当然，也有引入平方形式的变体。

soft multi-class dice score

前面直接讨论的是二值分割的情形，而对于多分类情况则需要考虑不同类别计算的整合方式。

最简单的方式就是直接考虑所有类别的平均。

可以称为 mean dice score，这里对应包含 $|\mathbf{L}|$ 个不同的类别：

$$D_{mean}(p,g)=\frac{1}{|\mathbf{L}|}\sum _{l\in \mathbf{L}}\frac{2{\sum }_i{g}_l^i{p}_l^i}{{\sum }_i{g}_l^i+p_l^i}$$

上式的推广形式可以通过引入类别权重参数 $w_l=\frac{1}{({\sum }_i{g}_l^i{)}^2},l\in \mathbf{L}$ 而得到。即从而将上式转化为加权平均的形式。这被称为 generalised soft multi-class dice score。

最终可以表示为：

$$D_{generalised}(p,g)=\frac{2{\sum }_l{w}_l{\sum }_i{g}_l^i{p}_l^i}{{\sum }_l{w}_l{\sum }_i(g_l^i+p_l^i)}$$

soft multi-class wasserstein dice score

前面的 dice score 的形式中，对于 $p^i$ 和 $g^i$ 的相似性的度量方式可以看做是 L1 距离，而这里将 wasserstein distance 引入来自然地以一种语义上有意义的方式比较两个标签概率向量。

这里首先介绍 wasserstein distance。

wasserstein distance

这也被称为 earth mover’s distance。用于表示将一个概率向量 $p$ 变换为另一个概率向量 $q$ 所需要的最小成本。

对于所有的 $l,l^{\prime }\in \mathbf{L}$，从 $l$ 移动到 $l^{\prime }$ 的距离的集合定义为 $l$ 和 $l^{\prime }$ 之间的距离矩阵 $M_{l,l^{\prime }}$，这一矩阵是固定的，可以认为是已知的。

这是一种将 $\mathbf{L}$ 上的距离矩阵 $M$（通常亦可以称为 ground distance matrix）映射为 $P(\mathbf{L})$ 上的距离的方式，这里用了关于 $\mathbf{L}$ 的先验知识。

在 $\mathbf{L}$ 为有限集合的情况下，对于 $p,q\in P(\mathbf{L})$，二者关于 $M$ 的 wasserstein distance 可以被定义为一个线性规划问题的解。

W M ( p , q ) = min ⁡ T l , l ′ ∑ l , l ′ ∈ L T l , l ′ M l , l ′ subject to  ∀ l ∈ L , ∑ l ′ ∈ L T l , l ′ = p l ,  and  ∀ l ′ ∈ L , ∑ l ∈ L T l , l ′ = q l ′ WM(p,q)=minTl,l′∑l,l′∈LTl,l′Ml,l′subject to ∀l∈L,∑l′∈LTl,l′=pl, and ∀l′∈L,∑l∈LTl,l′=ql′

WM(p,q)subject to ∀l∈L,l′∈L∑​Tl,l′​ and ∀l′∈L,l∈L∑​Tl,l′​​=Tl,l′​min​l,l′∈L∑​Tl,l′​Ml,l′​=pl​,=ql′​​​

这里的 $T=(T_{l,l^{\prime }}{)}_{l,l^{\prime }\in \mathbf{L}}$ 是 $(p,q)$ 的联合概率分布，且有着边界分布 $p$ 和 $q$。

上式最小的 $\hat{T}$ 被称作对于距离矩阵 $M$ 在 $p$ 和之间 $q$ 的最优传输。

关于 wasserstein distance 的解释可以阅读：

• Wasserstein GAN and the Kantorovich-Rubinstein Duality
• https://chih-sheng-huang821.medium.com/%E9%82%84%E7%9C%8B%E4%B8%8D%E6%87%82wasserstein-distance%E5%97%8E-%E7%9C%8B%E7%9C%8B%E9%80%99%E7%AF%87-b3c33d4b942

soft multi-class wasserstein dice score

这里使用 wasserstein distance 来扩展标签概率向量对之间的差异性度量，从而得到如下扩展形式：

• ${\Theta }_{AE}={\sum }_{i\in \mathbf{X}}{W}^M(p^i,g^i)$
• Θ T P l = ∑ i ∈ X g l i ( M l . b − W M ( p i , g i ) ) , ∀ l ∈ L ∖ { b } \Theta^l_{TP}=\sum_{i \in \mathbf{X}}g^i_l(M_{l.b}-W^M(p^i,g^i)), \forall l \in \mathbf{L} \setminus \{b\} ΘTPl​=∑i∈X​gli​(Ml.b​−WM(pi,gi)),∀l∈L∖{b}

$M$ 选择为使得背景类别 $b$ 总是离其他类最远的情况。

${\Theta }_{TP}={\sum }_{i\in \mathbf{X}}{\alpha }_l{\Theta }_{TP}^l$

这里同样使用加权的方式对各个类别的统计结果进行了组合。

通过选择 ${\alpha }_l=W^M(l,b)=M_{l,b}$ 来使得背景位置并不对 ${\Theta }_{TP}$ 发挥作用。

最终，关于 $M$ 的 wasserstein dice score 可以定义为：

$$D^M(p,q)=\frac{2{\sum }_l{M}_{l,b}{\sum }_i{g}_l^i(M_{l,b}-W^M(p^i,g^i))}{2{\sum }_l{M}_{l,b}{\sum }_i{g}_l^i(M_{l,b}-W^M(p^i,g^i))+{\sum }_i{W}^M(p^i,g^i)}$$

对于二值情况，可以设置：

M = [ 0 1 1 0 ] M = [0110]

M=[01​10​]

由此有

$$W^M(p^i,g^i)=|p^i-g^i|,M_{l,b}\to l\ne b$$

此时 wasserstein dice score 就退化为了 soft binary dice score：

D M ( p , q ) = 2 ∑ i g i ( 1 − ∣ p i − g i ∣ ) 2 ∑ i g i ( 1 − ∣ p i − g i ∣ ) + ∑ i ∣ p i − g i ∣ = 2 ∑ i p i g i 2 ∑ i p i g i + ∑ i [ p i ( 1 − g i ) + ( 1 − p i ) g i ] = 2 ∑ i g i p i ∑ i ( g i + p i ) DM(p,q)=2∑igi(1−|pi−gi|)2∑igi(1−|pi−gi|)+∑i|pi−gi|=2∑ipigi2∑ipigi+∑i[pi(1−gi)+(1−pi)gi]=2∑igipi∑i(gi+pi)

DM(p,q)​=2∑i​gi(1−∣pi−gi∣)+∑i​∣pi−gi∣2∑i​gi(1−∣pi−gi∣)​=2∑i​pigi+∑i​[pi(1−gi)+(1−pi)gi]2∑i​pigi​=∑i​(gi+pi)2∑i​gipi​​​

曾经的基于 wasserstein distance 的损失受限于其计算成本，然而，对于这里主要考虑的分割情形中，优化问题的闭式解存在。

对于 $\forall l,l^{\prime }\in \mathbf{L}$，最优传输为 $T_{l,l^{\prime }}=p_l^i{g}_{l^{\prime }}^i$，并且因此 wasserstein distance 可以简化成：

$$W^M(p^i,g^i)=\sum _{l,l^{\prime }\in \mathbf{L}}{M}_{l,l^{\prime }}{p}_l^i{g}_{l^{\prime }}^i$$

wasserstein dice loss

基于 $M$ 可以定义为：

$$L_{D^M}:=1-D^M$$

参考

• 代码：https://github.com/LucasFidon/GeneralizedWassersteinDiceLoss/blob/master/generalized_wasserstein_dice_loss/loss.py