【多模态（影像）自监督学习】Uni4Eye: Unified 2D and 3D Self-supervisedPre-training via Masked Image ModelingTran

Abstract

大规模标记数据集是计算机视觉中监督深度学习成功的关键因素。然而，标注的数据数量有限是非常常见的，特别是在眼科图像分析中，因为手动标注是费时费力的。自监督学习(SSL)方法为更好地利用未标记数据带来了巨大的机会，因为它们不需要大量的注释。为了尽可能多地使用未标记的眼科图像，有必要打破尺寸障碍，同时使用2D和3D图像。在本文中，我们提出了一个通用的自监督Transformer框架，名为Uni4Eye，用于发现眼科图像的固有属性并捕获嵌入的特定领域特征。Uni4Eye可以作为一个全局特征提取器，它建立在一个具有视觉转换(ViT)架构的蒙面图像建模任务的基础上。我们采用统一的Patch Embedding模块来代替ViT中的原始Patch Embedding模块，对二维和三维输入图像进行联合处理。此外，我们设计了一个双分支多任务解码器模块，在输入图像及其梯度图上同时执行两个重建任务，提供判别表示以获得更好的收敛性。我们通过在六个下游眼科图像分类任务上对预训练的Uni4Eye编码器进行微调来评估其性能。通过与其他最先进的SSL预训练方法的比较，Uni4Eye的优势得以成功确立。

关键词:自监督预训练·统一的2D和3D·视觉转换器·眼科疾病分类·多任务

1 Introduction

我们提出了一个简单而有效的框架，可以从2D和3D眼科图像中学习通用表示，称为Uni4Eye。Uni4Eye设计用于执行具有ViT架构的双MIM任务。我们在Uni4Eye中设计了一个双分支可切换的补丁嵌入层来代替原始的补丁嵌入层，使其可以切换到不同的分支来进行2D和3D图像的补丁嵌入。此外，我们在我们的网络中使用双分支解码器，并使用不同的建模/重构任务对其进行训练，从而实现更鲁棒的收敛和更好的表示。此外，我们创建了迄今为止最大的多模态和多维眼科图像数据集，共包含95,978个样本。我们将其命名为mmOphth-v1，我们提出的Uni4Eye在其上进行预训练。

总的来说，我们的主要贡献有三个方面:(1)据我们所知，这是第一次提出一个自监督预训练框架来学习2D和3D眼科图像的一般视觉表示。(2)收集并创建了国内最大的二维和三维多模态眼科图像数据集，命名为mmOphth-v1。该数据集将向公众开放。(3)我们对涉及常见眼病的4个数据集的6个下游分类任务进行了广泛的实验。我们提出的Uni4Eye优于其他最先进的(SOTA)自监督预训练方法，成功地建立在这些任务上。源代码可从https://github.com/Davidczy/Uni4Eye获得。

2 Methodology

我们的Uni4Eye概述如图1所示。有三个主要组成部分，包括统一补丁嵌入(UPE)模块，ViT编码器和双支路解码器。在预训练阶段，我们首先在两个MIM自监督任务上预训练我们的编码器，然后在不同的下游任务上微调我们的模型。如图1所示，预训练阶段和下游微调阶段分别记为P和d。P阶段的目的是训练编码器从不同输入的眼科图像中生成更广义和判别的表征。然后，利用UPE模块和D中的ViT编码器加载预训练的参数，并在不同的下游任务上继续微调，以获得更好的性能。对于下游分类任务，我们采用全连接层作为分类头，对编码器生成的特征和输出预测进行处理。

我们现在深入研究UPE和双分支解码器的细节。

2.1 Unified Patch Embedding Module

为了使我们的自监督框架兼容2D和3D数据，并适应不同的下游场景，我们使用UPE作为补丁嵌入模块。如图1所示，mmOphth-v1数据集中的不同图像，无论其维度如何，都可以直接馈入UPE模块。

UPE的结构如图2所示。设从mmOphth-v1中采样的未标记训练图像用x d表示，其中d 2f2;3g表示图像的尺寸。然后，将数据增强τ应用于x d以生成UPE模块的输入x^ d。UPE根据x^ d的维度将x^ d转换为特定的补丁嵌入。随后，在P阶段采用随机掩蔽策略生成掩蔽补丁嵌入f d，在d阶段跳过该策略。更具体地说，由于我们将图像划分为规则的不重叠的补丁(2D输入为2D方形补丁，3D输入为3D立方补丁)，因此我们遵循均匀分布对随机补丁进行采样而不进行替换。然后，剩下的被屏蔽掉，这意味着这些补丁不会被馈送到编码器中。因此，ViT编码器只对可见的补丁进行操作，而不对被掩盖的补丁进行操作，这与我们提出的方法不同。

2.2 Dual-Decoder for Intensity and Edge Reconstruction

图3:梯度图的生成。请注意，原始图像是强度解码器的重建目标，梯度图是边缘解码器的重建目标。关键词:水平方向Gx - Sobel算子;垂直方向上的Gy - Sobel算子。

我们使用两个解码器，即强度解码器和边缘解码器，以鼓励网络学习包含局部和全局信息的表示。强度解码器和边缘解码器具有相同的网络结构和相同的输入g~ d。如图1所示，g~ d表示由编码后的可见斑块g~ d和被遮挡的斑块组成的完整斑块集。G ~ d同时输入强度解码器和边缘解码器。两种解码器的区别在于重构目标的不同。

如图3所示，以眼底图像为例，左侧为原始输入，为强度解码器的重建目标。我们对原始图像在水平Gx和垂直Gy方向分别应用Sobel算子[18]，得到水平梯度图和垂直梯度图。然后，我们对这两个梯度图进行积分，得到眼底图像的梯度图，这是边缘解码器的重建目标。

我们将此操作应用于所有2D图像和每个3D体的2D切片。

与原始图像相比，梯度图均匀地表征了视网膜结构的边缘，更清晰地描绘了微小的视网膜血管。综上所述，在P阶段，权衡参数为λ1和λ2，我们的自监督学习框架的总目标函数为

其中Li和Le是强度/边缘解码器的预测与相应目标之间的掩蔽补丁的均方误差(MSE)损失。λ1和λ2分别设为0.5和0.5，使网络均匀地集中于感兴趣的眼科图像的全局强度信息和局部边缘信息。

3 Experiments and Results

3.1 Experimental Setup

在预训练阶段，mmOphth-v1的输入图像对2D图像下采样为224 × 224，对3D图像下采样为112 × 224 × 112。批量大小为64个2D和4个3D。数据增强策略是随机颜色抖动、随机灰度化、随机裁剪和随机水平翻转的组合。模型由AdamW优化器[23]进行优化，初始学习率为0.0005。我们的模型在PyTorch[25]中使用2个NVIDIA GeForce RTX 3090 gpu实现，需要50个epoch和20个小时才能收敛。在微调阶段，输入与上述设置保持一致。我们也使用AdamW作为优化器，初始学习率为0.0001，对于2D和3D图像，批大小分别设置为8和1。由于所有下游任务都是分类任务，因此我们采用曲线下面积(AUC)，准确度，精密度，召回率，f1分数和Kappa作为我们的评估指标。mmOphth-v1眼科数据集和评估数据集的详细信息见附录图A1和表A1。

表1:在四个2D数据集上进行微调得到的结果。Rand为随机初始化模型参数。viti -base和viti -large分别表示viti -base-patch16-224和viti -large-patch16-224。-表示结果无法从原始文章中获得。(单位:%)

3.2 Comparison with State-of-the-art

我们将Uni4Eye与其他采用卷积神经网络(CNN)或ViT作为主干的SOTA SSL方法进行了比较。不同预训练方法对四个二维数据集的二值分类结果如表1所示。Li等[20]将配对的眼底和FFA数据输入到CNN中进行自监督对比学习，并在icchallenge - amd和icchallenge - pm数据集上实现了SOTA性能。自监督视觉变压器(Self-supervised Vision Transformer, SiT)[1]进行图像重建、旋转预测和对比学习任务进行预训练，优于随机加权初始化和ImageNet预训练。虽然这些SSL方法有利于提高分类性能，但值得强调的是，无论骨干是什么，我们的Uni4Eye都优于所有比较的方法。在icchallenge - amd数据集上，我们的方法在f1得分方面比第二好的方法高出2.2%。

表2:通过GAMMA和OCTA-500数据集对3D OCT体积进行微调获得的结果。(单位:%)

对于3D下游任务，我们对来自GAMMA数据集和OCTA-500数据集的OCT卷进行微调Uni4Eye。如表2所示，我们提出的Uni4Eye性能优于随机初始化和ImageNet预训练。请注意，ImageNet预训练意味着我们只将ViT的补丁嵌入模块替换为3D版本，并保持ViT的其他预训练参数。由于3D眼科数据相对较少，3D模型的分类性能比2D模型差。

表3:首先在不同掩码比α的mmOphthv1上训练自监督模型，然后在icchallenge - amd数据集上进行微调得到的结果。(单位:%)

3.3 Reconstruction Results

我们将在mmOphth-v1上预训练的相同网络的不同眼科模式的重建结果可视化，以突出我们学习到的特征的普遍性。如图4所示，我们将不同模态的输入输入到网络中，得到重构结果。我们将UPE中的掩码比例设置为25%，50%，75%。很明显，较小的掩模比可以使模型产生更好的重建结果。然而，更好的重构并不等同于更好的下游任务性能。我们在icchallenge - amd上用相同的设置微调了这三个模型。如表3所示，50%掩码比预训练的网络在特定下游任务上的性能最好。对于眼科图像分析，这一结果可能表明编码器无法通过太简单(掩码比= 25%)或太难(掩码比= 75%)的重建任务生成判别表示。消融分析结果如附录表A2-A3所示，说明了解决尺寸障碍和采用双支路解码器的重要性。

4 Conclusion

本文提出了一种简单、统一、功能强大的眼科图像分析自监督框架Uni4Eye。具体来说，通过修改贴片嵌入模块，在ViT中生成UPE, Uni4Eye可以轻松打破维度障碍，同时处理二维和三维图像。我们还设计了一种基于MIM任务的双解码器结构，使Uni4Eye既能利用眼科图像中的强度信息，又能利用边缘信息。

在4个2D数据集和2个3D数据集上进行的大量实验表明，我们的Uni4Eye在眼病诊断方面取得了比代表性SOTA方法更好的分类性能。我们的结果也证明了MIM在各种医学图像分析中的自我监督预训练潜力。我们未来的工作将包括研究我们的框架对其他类型医学图像的可行性，并探索进一步提高我们框架效率的方法。