各种 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$ 可以被定义为标签空间 $\mathbf{L}=\{0,1\}$ 上的随机变量。

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

• $p=\{p^i:=P(S_i=1){\}}_{i\in \mathbf{X}}$
• $g=\{g^i:=P(G_i=1){\}}_{i\in \mathbf{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$，即 $\forall i\in \mathbf{X},g^i\in \{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 可以被定义为一个线性规划问题的解。

$$\begin{array}{rl}{W}^M(p,q)& =\underset{T_{l,l^{\prime }}}{\min }\sum _{l,l^{\prime }\in \mathbf{L}}{T}_{l,l^{\prime }}{M}_{l,l^{\prime }}\\ \text{subject to }\forall l\in \mathbf{L},\sum _{l^{\prime }\in \mathbf{L}}{T}_{l,l^{\prime }}& =p_l,\\ \text{ and }\forall l^{\prime }\in \mathbf{L},\sum _{l\in \mathbf{L}}{T}_{l,l^{\prime }}& =q_{l^{\prime }}\end{array}$$

这里的 $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)$
• ${\Theta }_{TP}^l={\sum }_{i\in \mathbf{X}}{g}_l^i(M_{l.b}-W^M(p^i,g^i)),\forall l\in \mathbf{L}\setminus \{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=\left[\begin{array}{cc}0& 1\\ 1& 0\end{array}\right]$$

由此有

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

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

$$\begin{array}{rl}{D}^M(p,q)& =\frac{2{\sum }_i{g}^i(1-|p^i-g^i|)}{2{\sum }_i{g}^i(1-|p^i-g^i|)+{\sum }_i|p^i-g^i|}\\ & =\frac{2{\sum }_i{p}^i{g}^i}{2{\sum }_i{p}^i{g}^i+{\sum }_i[p^i(1-g^i)+(1-p^i)g^i]}\\ & =\frac{2{\sum }_i{g}^i{p}^i}{{\sum }_i(g^i+p^i)}\end{array}$$

曾经的基于 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