[论文解读] Unpaired Multi-modal Segmentation via Knowledge Distillation

Qi Dou, Quande Liu+2 more|arXiv (Cornell University)|Jan 6, 2020

Domain Adaptation and Few-Shot Learning参考文献 34被引用 10

一句话总结

该论文提出了一种参数高效、无配对的多模态分割框架，该框架在CT和MRI之间共享所有卷积核，同时使用模态特定的批量归一化层。它引入了一种基于预测分布KL散度的新知识蒸馏损失，以对齐跨模态的语义特征，在2D和3D网络上均实现了心脏和腹部多器官分割任务的最先进性能。

ABSTRACT

Multi-modal learning is typically performed with network architectures containing modality-specific layers and shared layers, utilizing co-registered images of different modalities. We propose a novel learning scheme for unpaired cross-modality image segmentation, with a highly compact architecture achieving superior segmentation accuracy. In our method, we heavily reuse network parameters, by sharing all convolutional kernels across CT and MRI, and only employ modality-specific internal normalization layers which compute respective statistics. To effectively train such a highly compact model, we introduce a novel loss term inspired by knowledge distillation, by explicitly constraining the KL-divergence of our derived prediction distributions between modalities. We have extensively validated our approach on two multi-class segmentation problems: i) cardiac structure segmentation, and ii) abdominal organ segmentation. Different network settings, i.e., 2D dilated network and 3D U-net, are utilized to investigate our method's general efficacy. Experimental results on both tasks demonstrate that our novel multi-modal learning scheme consistently outperforms single-modal training and previous multi-modal approaches.

研究动机与目标

• 为解决CT与MRI之间无配对多模态医学图像分割的挑战，特别是当缺乏配对数据和图像配准时。
• 通过利用跨模态知识来提升分割性能，而无需配对训练数据或复杂的网络架构。
• 开发一种高度参数高效的模型，共享模态之间的卷积核，同时通过模态特定的归一化和新型蒸馏损失最小化分布偏移。
• 在真实临床分割任务中验证该方法，涵盖多种网络架构和数据质量条件。

提出的方法

• 该方法在CT和MRI之间共享所有卷积核，减少模型参数并实现参数效率。
• 使用模态特定的内部归一化层（如批量归一化）以适应CT和MRI之间不同的统计分布。
• 通过最小化来自同一输入在不同模态间预测分布的KL散度，引入一种新型知识蒸馏损失。
• 损失在预-softmax特征上计算，以对齐语义表征，促使模型在视觉差异下学习共享的、鲁棒的特征。
• 该框架被应用于2D空洞卷积神经网络和3D U-Net，证明其在不同架构和任务中的泛化能力。
• 该方法通过标准交叉熵损失和所提出的KD损失进行端到端训练，实现模态特定与共享表征的联合优化。

实验结果

研究问题

• RQ1 一个单一、紧凑的神经网络能否在无需图像配准或配对数据的情况下，有效从无配对的CT和MRI数据中学习？
• RQ2 当两种模态具有截然不同的强度分布和视觉外观时，如何有效蒸馏跨模态知识？
• RQ3 在共享所有卷积权重的同时使用模态特定的归一化层，是否能比无配对多模态学习中的独立编码器获得更好的性能？
• RQ4 基于预测分布对齐的知识蒸馏损失，是否能提升在低质量或异常成像场景下的鲁棒性和分割精度？
• RQ5 所提出的方法是否在不同网络架构（2D与3D）和多类别分割任务中具有泛化能力？

主要发现

• 在存在图像伪影的异常情况下，该方法在腹部多器官分割任务中取得了81.1的Dice分数，优于所有基线方法。
• 在同一异常情况下，该方法的Hausdorff距离为3.78，显著低于次优方法（4.89），表明边界定位能力更优。
• 在一般情况下，该方法在心脏和腹部分割任务中均优于单模态训练和最先进多模态方法（包括“Y”形和“X”形网络）。
• 该方法在不同网络架构（如2D空洞CNN和3D U-Net）上均表现出一致的性能提升，证实了其泛化能力。
• 消融实验证实，所提出的知识蒸馏损失对性能至关重要，尤其是在成像条件具有挑战性时。
• 该方法对图像伪影和低质量数据表现出鲁棒性，在其他方法显著退化时仍保持高性能。