[论文解读] A unified representation network for segmentation with missing modalities

Kenneth K. Lau, Jonas Adler+1 more|arXiv (Cornell University)|Aug 19, 2019

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

一句话总结

本文提出一种统一表示网络（URN）与模态丢弃策略，以在输入模态缺失时提升医学影像深度学习分割的性能。URN 能够在不同输入模态之间学习共享表示，从而在数据不完整的情况下仍保持鲁棒的分割能力，且在使用 BRATS 和 HCP 等多模态数据集进行预训练后，其性能优于标准 U-net 及仅使用模态丢弃的方法，尤其在模态缺失场景下表现更优。

ABSTRACT

Over the last few years machine learning has demonstrated groundbreaking results in many areas of medical image analysis, including segmentation. A key assumption, however, is that the train- and test distributions match. We study a realistic scenario where this assumption is clearly violated, namely segmentation with missing input modalities. We describe two neural network approaches that can handle a variable number of input modalities. The first is modality dropout: a simple but surprisingly effective modification of the training. The second is the unified representation network: a network architecture that maps a variable number of input modalities into a unified representation that can be used for downstream tasks such as segmentation. We demonstrate that modality dropout makes a standard segmentation network reasonably robust to missing modalities, but that the same network works even better if trained on the unified representation.

研究动机与目标

• 为解决测试时输入模态缺失这一现实问题，该问题违反了标准训练-测试分布假设。
• 开发一种方法，避免为每种可能的输入模态组合单独训练网络，从而规避组合爆炸问题。
• 通过无监督预训练与统一表示学习，提升在模态缺失场景下的分割鲁棒性与性能。
• 评估模态丢弃与统一表示网络是否能增强泛化能力并降低对缺失数据的敏感性。

提出的方法

• 模态丢弃在训练过程中随机将整个输入通道置零，迫使网络学习模态无关的特征，从而提升对缺失输入的鲁棒性。
• 统一表示网络（URN）为每种输入模态使用独立的 U-net 编码器，通过可学习的编码器将各模态映射到共享的潜在表示空间。
• URN 在具有部分模态重叠的数据集（如 BRATS 和 HCP）上以无监督方式预训练，学习共享表示，且在预训练阶段无需使用分割标签。
• 下游分割通过在统一表示上训练的解码器头完成，支持不同输入模态组合的端到端分割。
• 使用截断几何分布采样模态丢弃的数量：$ p_\theta(k) = \frac{(1-\theta)\theta^k}{1-\theta^{N+1}} $，其中 $ k = 0, \dots, N_{\text{max}} $。
• 预训练在 HCP 数据集中 1108 名健康受试者的 T1 和 T2 MR 扫描上进行，每名受试者及每种模态均应用了强度归一化与偏置场校正。

实验结果

研究问题

• RQ1 模态丢弃能否提升标准 U-net 在医学图像分割中对缺失输入模态的鲁棒性？
• RQ2 在多种模态上学习统一表示是否能带来优于仅使用模态丢弃的分割性能？
• RQ3 在多个具有部分模态重叠的数据集上进行无监督预训练，能否提升下游分割性能？
• RQ4 在不完整输入组合的测试中，统一表示网络是否比标准模型更有效？

主要发现

• 模态丢弃显著提升了标准 U-net 对缺失输入模态的鲁棒性，有效缓解了测试时模态缺失导致的性能下降。
• 统一表示网络（URN）在所有输入组合下均实现了全肿瘤分割的最高 Dice 评分，优于基线 U-net 及仅使用模态丢弃的变体。
• 在合并的 BRATS 与 HCP 数据上对 URN 进行预训练，提升了增强肿瘤与肿瘤核心区域的分割性能，尽管全肿瘤分割性能与未预训练的 URN 相当。
• 即使所有模态均存在，使用模态丢弃进行训练仍能提升基线 U-net 与 URN 的性能，表明其具有正则化作用，能促进模态间特征共享。
• 在 BRATS 中，FLAIR 被发现是全肿瘤分割中最具信息量的模态，因其缺失导致基线模型性能下降最大。
• 定性结果表明，仅使用 T1c 模态，URN 也能可靠分割所有肿瘤结构，尽管在临床实践中 T1c 主要用于增强肿瘤的显示。