[论文解读] Revisiting Rubik's Cube: Self-supervised Learning with Volume-wise Transformation for 3D Medical Image Segmentation

Tao Xing, Yuexiang Li+3 more|arXiv (Cornell University)|Jul 17, 2020

Advanced Neural Network Applications参考文献 19被引用 6

一句话总结

本文提出 Rubik’s cube++，一种新颖的自监督学习框架，用于3D医学图像分割，通过体素级变换（如3D体素旋转）对3D U-Net架构的编码器和解码器组件进行预训练。通过将上下文恢复定义为先验任务，并采用空间一致的子立方体置换，该方法在无需额外标注数据的情况下，显著提升了胰腺和脑组织数据集上的分割精度，实现了最先进性能。

ABSTRACT

Deep learning highly relies on the quantity of annotated data. However, the annotations for 3D volumetric medical data require experienced physicians to spend hours or even days for investigation. Self-supervised learning is a potential solution to get rid of the strong requirement of training data by deeply exploiting raw data information. In this paper, we propose a novel self-supervised learning framework for volumetric medical images. Specifically, we propose a context restoration task, i.e., Rubik's cube++, to pre-train 3D neural networks. Different from the existing context-restoration-based approaches, we adopt a volume-wise transformation for context permutation, which encourages network to better exploit the inherent 3D anatomical information of organs. Compared to the strategy of training from scratch, fine-tuning from the Rubik's cube++ pre-trained weight can achieve better performance in various tasks such as pancreas segmentation and brain tissue segmentation. The experimental results show that our self-supervised learning method can significantly improve the accuracy of 3D deep learning networks on volumetric medical datasets without the use of extra data.

研究动机与目标

• 为解决深度学习中3D医学图像标注数据有限这一关键瓶颈问题。
• 通过利用固有的3D解剖结构，采用自监督先验任务对深层网络进行预训练，以提升3D医学图像分割性能。
• 克服先前自监督方法仅预训练编码器层的局限性，通过上下文恢复任务联合预训练编码器和解码器。
• 验证体素级变换（如3D体素旋转）相较于补丁级或网格基变换，能生成更鲁棒的特征表示。

提出的方法

• 该方法将3D体积分割为大小为 n×n×n 的子立方体网格，形成类似魔方的结构。
• 通过沿矢状面、冠状面或轴面将整个立方体层旋转90°、180°或270°，施加体素级变换，保持空间一致性并增加任务难度。
• 先验任务为上下文恢复：网络需预测变换后3D体素的原始空间配置，从而鼓励学习3D结构先验。
• 该框架应用于全卷积网络（FCNs），如3D U-Net，实现编码器和解码器路径的联合预训练。
• 采用对比损失进行端到端训练，以鼓励从变换版本中正确重建原始体素。
• 通过在公开数据集（包括 Pancreas-101 和 MRBrainS18）上进行交叉验证评估该方法，主要指标为Dice相似系数（DSC）。

实验结果

研究问题

• RQ1 与补丁级或网格基变换相比，上下文恢复任务中的体素级变换是否能提升3D解剖特征学习？
• RQ2 通过上下文恢复任务联合预训练编码器和解码器，是否能带来优于仅预训练编码器的3D医学图像分割性能？
• RQ3 先验任务的难度（通过子立方体大小 n 控制）如何影响所学表征质量和下游分割精度？
• RQ4 所提出的 Rubik’s cube++ 方法在3D医学图像分割任务中是否在统计上显著优于现有自监督基线方法（如 Models Genesis 和 Rubik’s cube [20]）？

主要发现

• 在脑组织分割任务中，Rubik’s cube++ 相较于从零开始训练，平均Dice相似系数（DSC）提升了5.34%，优于UCF101预训练和其他自监督基线方法。
• 在胰腺分割任务中，该方法在100%训练数据下达到84.08%的DSC，显著优于次佳基线方法（83.72%），且在所有数据设置下均表现出一致的性能增益。
• t检验结果表明，与 Models Genesis 的性能提升在5%显著性水平下具有统计显著性（p值 = 3.42%），验证了该方法的鲁棒性。
• 最优子立方体大小 n=7 在任务难度与特征学习之间实现了最佳平衡，DSC在50%训练数据下达到82.80%，在100%数据下达到84.08%。
• 当 n 增加至9时，性能下降，表明过度破坏3D结构会损害解剖信息提取，降低表征鲁棒性。
• 该方法表明，通过上下文恢复任务联合预训练编码器和解码器，可实现更优的分割性能，而其他仅预训练编码器的方法则不具备此优势。