【多模态（影像）自监督学习】UniMiSS: Universal Medical Self-Supervised Learning via Breaking Dimensionality Barrie

最新推荐文章于 2025-09-14 10:42:19 发布

原创最新推荐文章于 2025-09-14 10:42:19 发布 · 2k 阅读

20 ·

本内容遵循CC 4.0 BY-SA版权协议

GEO检测

收录于

深度学习

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实验，这些CT体由108,948张2D x射线图像增强。得益于庞大的增强2D数据，本文提出的UniMiSS在下游3D分类/分割任务上实现了明显的性能提升。此外，UniMiSS的预训练模型可以自由地应用于2D下游任务，这击败了ImageNet等强大的竞争对手在2D下游医疗任务上的预训练。

综上所述，我们的贡献有三个方面:(1)我们是第一个将3D医学图像与SSL目的的易于访问的未配对2D图像增强，旨在解决SSL过程中3D数据量的限制;(2)所提出的MiT打破了维度障碍，实现了二维和三维图像的SSL联合训练;(3)我们的UniMiSS预训练在六个下游任务上达到了先进的性能，涵盖了3D/2D医学图像的分类/分割。

2 Related Work

2.1 Self-supervised Learning

最近，Transformer已经成为计算机视觉中越来越受欢迎的替代架构。将变压器和SSL的优点结合起来，推进自监督视觉变压器已经成为一种趋势。开创性的工作是iGPT[10]，它遵循屏蔽自回归语言建模来预训练自监督视觉变压器。此外，也有人尝试使用对比学习[7]或暹罗蒸馏[13]对视觉变形器进行预训练，其性能优于基于cnn的SSL方法，在ImageNet上创造了新的记录。

SSL在计算机视觉方面的成功也使医学界受益[8,9,40,41,47,55,57,58]。典型的尝试包括通过恢复原始图像的内容来预训练CNN[9,41,55,57,58]，以及为医学图像定制对比SSL[8,40,41,47]。这些努力是朝着更好的SSL医学图像分析方法迈出的重要和及时的一步。然而，它们有两个限制。首先，CNN架构只能对二维或三维医学图像进行预训练，不能同时对两者进行处理。生成的表示将被困住，特别是对于有限的3D数据。因此，预训练的CNN只能转移到特定维度的下游任务。其次，上述SSL方法从切片[8]或体[41,57,58]中捕获三维医学图像的空间背景。很少有人考虑体积与其切片之间的内在一致性关系。

2.2 Cross-Domain Training for Medical Imaging

在医学背景下，跨域训练通常联合使用在不同地点获取的两个或多个数据集[24,30]或使用不同的成像方式[16,29,54]来训练单个模型，该模型可以在不同的数据集上表现良好。Karani等人[24]和Liu等人[30]分别使用在每个站点获取的MRI数据训练了具有共享卷积层和特定批处理归一化层的单个CNN，旨在明确解决统计差异。Zhang等人[54]同时使用未配对的CT和MRI数据学习了体到体的转换，并使用从另一种模态翻译的合成数据学习了强分割器。Dou等人[16]推导了一种知识蒸馏(KD)的变体，利用CT和MRI之间共享的跨模态信息来准确分割解剖结构。Li等[29]也将KD引入到CT和MRI数据的交叉模态分析中，但他们同时利用了大量未标记的数据。这些研究致力于分析多模态/地点但固定维度(3D)的医学图像，未能解决我们场景中的维度障碍。

3 Methods
3.1 Overview

UniMiSS是一个通用的医学SSL框架，它优于学习大规模混合二维和三维无标记医学图像的一般图像表示。图2说明了联苏特派团的管道。让我们用{D 2D, D 3D}表示混合的2D和3D数据池。为了使UniMiSS能够同时处理二维和三维医学图像，我们构建了MiT作为其主干，该主干主要由自适应维度的SPE模块和Transformer层组成。我们以自蒸馏的方式执行SSL过程，并利用标准的交叉熵损失来最大限度地提高学生和教师输出之间的一致性。此外，为了最大限度地利用三维体积信息，我们引入了体片一致性约束，鼓励UniMiSS对跨维度一致性进行建模。直观上有利于从体积图像中学习强特征表示。现在我们深入研究这个框架的细节。

图2所示。说明拟议的联苏支助团框架。它具有双路径架构，即学生和教师。以2D x射线和3D ct作为输入，UniMiSS通过自蒸馏策略进行训练，即最大化两条路径的一致性。为了打破x射线和ct之间的维度障碍，麻省理工学院的网络，由可切换贴片嵌入(SPE)模块和变压器组成，以序列对序列的方式处理3D/2D数据。

3.2 MiT: A Dimension-free Architecture

虽然在计算机视觉方面取得了巨大的成功，但由于计算成本和内存要求高，视觉变压器[15]在处理高分辨率3D图像方面仍然具有挑战性。受[45]的启发，我们设计了一个金字塔结构的麻省理工学院，以有效地处理二维和三维图像。为了打破维度障碍，我们提出了一个简单而高效的SPE模块，可以根据输入类型自适应选择2D或3D补丁嵌入。

MiT有一个编码器-解码器架构，可以促进各种应用，包括分割和分类。我们现在描述MiT的每个部分，更多的细节可以在附录中找到。

SPE。如图2所示，SPE模块在获得特定维度嵌入方面起着重要作用，即对2D输入使用2D补丁嵌入操作，对3D输入使用3D补丁嵌入操作。注意，编码器和解码器中的SPE实现是不同的。编码器中的SPE是指一个可切换的2D和3D卷积块，其步幅为2，这降低了特征分辨率。相比之下，解码器中的SPE是一个可切换的2D和3D转置卷积块，这增加了特征分辨率。

Encoder-Decoder。MiT编码器遵循一个渐进收缩的金字塔变压器，如[45]所做的那样。它明确地由四个级组成，每个级由一个SPE模块和几个堆叠的变压器组成。在每个阶段，SPE模块对输入特征进行下采样，并生成特定维度的嵌入序列。值得注意的是，我们在补丁嵌入序列中附加了一个额外的可学习SSL令牌[7,13]。SSL令牌类似于ViT中的[CLS]令牌，它能够通过自关注聚合来自整个补丁嵌入令牌的信息。所得到的序列，结合可学习的位置嵌入，被输入到以下的transformer中进行长期依赖建模。每个Transformer层包括一个自关注模块和一个具有两个隐藏层的前馈网络(FFN)。

为了使MiT能够处理高分辨率图像，我们遵循了空间减少注意(SRA)层[45]。给定一个查询q、一个键k和一个值v作为输入，SRA首先降低k和v的空间分辨率，然后将q、约简k和约简v馈送到一个多头自关注(MSA)层，以产生精细化的特征。这个过程可以正式表示如下

其中σ(·)表示降低特征图分辨率的线性投影，即跨行二维或三维卷积运算，R(·)将输入序列重塑为原始空间大小的特征图，F(·)将输入平坦化为一维序列。MiT有一个对称的解码器结构，由三个阶段组成。在每个阶段，输入特征映射首先由SPE模块上采样，然后由堆叠的Transformer层进行细化。此外，我们还在编码器和解码器之间添加了跳过连接，以保留更多低层次但高分辨率的信息。

3.3 Objective of UniMiSS

拟议的综合管理信息系统框架以学生-教师模式为基础。

每条路径由一个MiT网络Fθ(·)和一个投影仪Pθ(·)组成。Pθ(·)是n层多层感知器(MLP)头部，θ表示该路径的参数集。在前馈计算过程中，SPE层可以切换到二维或三维补丁嵌入，表示为Fθ(·;2D)和Fθ(·;分别3 d)。在SSL过程中，我们只从Fθ(·;2D/3D)作为投影仪的输入。由于Transformer使维度自由，我们的UniMiSS能够从未标记的2D和3D医学图像中学习图像表示。

两条路径共享相同的架构。然而，它们在以下两个项目上有所不同。首先，将教师网络表述为学生网络的动量版本，学生网络通过指数移动平均策略更新，定义为

其中λ在训练过程中使用余弦调度从0.996增加到1[7]。

其次，对教师网络进行停止梯度算子，避免模型崩溃。

目的:获取二维域数据。以一小批2D数据x为例，我们首先使用数据增强模块T创建两个增强视图x1和x2，然后将它们输入到学生和教师网络中。

将得到的SSL令牌输入到投影仪中产生输出矢量，记为f1 = Pθ(Fθ(x1;2D))， f2 = Pµ(Fµ(x2;2 d))。UniMiSS的目标是最大限度地提高学生和教师网络的输出向量之间的一致性

其中C为教师输出的中心，τt和τs为学生和教师网络的锐化温度参数。定心操作使模型趋于均匀分布，而锐化则相反，使一维占主导地位。两者共同使用，避免模型崩溃[7]。具体来说，为了锐化的目的，在教师路径中温度τt被设置为一个小的值。中心C首先通过平均最小批数据的教师输出来计算，然后使用指数移动平均策略进行更新，以便在整个批次中聚合中心，如下所示

其中ω为速率参数，f珞2为小批量教师输出的均值。我们将二维图像的对称损失定义为:

目的:获取三维领域数据。在医学领域，三维体可以看作是二维图像随层间维数的叠加。卷数据与其片具有固有的一致性，这启发我们对SSL的卷片一致性进行建模。给定从3D医学数据集中采样的3D数据x，我们将其两个增强视图表示为x1和x2，每个视图包含m个2D切片。我们在3D模式下计算学生和教师网络的全局体积表示，即f1 = Pθ(Fθ(x1;3D))， f2 = Pµ(Fµ(x2;3 d))。同时，我们将每个增强视图的m个切片进行批量堆叠，并将其作为二维输入，以二维模式计算逐片表示，然后将所有切片的平均输出视为整体切片表示，即f ' 1 =1 m m i=1 Pθ(Fθ([x1] i);(2D))，且f′2 =1 m f′i=1 P μ (f μ ([x2] i;2D))，其中[x] i表示从3D数据x中提取的第i个切片。之后，我们构建如下目标函数

上述目标函数鼓励我们从三个方面学习与三维医疗数据的精细化一致性，即体积与体积、切片与切片、体积与切片。

我们引入了一种替代训练方案来解决这一多目标优化问题。如图3所示，我们首先对2D图像进行采样，以从步骤0到步骤υ训练UniMiSS，然后在接下来的步骤υ中轮流对3D体积进行采样。下面的训练过程将以循环的方式继续，直到模型收敛。提出的迭代训练方案有两个优点:(1)克服了在同一批图像中同时使用二维和三维图像的困难;(2)减少了二维和三维数据分布差异带来的不稳定性。

4 Experiments

4.1 Datasets

训练数据集。我们从5个数据集(即MOTS数据集[52]、LIDC-IDRI数据集[5]、天池数据集[2]、RibFrac数据集[23]、TCIACT数据集[4])中收集了5022张3D CT扫描图，并从NIH ChestX-ray8数据集[46]中收集了108948张2D图像，以自监督的方式训练UniMiSS。

下游数据集。表1给出了6个下游任务的详细信息，可分为(1)3D下游:基于ct的分割(BCV)和分类(RICORD)，基于mri的分割(CHAOS);(2) 2D下游:多器官分割(JSRT)、肺炎分型(ChestXR)、皮肤病变分割(ISIC)。请注意，CHAOS和ISIC数据集在模式方面与预训练数据不同(即3D CT与MRI, 2D x射线与皮肤镜检查)。它们被用来评估无形模态的可转移性。

4.2 Experimental Details

训练的设置。我们将输入2D补丁的大小设置为224 × 224, 3D补丁的大小设置为16 × 96 × 96，旨在权衡为SSL保留足够的信息和将计算和空间复杂性降低到可承受的水平之间的平衡。我们应用了一组丰富的数据增强来创建正面视图，包括颜色抖动，高斯模糊/噪声，随机裁剪，缩放和翻转到产生两个视图的输入。随后[7]，我们采用了具有余弦衰减学习率[31]的AdamW优化器[32]，预热期为10次，来训练我们的UniMiSS。我们将初始学习率设置为0.0008，批量大小设置为192，最大epoch设置为200，速率参数ω设置为0.9，温度参数τt和τs分别设置为0.04和0.1。使用8个NVIDIA V100 gpu对UniMiSS进行预训练大约需要2.5天。我们知道这是一个很大的GPU消耗，但它节省了大量的时间和金钱来收集3D医学图像数据，因为我们使用容易收集的2D数据作为燃料。

下游培训设置。对于分类，我们提取了预训练的MiT编码器，并附加了一个FC层，输出通道作为预测的类数。对于分割，我们在删除SSL令牌的同时使用预训练的MiT编码器和解码器，并附加一个分割头用于预测。该头部包括一个转置卷积层、一个卷积- in - leakyrelu和一个卷积层，其核大小为1，输出通道为类数。分割性能由Dice系数得分来衡量。分类性能通过接收算子曲线下面积(AUC)来衡量。注意，我们随机分割25%的训练样本作为验证集，以选择消融研究中UniMiSS的超参数。每个下游任务的详细培训设置见附录

4.3 Results on 3D downstream tasks

特定维度的SSL与跨维度的SSL。在本节中，我们评估了SSL在两个下游3D测试中的性能，即多器官分割(BCV)和COVID-19筛查(RICORD)。将UniMiSS预训练与随机初始化(Rand)进行了比较。)和五种高级SSL方法，包括MoCo v2/v3[12,13]、PGL[47]、PCRL[55]和DINO[7]。

注意，MoCo v2、PGL和PCRL采用CNN作为编码器骨干，即一个具有50个可学习层的3D ResNet。在SSL过程中，MoCo v2和PGL仅对编码器部分进行预训练，而PCRL通过重构任务对解码器部分进行预训练。此外，MoCo v3, DINO和我们的UniMiSS使用Transformer模型作为主干，其中包含编码器和解码器。我们采用类似u的PVT作为MoCo v3和DINO的主干，它们具有与MiT相似的架构，但具有不同的补丁嵌入模块。由于缺乏SPE，它们无法同时处理2D和3D图像，从而导致仅使用3D数据的特定维度SSL。为了公平比较，所有这些SSL方法都是在5022个未标记的3D CT扫描上进行预训练的。略有不同的是，拟议的联苏特派团在3D SSL培训中增加了2D x射线，受益于普遍性。

我们对提议的UniMiSS和其他特定维度的CNN/Transformer SSL方法进行了更详细的比较。表2显示了三种标签设置(20%、40%和100%标签可用)的结果。我们将该表总结为以下几点:(1)基于transformer的模型明显优于基于cnn的方法，这主要归功于SSL预训练。它反映了Transformer是一个竞争性的体系结构，并且SSL预训练对于Transformer实现良好的性能至关重要。(2)所提出的UniMiSS优于MoCo v3和DINO。当100%的标签可用时，分割性能比DINO提高+2.38%，分类性能比DINO提高+3.15%。它证明了使用丰富的2D医学图像来辅助3D SSL过程的有效性。(3)此外，令人鼓舞的是，所提出的UniMiSS能够在更少的注释，甚至一半的情况下实现相当甚至更好的性能。以BCV为例，40%标签的UniMiSS分割概率达到80.97%，优于100%标签的随机初始化方法的79.97%。

BCV在线测试集对比为了更有说服力，我们还在BCV在线测试集上将所提出的UniMiSS与其他最先进的分割方法进行了比较。如表3所示，这些比较的方法包括PaNN[56]、UNETR[18]、nnUnet[22]和DoDnet[52]。请注意，这些竞争者的表现记录来自他们的原始论文。结果表明，我们的UniMiSS在不使用任何集成策略的情况下，仍然取得了竞争性能、最佳Hausdorff距离(HD)和平均表面距离(SD)，以及在线测试集上第二高的Dice，优于带有监督预训练的DoDNet。当使用像[22]这样的从粗到精的集成策略时，我们的UniMiSS可以在所有指标方面获得最佳性能。

2D下游任务的结果。由于可以对二维和三维医学图像进行预训练，因此我们的UniMiSS可以自由地应用于二维下游任务。

表4给出了二维医学图像分割和分类任务的比较。比较的方法包括Rand。init。， ImageNet预训练(INpre)，基于cnn的SSL方法(即MoCo v2, PGL和PCRL)，以及基于transformer的SSL方法(即MoCo v3和DINO)。与3D场景不同的是，MoCo v2、PGL和PCRL中使用2D ResNet-50作为骨干网。MoCo v3和DINO仍然以u型PVT为主干，但修改了补丁嵌入以适应二维输入。在这里，所有比较的SSL方法都是在相同的2D未标记医学图像上进行预训练的。对于UniMiSS，我们直接将之前预训练好的模型应用到2D任务中，不做任何修改，也不做进一步的再训练。结果表明:(1)SSL方法在两项任务上均优于INpre方法，表明在大规模医学图像数据集上的预训练比在自然图像上的预训练对医学领域下游任务更友好;(2)尽管3D数据的数量远少于2D数据，约为1 / 20，但UniMiSS预训练仍然可以获得优于纯2D SSL方法(如DINO)的性能增益。这可能部分解释了x射线和ct之间的内在相关性。联卢综合管理系统可以获得这种相关信息，从而有助于提高性能。