[论文解读] Joint Learning of Brain Lesion and Anatomy Segmentation from Heterogeneous Datasets

Nicolas Roulet, Diego Fernández Slezak+1 more|arXiv (Cornell University)|Mar 8, 2019

Advanced Neural Network Applications参考文献 15被引用 5

一句话总结

本文提出一种自适应交叉熵（ACE）损失函数，用于联合训练单个U-Net模型，实现基于异质数据集（标注不重叠）的脑部病变与解剖结构分割。ACE损失通过动态加权类别特定的交叉熵损失，解决病变与组织标签之间的标签矛盾，从而在标准交叉熵和Dice损失失效的情况下实现有效的联合学习。该方法在脑组织分割方面表现优于多网络方法，且在多种病变类型下保持鲁棒性。

ABSTRACT

Brain lesion and anatomy segmentation in magnetic resonance images are fundamental tasks in neuroimaging research and clinical practice. Given enough training data, convolutional neuronal networks (CNN) proved to outperform all existent techniques in both tasks independently. However, to date, little work has been done regarding simultaneous learning of brain lesion and anatomy segmentation from disjoint datasets. In this work we focus on training a single CNN model to predict brain tissue and lesion segmentations using heterogeneous datasets labeled independently, according to only one of these tasks (a common scenario when using publicly available datasets). We show that label contradiction issues can arise in this case, and propose a novel adaptive cross entropy (ACE) loss function that makes such training possible. We provide quantitative evaluation in two different scenarios, benchmarking the proposed method in comparison with a multi-network approach. Our experiments suggest that ACE loss enables training of single models when standard cross entropy and Dice loss functions tend to fail. Moreover, we show that it is possible to achieve competitive results when comparing with multiple networks trained for independent tasks.

研究动机与目标

• 解决使用异质数据集（标注不重叠）联合训练单个深度学习模型进行脑部病变与解剖结构分割的挑战。
• 识别并解决当病变与组织标签在不同数据集中冲突时产生的标签矛盾问题，特别是当一个数据集中的背景在另一数据集中被标记为组织时的情况。
• 开发一种损失函数，使单个U-Net模型在存在此类冲突标注的情况下仍能实现有效的联合训练。
• 在不同病变类型下的分割精度与泛化能力方面，评估所提方法与多网络方法及标准损失函数方法的性能。
• 证明即使在输入模态有限的情况下，使用ACE损失进行联合学习的性能仍可与针对特定任务设计的多网络模型相媲美。

提出的方法

• 提出一种自适应交叉熵（ACE）损失函数，根据标签在不同数据集中的频率与一致性动态调整类别权重，以解决标签矛盾问题。
• ACE损失计算类别特定的交叉熵项，并应用自适应加权以最小化冲突标注的影响，尤其针对背景与组织类别。
• 在合并的异质数据集上端到端训练单个U-Net模型，其中每个数据集仅针对一项任务（病变或解剖结构）进行标注。
• 采用多任务学习框架，使同一网络同时预测病变与组织分割结果，共享任务间的特征表示。
• 使用公开数据集中的数据，如WMH（白质高信号）、BRATS（肿瘤）和T1加权MRI（用于解剖结构），并进行仔细的预处理以对齐空间与强度域。
• 通过消融研究比较ACE损失与标准分类交叉熵损失及Dice损失，并与采用独立网络的Multi-UNet基线模型进行性能评估。

实验结果

研究问题

• RQ1 当训练数据来自标注不重叠的异质数据集时，能否有效训练单个深度神经网络，实现脑部病变与解剖结构的联合分割？
• RQ2 在使用标准损失函数（如交叉熵或Dice损失）时，病变与组织标注之间的标签矛盾如何影响训练过程？
• RQ3 所提出的ACE损失函数在多大程度上缓解了标签矛盾问题，并提升了联合学习场景下的分割性能？
• RQ4 单个ACE训练模型的性能与采用独立模型进行病变与解剖结构分割的多网络方法相比如何？
• RQ5 当输入模态被限制为病变与解剖结构数据集共有的模态时，联合学习使用ACE方法存在哪些局限性？

主要发现

• ACE损失成功实现了单个U-Net模型在病变与解剖结构联合分割任务上的训练，而标准交叉熵损失与Dice损失因标签矛盾而失败。
• 在WMH分割方面，单个ACE模型的性能与Multi-UNet基线模型相当，Wilcoxon signed-rank检验未发现显著差异。
• 在脑组织分割方面，ACE模型在多个结构上的表现优于或至少不逊于Multi-UNet，表明通过联合优化可实现更优的特征学习。
• ACE模型在水肿与肿瘤分割方面性能下降，主要由于输入模态受限（仅使用T1图像），而Multi-UNet则利用了多种序列（T1、T1g、T2、FLAIR）。
• 该方法展示了层次化多任务学习的可行性，通过联合预测病变与组织类别，提升了整体分割精度，捕捉了标签间的相互作用。
• 作者观察到ACE损失能有效处理类别不平衡与标签冲突问题，表明其在多任务医学图像分割中的广泛应用潜力。