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

Abstract

在计算机视觉领域中实现监督深度学习的成功往往依赖于大规模标注数据集。然而，在标注数据量有限的问题上存在显著挑战，在眼科医学影像分析领域尤为突出。由于手动注释过程耗时耗力，在此场景下问题更加突出。自监督学习（SSL）提供了一种有效途径来充分利用未标注的数据资源。为了最大化地利用未标注的眼科图像数据并克服尺寸限制问题，在本文中我们开发出一种统一架构的自监督Transformer模型框架——Uni4Eye系统。该框架旨在自动发现眼科医学影像固有属性并提取具有特定领域特性的嵌入表示。Uni4Eye可作为一个全局特征提取器，在遮蔽建模任务基础上构建于Visual Transformer架构之上，并采用统一化的Patch Embedding模块取代ViT中的标准Patch Embedding模块以实现对二维与三维输入信号的有效联合处理。此外我们创新性地设计了一个双分支多任务解码器模块该模块不仅能够同时作用于输入影像及其梯度图完成两个重建任务还能通过判别性表示促进模型收敛速度和效果提升。具体而言 Uni4Eye编码器经过微调优化后被应用于六个不同眼科医学影像分类下游任务均取得了令人满意的性能表现指标与现有的最先进自监督预训练模型进行比较实验结果表明 Uni4Eye系统展现出显著的优势

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

1 Introduction

我们开发出一种简洁且高效的架构，在分析2D与3D眼科医学图像时能够提取具有普适性的特征表示，并命名为Uni4Eye系统。该系统专为执行包含Vision Transformer架构的两项互信息最大化任务而进行设计。在Uni4Eye中开发出一种创新性的模块化架构，并实现对2D与3D图像进行特征提取与重建功能的统一处理。具体而言，在我们的网络结构中采用包含两个解码器组件的设计，并基于多种建模与重构任务进行联合优化以提升性能表现。此外，在 Uni4Eye 系统上构建了一个规模庞大且多模态的眼科医学图像数据集，并命名为mmOphth-v1版本（共包含95,978份样本），随后对该系统进行了预训练以获得更好的参数初始化效果。

总的来讲，我们的主要贡献有三个方面:(1)基于我们目前掌握的信息，该研究首次提出了一种基于自监督预训练的框架来学习2D和3D眼科图像的一般视觉表示。(2)通过系统地收集与整理了国内最大的二维和三维多模态眼科图像数据集，并命名为mmOphth-v1作为开放数据集发布。(3)针对涉及常见眼病的4个数据集的6个下游分类任务进行了系统性评估，并验证了该方法的有效性与优越性。我们的研究结果表明，在这些任务上所提出的Uni4Eye方法显著优于当前最先进的人工智能自监督预训练方法，并已将源代码公开 available at 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

以实现

UPE架构示意图见图2。采用自mmOphth-v1的数据集中的未标注训练图像样本表示为x d，在其中d=2f, 3g分别代表二维和三维空间中的尺寸参数。随后对输入样本x d进行预处理操作τ后得到用于训练UPE模块的数据样本x^ d。随后根据输入样本维度将补丁嵌入空间映射到特定维度，在P阶段则采用随机掩蔽策略生成掩蔽补丁嵌入f_d，并在d阶段跳过此策略处理后续数据流。具体而言，在本研究中由于我们基于规则划分的方法将图像划分为不重叠的小块（二维为正方形块体二维为立方体块体），因此我们在选取随机块时采用了均匀分布采样而不涉及任何替换操作其余未被选中的部分则被屏蔽掉不影响编码器后续处理过程这一特性与我们提出的新型网络架构存在显著区别

2.2 Dual-Decoder for Intensity and Edge Reconstruction

本节将介绍梯度图的生成过程及其在图像处理中的作用。特别地，请注意：本研究中所采用的核心算法设计基于深度学习框架，并结合多尺度特征提取策略以提升图像恢复质量；在此过程中涉及的关键步骤包括图像预处理阶段的质量控制机制设计以及深度学习模型参数的有效配置策略探讨；其中关键术语包括水平方向Gx-Sobel算子以及垂直方向上的Gy-Sobel算子等基本组件构建；

我们采用两个独立的设计单元——强度编码模块与边缘编码模块——来引导网络去学习能够整合局部与全局空间信息的特征表达。这两个模块拥有相同的设计架构，并共同接收编码后的可见区域g_{d的信息（如图1所示）。在这个方案中,g}d代表了一个完整斑块集,它包含了编码后的可见区域以及被遮挡的部分.无论是强度编码模块还是边缘编码模块,它们都共同接收同样的输入参数g~d.值得注意的是,这两种类型的模块在实现上存在差异,主要体现在它们对重建目标的不同追求上.

如图3所示，在本研究中选取眼底图像作为示例分析对象。左侧显示原始输入数据,该数据是强度解码器所追求的目标输出结果。随后,我们分别对原始图像沿水平(Gx)和垂直(Gy)方向施加Sobel算子[18],从而生成水平梯度分布和垂直梯度分布图形。接着,将这两个梯度分布图形进行积分运算,最终获得眼底图像的整体梯度分布特征图形,这被用作边缘解码器所期望的重建目标

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

相较于原始图像而言，在处理后的梯度图中能够更加准确且均匀地反映了视网膜结构的边缘特征，并显著提升了图像细节的表现能力。综合考虑各因素后，在P阶段阶段权衡参数选择λ₁和λ₂作为关键超参数，在此基础上我们所提出的自监督学习框架旨在优化以下目标函数：

其中Li和Le分别代表强度/边缘解码器在预测阶段与对应目标之间的掩蔽补丁所对应的均方误差(MSE)损失函数。参数λ₁与λ₂都被设定为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展示了在四个二维数据集上的微调结果.其中' Rand'代表使用随机初始化的模型参数.其中' viti-base'与' viti-large'分别代表' viti-base-patch16-224'与' viti-large-patch16- 00'.其中'- '标记的数据点无法在原始文献中找到结果（单位:%）.

3.2 Comparison with State-of-the-art

我们评估了Uni4Eye与其他基于卷积神经网络(CNN)或视觉 Transformer(ViT)作为主干且达到当前最佳水平(SOTA)的自监督学习(SSL)方法之间的性能差异。通过表1的数据展示可以看出，在四个二维数据集上的二值分类任务中，在AICChallenger-AMD 和 AICChallenger-PM 数据集中表现优异的结果已被记录下来。Li 等人[20]通过将配对的眼底 OCT 和 FFA 数据输入至 CNN 进行自监督对比学习，并在此基础上实现了AICChallenger-AMD 和 AICChallenger-PM 数据集中 SOTA 的性能水平。此外，在图像重建、旋转预测以及对比学习任务中的预训练阶段,SiT 模型超越了随机加权初始化及 ImageNet 预训练策略。尽管这些 SSL 方法有助于提升分类性能,但我们必须强调的是,不论所采用的是何种主干架构,Uni4Eye 均表现出色,并超越了所有其他比较方案的方法。特别是在 AICChallenger-AMD 数据集中,我们的模型在 F1 得分方面比第二名方案高出 2.2%左右

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

基于3D应用场景的研究中，在GAMMA数据集以及OCTA-500数据集上开展的OCT卷经过微调训练后得到了Uni4Eye模型。从表2可以看出，在随机初始化方法的基础上以及ImageNet预训练方法的基础上（后者指仅更换ViT拼块嵌入层为3D版本而保持ViT其他预训练参数不变），所提出的Uni4Eye方法表现更为突出。值得注意的是，在ImageNet预训练这一前提下（即仅更换ViT拼块嵌入层为3D版本而保持ViT其他预训练参数不变），因为用于获取3D眼科图像的数据样本数量相对较少而导致的结果分类性能相较于传统二维模型有所下降

表 3：在多个不同的 mask ratio α 的 mmOphthv1 上训练自监督模型；然后基于 icchallenge-amd 数据集对其进行微调训练，并获得最终结果（单位：%）。

3.3 Reconstruction Results

基于mmOphth-v1预训练相同的网络，在不同眼科模式下的重建结果中进行可视化展示

4 Conclusion

本研究开发了一款简洁高效的眼科图像分析自监督学习框架Uni4Eye。具体而言，在ViT架构中对嵌入块进行优化后构建UPE组件后 Uni4Eye能够有效地解决维度问题。此外我们还开发了一种基于MIM任务的双分支解码器模块 该框架不仅能够提取眼底图像中的明暗特征还能有效捕捉边缘细节。

基于4个二维数据集以及2个三维数据集上的全面实验结果表明，在眼科疾病识别任务中，Uni4Eye模型展现了超越具有代表性的SOTA方法的分类效果。实验结果证实了MIM算法在医学图像自监督预训练任务中的广泛适用性。未来的研究工作计划将聚焦于评估该框架在其他类型的医学图像处理中的适用性，并寻求进一步优化其性能的可能性。