[论文解读] Few-shot brain segmentation from weakly labeled data with deep heteroscedastic multi-task networks

Richard McKinley, Michael Rebsamen+4 more|arXiv (Cornell University)|Apr 4, 2019

Brain Tumor Detection and Classification参考文献 17被引用 6

一句话总结

本文提出一种少样本脑部分割方法，仅使用ABCD研究中的9例弱标签T1加权MRI扫描，通过深度异方差多任务网络建模偶然不确定性。新颖的标签翻转不确定性网络可学习每个体素的误标概率，实现0.87±0.11的平均Dice分数，显著优于标准网络和先前的异方差方法，表明通过不确定性估计实现的模型标签平滑可显著提升性能，即使在极少量且噪声较大的训练数据下亦然。

ABSTRACT

In applications of supervised learning applied to medical image segmentation, the need for large amounts of labeled data typically goes unquestioned. In particular, in the case of brain anatomy segmentation, hundreds or thousands of weakly-labeled volumes are often used as training data. In this paper, we first observe that for many brain structures, a small number of training examples, (n=9), weakly labeled using Freesurfer 6.0, plus simple data augmentation, suffice as training data to achieve high performance, achieving an overall mean Dice coefficient of $0.84 \pm 0.12$ compared to Freesurfer over 28 brain structures in T1-weighted images of $\approx 4000$ 9-10 year-olds from the Adolescent Brain Cognitive Development study. We then examine two varieties of heteroscedastic network as a method for improving classification results. An existing proposal by Kendall and Gal, which uses Monte-Carlo inference to learn to predict the variance of each prediction, yields an overall mean Dice of $0.85 \pm 0.14$ and showed statistically significant improvements over 25 brain structures. Meanwhile a novel heteroscedastic network which directly learns the probability that an example has been mislabeled yielded an overall mean Dice of $0.87 \pm 0.11$ and showed statistically significant improvements over all but one of the brain structures considered. The loss function associated to this network can be interpreted as performing a form of learned label smoothing, where labels are only smoothed where they are judged to be uncertain.

研究动机与目标

• 探究是否仅使用少量弱标签MRI扫描即可实现高性能脑部分割，挑战大规模标注数据集为必需的假设。
• 评估在训练数据稀缺且标签存在噪声的情况下，异方差不确定性建模对分割性能的影响。
• 提出并验证一种新型异方差网络，可对每个体素的标签误分类概率进行估计，实现学习型标签平滑。
• 在少样本、弱监督设置下，比较标准网络、预测方差网络与所提出的标签翻转不确定性网络的性能。

提出的方法

• 在ABCD研究的9例弱标签T1加权MRI体积上训练深层多任务U-Net架构，使用Freesurfer 6.0分割结果作为辅助标签。
• 引入两种异方差变体：一种使用蒙特卡洛Dropout估计预测方差，另一种为新方法，可预测每个体素的标签不一致概率（即标签翻转概率）。
• 标签翻转不确定性网络通过降低高预测误标概率体素的损失，实现学习型标签平滑，有效降低不确定或噪声标签的影响。
• 损失函数结合焦点损失（γ=2）以处理严重的类别不平衡问题，并引入基于预测不确定性的感知损失项，根据预测不确定性调节梯度流动。
• 数据增强包括中线反射、随机旋转和强度重标定，以在数据极少的情况下提升泛化能力。
• 模型推理采用轴向、矢状和冠状视图的3D集成平均，对logits进行平均以生成最终分割结果。

实验结果

研究问题

• RQ1 仅使用九个弱标签训练样本，能否实现高精度脑部分割，而无需人工标注？
• RQ2 在少样本、弱监督设置下，通过异方差网络建模偶然不确定性是否能提升分割性能？
• RQ3 一种可预测标签误分类概率（即标签翻转概率）的新颖不确定性头，是否优于标准预测方差估计？
• RQ4 学习型标签平滑（仅在不确定性高时对标签进行平滑）是否比均匀平滑或不平滑更有效，从而降低噪声引起的误差？

主要发现

• 普通网络在28个脑结构上的平均Dice系数为0.84±0.12，性能在160,000个训练样本左右达到峰值后因过拟合而下降。
• 预测方差异方差网络的平均Dice为0.85±0.14，在28个脑结构中的25个上显著优于普通网络（p<10⁻⁶）。
• 所提出的标签翻转异方差网络达到最高性能，平均Dice为0.87±0.11，在除一个结构外的所有结构上均显著优于普通网络和预测方差模型（p<10⁻⁶）。
• 标签翻转模型表现出更优的训练稳定性，所有脑区（包括难分割结构如伏隔核和第三脑室）均保持高性能，未出现早期性能崩溃。
• 标签翻转图可识别高误标概率体素，显示出作为分割输出质量度量的潜力。
• 初步结果表明，通过精心选择样本，单样本训练可能可行，从而拓展了少样本学习在医学影像中的适用性。