用复原魔方的思想对3D数据进行自监督学习（MICCAI 2020）
用复原魔方的思想对3D数据进行自监督学习（MICCAI 2020）
今天分享一篇发表在MICCAI 2020上的论文：Revisiting Rubik’s Cube: Self-supervised Learning with Volume-wise Transformation for 3D Medical Image Segmentation (原文链接：[1])。
1 研究背景
现有的大部分深度学习算法都高度依赖于标签数据的质量，然而对于医学影像中的3D数据进行标注，不止需要大量的人工成本同时还需要一定的专业知识。自监督学习是一种摆脱对大量有标注的训练数据高度依赖的方法。这篇文章提出了一种新颖的用于3D医学影像数据的自监督方法，利用这种自监督方法能够对3D神经网络进行预训练，而后只需少量数据和修改就可以进行迁移学习。
2 方法
2.1 问题定义
3D的医学影像数据可以看成一个由一系列立方体 (Subcube)构成的高阶魔方，假设3D数据的维度为
，每一个立方体定义为包含三维解剖信息的最小组件，它们的棱长为
个体素，那么3D数据可以描述为由
个立方体构成的魔方，其中
表示向下取整操作。
根据魔方的规则，将在同一解剖平面上的一组子方体组成的面称为立方体层 (Cube layer)，这些层可以沿着大轴(矢状面、冠状面和轴向面)进行特定角度的旋转变化，具体如下图所示。
2.2 自监督：恢复魔方
根据上一小节的定义，3D的医学影像数据可以看成是由一系列立方体构成的魔方，根据魔方的规则随机在矢状面、冠状面和轴向面上对立方体层进行旋转，就可以获得打乱后的3D数据。之监督任务的目标就是将3D数据从打乱后的无序状态恢复到打乱前的初始状态。
具体算法可见下图，算法核心是如下两步：
- 迭代选取矢状面、冠状面和轴向面中的一个作为基准方位
- 对于选定的方位，选取
个立方体层 (cube layers)，随机旋转
角度
2.3 网络结构
文章中采用GAN结构的框架来学习魔方的复原，框架中包含了一个生成器
和一个判别器
。如下图 (Fig.2)所示，生成器由3D的编码器-解码器以及跳层连接构成，判别器由4个卷积核为
的3D卷积层构成。将生成器还原后的状态
和初始状态
分别与打乱后的无序状态
一起送入判别器进行真/假分类。
2.4 损失函数
整个框架的损失函数由两部分构成：重建损失函数 (Reconstruction Loss)和对抗损失函数 (Adversarial Loss)。
2.4.1 重建损失函数
如上图 (Fig.2)所示，将初始状态
根据魔方规则打乱得到的无序状态
作为生成器
的输入，产生的输出记为
，通过计算
和
之间体素的
损失来优化重建质量：
注：文中提到尝试过采用L2损失函数，但发现
损失函数重建出来的图像模糊，很容易丢失掉器官的边缘信息。
2.4.2 对抗损失函数
对抗损失函数用来确保生成器重建的图像能够与真实图像尽可能地接近。生成器学习从无序状态到初始状态的映射
，判别器学习真（原数据）假（生成器生成的数据）分类，损失函数表达如下，
2.4.3 学习目标
在训练过程中，鼓励生成器生成能够糊弄判别器的图像，同时判别器需要尽可能正确地分类真假数据。因此生成器
期望最小化
和
，同时判别器
期望最大化
。综上，整个模型的学习目标可以表示为，
其中
在文章的实验中被设置为10.
2.5 迁移学习
框架在魔方恢复任务中完成自监督学习后，生成器已经能够学到3D医学影像数据中的解剖信息，再通过将生成器的最后一层替换为分割输出层，进行微调后即可适应目标任务。
3 实验结果
这里我只给出论文中的部分实验结果，具体的实验结果分析以及实验和参数的设置请看原文。
[1] https://arxiv.org/pdf/2007.08826