[论文解读] Self-supervised Feature Learning for 3D Medical Images by Playing a Rubik's Cube

Xinrui Zhuang, Yuexiang Li+4 more|arXiv (Cornell University)|Oct 5, 2019

Medical Imaging and Analysis参考文献 12被引用 15

一句话总结

该论文提出了一种用于医学图像的自监督3D特征学习框架，采用一种新颖的代理任务——魔方复原，通过网络重建被置换和旋转的3D体素立方体，学习平移和旋转不变性。该方法在脑出血分类任务（83.8%准确率）和脑肿瘤分割任务（77.3% mIoU）中均达到最先进性能，显著优于从零开始训练和在自然视频数据集上预训练的方法。

ABSTRACT

Witnessed the development of deep learning, increasing number of studies try to build computer aided diagnosis systems for 3D volumetric medical data. However, as the annotations of 3D medical data are difficult to acquire, the number of annotated 3D medical images is often not enough to well train the deep learning networks. The self-supervised learning deeply exploiting the information of raw data is one of the potential solutions to loose the requirement of training data. In this paper, we propose a self-supervised learning framework for the volumetric medical images. A novel proxy task, i.e., Rubik's cube recovery, is formulated to pre-train 3D neural networks. The proxy task involves two operations, i.e., cube rearrangement and cube rotation, which enforce networks to learn translational and rotational invariant features from raw 3D data. Compared to the train-from-scratch strategy, fine-tuning from the pre-trained network leads to a better accuracy on various tasks, e.g., brain hemorrhage classification and brain tumor segmentation. We show that our self-supervised learning approach can substantially boost the accuracies of 3D deep learning networks on the volumetric medical datasets without using extra data. To our best knowledge, this is the first work focusing on the self-supervised learning of 3D neural networks.

研究动机与目标

• 为解决在训练深度学习模型进行医学诊断时，3D医学图像标注数据有限的挑战。
• 开发一种自监督预训练框架，利用未标注的3D体积分块数据，提升下游任务的性能。
• 提出一种新颖的代理任务——魔方复原，强制3D卷积神经网络学习对平移和旋转的空间不变性。
• 证明仅使用医学数据进行预训练，无需外部数据集，其性能优于在ImageNet或UCF101上预训练的模型。

提出的方法

• 该方法将3D医学体积分割为3D网格（例如2×2×2）的子体积（立方体）。
• 对立方体应用随机置换（立方体重排）和随机旋转（立方体旋转），生成原始体积的损坏版本。
• 代理任务是重建立方体的原始配置——包括空间排列顺序和方向。
• 使用对比学习或分类损失，训练3D CNN预测正确的置换和旋转。
• 将预训练的编码器在下游任务（如脑出血分类和肿瘤分割）上进行微调，仅使用有限的标注数据。
• 该框架与基线方法进行对比，包括从零开始训练和使用UCF101预训练权重进行微调。

实验结果

研究问题

• RQ1 基于3D空间操作的自监督代理任务是否能改善体积分块医学图像的特征学习？
• RQ2 仅在医学数据上进行预训练，不使用外部数据集，其性能是否优于在自然视频数据（如UCF101）上预训练的模型？
• RQ3 与仅考虑立方体排序相比，代理任务中引入立方体旋转如何影响所学特征的泛化能力和鲁棒性？
• RQ4 魔方复原任务在多大程度上提升了下游3D医学图像任务（如分类和分割）的性能？

主要发现

• 所提出的魔方复原预训练方法在脑出血分类任务中达到83.8%的准确率，较从零开始训练（72.6%）提升了11.2个百分点。
• 在脑肿瘤分割任务中，该方法使用3D DUC模型达到77.3% mIoU，较仅使用立方体排序的基线模型提升了2.3个百分点。
• 从魔方复原预训练权重进行微调，在分类任务上比使用UCF101预训练权重微调高出+2.7%，在分割任务上高出+1.1%。
• 在代理任务中引入立方体旋转，使学习到的特征更具鲁棒性，具备更强的平移和旋转不变性，泛化能力优于仅使用立方体排序的情况。
• 3D DUC架构在所有预训练策略下均优于U-Net，表明其对随机解码器初始化的退化效应更小。
• 该方法是首个证明仅使用未标注医学数据即可实现有效自监督预训练的3D CNN方法，适用于体积分块医学图像。