[论文解读] Multi-Task Attention-Based Semi-Supervised Learning for Medical Image Segmentation

Shuai Chen, Gerda Bortsova+3 more|arXiv (Cornell University)|Jul 29, 2019

Advanced Neural Network Applications参考文献 13被引用 13

一句话总结

该论文提出了一种基于多任务注意力的半监督学习方法（MASSL），通过注意力机制联合优化分割与重建任务，实现对前景和背景的分离重建。该方法在仅使用少量标注数据训练时，性能优于完全监督和预训练的CNN模型，交替训练与注意力机制显著提升了脑肿瘤和白质高信号病灶数据集上的特征判别能力与分割精度。

ABSTRACT

We propose a novel semi-supervised image segmentation method that simultaneously optimizes a supervised segmentation and an unsupervised reconstruction objectives. The reconstruction objective uses an attention mechanism that separates the reconstruction of image areas corresponding to different classes. The proposed approach was evaluated on two applications: brain tumor and white matter hyperintensities segmentation. Our method, trained on unlabeled and a small number of labeled images, outperformed supervised CNNs trained with the same number of images and CNNs pre-trained on unlabeled data. In ablation experiments, we observed that the proposed attention mechanism substantially improves segmentation performance. We explore two multi-task training strategies: joint training and alternating training. Alternating training requires fewer hyperparameters and achieves a better, more stable performance than joint training. Finally, we analyze the features learned by different methods and find that the attention mechanism helps to learn more discriminative features in the deeper layers of encoders.

研究动机与目标

• 通过半监督学习利用大量未标注数据，解决标注医学图像有限的挑战。
• 克服标准自编码器重建方法在U-Net架构中的局限性，后者因跳跃连接而无法保留类别特定信息。
• 通过注意力机制引入分割感知的重建，改进深层编码器的特征学习。
• 评估两种多任务训练策略——联合训练与交替训练——对分割性能的影响。
• 证明注意力引导的重建相比标准重建，能提升编码器深层特征的判别性。

提出的方法

• 提出一种多任务框架，共享编码器同时连接分割头（带跳跃连接）与重建头（无跳跃连接），形成类似U-Net的结构。
• 引入注意力机制，从软分割预测中生成独立的前景与背景掩码，实现类别感知的重建。
• 定义重建损失（L2），使用注意力掩码隔离类别特定区域，对前景与背景区域分别应用加权MSE损失。
• 实现两种训练策略：联合训练（同时优化两个任务）与交替训练（按顺序优化，使用任务特定的学习率）。
• 两个任务共享编码器，使重建目标能够引导编码器学习更具判别性的特征，尤其在深层网络层。
• 将方法应用于两项医学影像任务：脑肿瘤（BraTS18）与白质高信号病灶（WMH17）分割，仅使用少量标注图像与大量未标注数据。

实验结果

研究问题

• RQ1 与标准自编码器重建相比，基于注意力的重建目标是否能提升半监督医学图像分割性能？
• RQ2 通过注意力机制整合分割与重建，是否能相比标准多任务学习，使编码器学习到更具判别性的特征？
• RQ3 联合训练与交替训练策略在多任务半监督分割中的性能与稳定性有何差异？
• RQ4 当仅使用少量标注图像时，所提方法在多大程度上优于完全监督CNN与预训练基线？
• RQ5 注意力机制如何影响编码器深层所学特征的质量，以线性判别能力为度量？

主要发现

• 在BraTS18数据集上（50张标注，70张未标注），采用交替训练的MASSL获得0.7111 ± 0.02的Dice分数，显著优于完全监督CNN基线（0.6670 ± 0.03）与MSSL变体（0.6646 ± 0.03）。
• 在WMH17数据集上，采用交替训练的MASSL获得0.7204 ± 0.02的Dice分数，超过完全监督CNN基线（0.6806 ± 0.03）与MSSL基线（0.6880 ± 0.02）。
• 注意力机制显著提升性能：采用交替训练的MASSL在BraTS18上获得0.7111 ± 0.02，而无注意力的MSSL基线仅获0.6670 ± 0.03。
• 交替训练在稳定性和性能上均优于联合训练，在所有数据集上均获得更高Dice分数，且无需对损失加权系数γ进行超参数调优。
• 特征分析显示，注意力机制增强了编码器深层的判别能力：在BraTS18的第5层（最深层）中，MASSL的R²达0.535 ± 0.03，而标准CNN为0.486 ± 0.04。
• 消融实验确认注意力机制是性能提升的关键因素，即使采用相同的训练策略与损失加权，MSSL（无注意力）仍表现逊于MASSL。