Abstract
自监督学习(SSL)为医学图像分析提供了巨大的机会,众所周知,医学图像分析缺乏注释。
然而,由于高成像成本和隐私限制,聚合大量(未标记的)3D医学图像(如计算机断层扫描(CT))仍然具有挑战性。在本文中,我们提倡引入丰富的2D图像,如胸部x光片,以弥补缺乏3D数据,旨在建立一个通用的医学自监督表示学习框架,称为UniMiSS。接下来的问题是如何打破维度障碍,即使2D和3D图像都可以执行SSL ?为了实现这一点,我们设计了一个金字塔形的u形医疗变压器(MiT)。它由可切换贴片嵌入(SPE)模块和变压器组成。SPE模块根据输入尺寸自适应切换到2D或3D补丁嵌入。嵌入的补丁被转换成一个序列,而不考虑其原始尺寸。变形金刚以序列到序列的方式对长期依赖关系进行建模,从而使unimis能够从2D和3D图像中学习表示。以MiT为主干,我们以自蒸馏的方式执行UniMiSS。我们对包括分割和分类在内的六个3D/2D医学图像分析任务进行了昂贵的实验。结果表明,所提出的UniMiSS在各种下游任务上取得了令人满意的性能,大大优于ImageNet预训练和其他高级SSL对象物。代码可从https://github.com/YtongXie/UniMiSS-code获得。
关键词:自监督学习;Cross-dimension;医学图像分析;变压器
1 Introduction
图1所示。(a)纯3D医用SSL仅通过3D CT扫描学习表征;(b)我们提出的UniMiSS带来了丰富的2D x射线,以弥补3D数据的不足,从而使大规模SSL能够获得更好的预训练性能。此外,预训练模型对各种下游应用具有通用性,不受维数障碍的限制。
医学图像分析是计算机辅助诊断的一个关键过程,众所周知,它缺乏训练标签,特别是对于3D任务。最近的研究表明,自监督学习(self-supervised learning, SSL)有望通过充分利用未标记的数据来降低标注成本[8,9,40,41,47,55,57,58]。
虽然去掉了标签,但SSL仍然严重依赖于大规模的未标记数据来探索特征表示。不幸的是,由于成像成本高和数据隐私,公开可用的3D医疗数据相对有限。
大多数3D医疗数据集只包含几千个病例。例如,Zhou等人[55]利用LUNA数据集[38],其中包含约1000个CT病例,进行自监督预训练。如此小的数据规模可能会限制SSL在三维医学图像分析中的潜力。
与3D数据相比,由于成像速度快,辐射低,成本低,因此很容易收集到数十万张二维医学图像,如x射线。因此,我们主张将丰富的2D医学图像引入到3D SSL过程中,旨在通过大尺度图像学习强表征,如图1所示。与纯3D医用SSL相比,这种做法在三个方面有利于医用SSL。首先,2D数据弥补了3D数据的不足,使大规模SSL预训练成为可能。其次,2D和3D图像之间存在解剖学相关性,比如胸部x光片和CT。这种内在关联可能有助于形成强关联表征。第三,预训练模型具有足够的通用性,可以同时应用于3D和2D下游任务。为了实现通用SSL的目的,在技术方面,我们需要构建一个能够处理2D和3D图像的通用模型。医学图像分析中常见的做法是分别为二维图像设计二维卷积神经网络[49,51,55],为三维图像设计三维卷积神经网络[47,48,52,55,57]。由于维数障碍的限制,设计无维CNN网络几乎是不可能的。
最近几个月见证了Transformer在计算机视觉方面的成功[15]。vision Transformer通常以学习到的线性嵌入表示的图像补丁序列作为输入,对序列元素之间的长期依赖关系进行建模。由于序列建模,Transformer可以接受任何维度的数据,包括但不限于2D图像和3D体积数据。因此,Transformer提供了打破维度障碍和构造通用SSL模型的可能性。
在本文中,我们提出了一个通用医学自监督表示学习框架(UniMiSS),它从2D和3D未标记的医学图像中学习一般表示。为此,我们设计了一种无量纲金字塔型医用变压器(MiT),该变压器主要由可切换贴片嵌入(SPE)模块和变压器组成。SPE模块通过使用2D或3D补丁嵌入将输入图像转换为序列,具体取决于输入维度。Transformer层以序列到序列的方式处理嵌入的令牌,而不考虑它们的原始维度。
我们通过学生网络和教师网络的自升华来实现自监督学习,这两个网络都以MiT为主干。学生网络学习预测与动量教师网络得到的输出分布,遵循视图一致性。此外,由于成像内容相同,三维体图像应该与它们的切片相同。采用体片一致性作为跨维正则化来增强表征。我们基于5,022个3D CT体进行了SSL实验