[论文解读] Generalised Wasserstein Dice Score for Imbalanced Multi-class Segmentation using Holistic Convolutional Networks

Lucas Fidon, Wen-Zhi Li+5 more|UCL Discovery (University College London)|Jul 3, 2017

Artificial Intelligence in Healthcare被引用 16

一句话总结

本文提出了一种新颖的广义Wasserstein Dice评分，通过标签空间树和全卷积神经网络（多尺度监督）来整合类别间的语义关系，以改善类别不平衡的多类别脑肿瘤分割。该方法在罕见肿瘤亚型上的分割结果比基线损失函数和网络架构更具语义意义，且Dice分数更高。

ABSTRACT

The Dice score is widely used for binary segmentation due to its robustness to class imbalance. Soft generalisations of the Dice score allow it to be used as a loss function for training convolutional neural networks (CNN). Although CNNs trained using mean-class Dice score achieve state-of-the-art results on multi-class segmentation, this loss function does neither take advantage of inter-class relationships nor multi-scale information. We argue that an improved loss function should balance misclassifications to favour predictions that are semantically meaningful. This paper investigates these issues in the context of multi-class brain tumour segmentation. Our contribution is threefold. 1) We propose a semantically-informed generalisation of the Dice score for multi-class segmentation based on the Wasserstein distance on the probabilistic label space. 2) We propose a holistic CNN that embeds spatial information at multiple scales with deep supervision. 3) We show that the joint use of holistic CNNs and generalised Wasserstein Dice scores achieves segmentations that are more semantically meaningful for brain tumour segmentation.

研究动机与目标

• 通过利用肿瘤亚型之间的语义关系，解决多类别脑肿瘤分割中的类别不平衡问题。
• 通过新型全卷积神经网络架构引入多尺度空间特征，提升分割鲁棒性。
• 开发一种可微分损失函数，以编码层次化的标签结构，实现更好的泛化能力。
• 证明将语义感知损失与多尺度监督相结合可获得更优的分割性能。

提出的方法

• 基于标签概率向量之间的Wasserstein距离，提出一种广义Dice评分，利用预定义的标签空间距离矩阵来编码语义关系。
• 设计一种全卷积神经网络（HCNN），可生成多尺度预测，并在每个尺度上应用深度监督，以改善特征学习。
• 采用两阶段训练策略：先使用均值软Dice损失进行预训练，再使用所提出的Wasserstein Dice损失进行微调，以稳定训练过程。
• 使用树状结构的标签空间，人工分配边权重以反映肿瘤类型之间的解剖学和病理学关系（例如，坏死核心、非增强、增强区域）。
• 应用Wasserstein距离同时比较所有类别的概率标签图，实现类别间误分类惩罚的联合优化。
• 通过学习或固定的融合策略，融合多尺度预测结果，以结合高分辨率细节与大感受野的上下文信息。

实验结果

研究问题

• RQ1 基于Wasserstein的Dice评分广义化是否能提升在类别不平衡的多类别医学图像分割中的性能？
• RQ2 通过距离矩阵引入层次化标签关系在多大程度上影响模型的泛化能力和语义一致性？
• RQ3 在脑肿瘤MRI中，通过深度监督实现的多尺度特征学习在多大程度上提升了分割精度？
• RQ4 两阶段训练策略（先使用均值Dice损失预训练，再使用Wasserstein Dice损失微调）是否相比仅使用任一损失的端到端训练，能带来更好的收敛性和性能？

主要发现

• 所提出的广义Wasserstein Dice损失（$\mathcal{L}_{D^{M_{\text{tree}}}}$）在罕见类别（如坏死核心和非增强肿瘤）上的Dice分数高于均值软Dice损失。
• 采用多尺度监督与深度监督的全卷积神经网络（HCNN）相比标准U-Net类架构，产生了更一致且更准确的分割结果。
• HCNN的融合预测结果达到整体肿瘤Dice分数89%（±6），核心肿瘤73%（±22），增强肿瘤74%（±23），以及$D^{M_{\text{tree}}}$为87%（±4），优于单一尺度和基线方法。
• 先使用均值Dice损失预训练，再使用Wasserstein Dice损失微调（$\mathcal{L}_{D^{M_{\text{tree}}}-\text{PT}}$）的策略实现了稳定的训练过程，并相比直接使用$\mathcal{L}_{D^{M_{\text{tree}}}}$训练，提高了真正例率。
• 第1个尺度因感受野不足，Dice分数较低；而第3和第4个尺度因分辨率降低，未能捕捉到细小的肿瘤区域，凸显了尺度与分辨率之间的权衡。
• 定性结果表明，仅融合预测和第2个尺度的预测正确分割了增强肿瘤区域，该区域为最小且最关键的结构。