浙江大学联合微软亚洲研究院提出了一种全新的统一医学图像预训练框架 UniMedI。它利用诊断报告作为公共语义空间,可为不同模态的医学图像创建统一的表示,成功整合了 2D 和 3D 图像,使复杂的医学数据被更好地利用。
让 AI 在某些条件下具备类似人类的反应能力,从而代替人类高效地从事特定工作,是 AI 领域研究人员孜孜不倦的追求。正如在医学图像和人工智能的交叉领域,基于视觉语言预训练的深度模型 (Visual-Language Pre-training, VLP) 凭借其自动化的特点,可以在大量图像及对应文本的数据集上进行预训练,并学会自动从新的图像中提取相关特征,可高效地解决费时费力的人工注释需求。
然而,尽管 VLP 在医疗领域已经取得了一定程度的成功,但在进一步扩大其应用的数据规模时,仍然面临着诸多挑战。
首先,现有模型训练大多主要以单模态数据(主要是 2D 图像,如 X 光片)为主,这与包含多模态图像(包含 2D 和 3D 图像,如 CT、MRI 图像等)的真实医学场景并不相符;其次,不同模态的医学图像固有的异质性也阻碍了它们的有效协作和整合。此外,医学图像不同模态的数据还存在维度差异、缺乏成对数据等。因此,如何构建一个统一的模型,并将这些不同模态的数据有效映射至共同空间,实现联合学习,成为了一个极具挑战性的课题。
为了解决上述问题,浙江大学胡浩基团队联合微软亚洲研究院邱锂力团队提出了一种全新的统一医学图像预训练框架 UniMedI。它利用诊断报告作为公共语义空间,可为不同模态的医学图像创建统一的表示,此外,它还引入创建「伪配对」 (Pseudo-Pairs) 的技术,在文本的指导下,UniMedI 能够从复杂的 3D 图像中选择与该文本相关的 2D 切片,这些切片充当桥接 2D 和 3D 数据的伪对,可增强各种医学成像模式之间的一致性,有效整合医学多模态图像。
相关研究成果以「Unified Medical Image Pre-training in Language-Guided Common Semantic Space」为题,收录于计算机视觉和机器学习领域的顶会 ECCV 2024。
研究亮点:
* 在实验中,UniMedI 在多个不同数据集上的 2D 和 3D 图像上的性能表现优异,在广泛的医学任务,如图像分类、分割和检索方面表现卓越
* UniMedI 可以统一采集 2D 和 3D 图像,解决了医学领域的数据稀缺问题
论文地址:
https://eccv.ecva.net/virtual/2024/poster/1165
真实医学数据,有效验证框架
预训练 UniMedI 框架所使用的数据来自 JPG 版本的 2D X 光片数据集 MIMIC-CXR 2.0.0 和 3D CT 扫描数据集 BIMCV。
其中,研究人员对 2D 数据集进行了预处理,消除了所有侧面图像,以便与仅使用正面图像的下游任务对齐。同时,为保持数据集的完整性,实验中未使用少于 3 句话的 2D 和 3D 数据集简短报告。
图像方面,2D 图像的大小为 224 × 224,3D 图像大小为 128 × 128 × 32。
研究团队在 8 个 Tesla V100 GPU 上对 UniMedI 框架进行了 50 次预训练,batch 大小为 144。
在实验评估中,首先团队对 2D 和 3D 数据集进行了医学图像分类,其中具有代表性的 2D 数据集有 3 个,分别是:CheXpert,包含 191,229 张正视图胸片;包含约 29,700 张正视图胸片的 RSNA 肺炎 stage 2 版本;来自 2,800 多名患者的 16,490 张阳性 COVID-19 图像。
然后团队对两个代表性的 3D 数据集进行了分类,分别是 CC-CCII 和 LUNA 16。其中,CC-CCII 使用了 Clean-CC-CCII 版本,该版本包含来自 2,698 例患者的 3,993 次扫描的 340,190 个切片;基于 LIDC-IDRI 建立的 LUNA 16,包含 888 个 CT 扫描,带有注释。实验删除了 LIDC-IDRI 数据库中切片厚度大于 3mm 的 CT 扫描。
层层协作机制,打破数据藩篱
本研究提出的 UniMedI 是一个视觉语言预训练框架,医学图像及其文本报告分别由视觉编码器 (Vision Encoder) 和文本编码器 (Text Encoder) 两个编码器进行编码,然后通过 VL (Vision-Language) 对比学习共同学习。UniMedI 的独特之处在于它可以用统一的方式有效地采集 2D 和 3D 图像,解决医疗领域的数据稀缺问题。UniMedI 总体框架如下图左边所示:
UniMedI 总体框架:左边为整体流程,右边为关键设计
实验中,视觉编码器采用的是 ViT-B/16,主要在 2D 和 3D 视觉数据的公共特征空间中提取表示,文本编码器使用的是 BioClinicalBERT,用来对文本特征进行编码,视觉编码器和文本编码器在 2D 和 3D 数据中是通用的。
为了克服不存在配对 2D 和 3D 图像数据的挑战。研究团队在 UniMedI 中引入创建「伪配对」的方法,该方法的设计基于一种新颖的以语言为指导的注意力切片选择策略。
比如,当输入为一个 3D 图像时,从其中提取与报告最相关的一部分 2D 切片,然后将选定的切片视为 2D 图像,从而形成 2D-3D 图像的伪配对关系。在此之后,通过将选定的 2D 切片与原始 3D 图像一起输入到网络中,可共同学习它们与报告 (Report) 之间的关系,最终形成一个统一的特征空间。当输入为 2D 图像时,则省略了切片选择处理。
之后,一个视觉编码器将所有多模态图像(包括原始 2D 和 3D 图像以及所选择的 2D 切片)映射到表示空间中。该视觉编码器分别拥有用于 2D 和 3D 图像的标记器 T2D 和 T3D,以及用于实现更好集成的共享骨干 Ev。通过对比学习损失 Lᵥₗ,在一个 VLP 中端到端学习包含视觉编码器和文本编码器 Eₗ 的模型。在这个过程中,2D 和 3D 图像都可以被编码到由报告中语言信息监督的公共语义空间中。
为了充分利用医学图像本身多模态数据,以及共享的一些公共信息,本研究中还引入了一个辅助任务设计,即掩蔽和恢复,并使用自蒸馏方法来完成任务。这使得 2D 和 3D 图像的 token 可以互相通信,并增强跨维度的交互和多模态图像的集成。
值得注意的是,UniMedI 的一大亮点是注意力切片选择策略和 VL 对比学习的协同效应。
* 一方面,VL 对比学习可以实现语言监督,这种监督直接应用于视觉 CLS token。该 token 在报告中包含重要信息,因此视觉 CLS token 的注意权重作为 2D 切片选择的基础,才会携带来自报告的监督信息,并与 3D 特征一起构建联合特征空间。
* 另一方面,细致的切片选择使 2D 和 3D 特征空间更加整合,即便是没有配对数据。这种共同空间可以放大医学图像和报告之间的详细信息,并且通过这种方式,推动了图像和报告之间的对齐。这两种设计使得多模态图像的表征结合在一起,并使其同时接近报告表征空间,在构建共同语义空间上达到了一加一大于二的效果。
多角度实验评估,性能赶超 UniMiss
为了对 UniMedI 进行全面有效的评估,本研究设置了多角度观察,并通过与各种医学 VLP 方法进行比较分析,进而验证其性能和有效性。
首先,研究团队将 UniMedI 与包括 ConVIRT、GLoRIA、MGCA、LOVT、PRIOR 等在内的方法进行了比较,后者是针对 X 光片及其相应的医学报告量身定制的;然后,研究团队又将 UniMedI 与几种 2D 和 3D 联合学习的方法进行比较,包括 UniMiss 和 Joint。
线性分类实验结果显示,在 2D 医学图像分类实验结果中(如下图示),相比使用 ViT 作为视觉编码器的最先进 MGCA (ViT-b/16) 方法,UniMedI 在不同训练数据下 (1%,10%,100%) 的 3 种 2D 医学图像分类中表现最好。
* 线性分类实验:用来评估 UniMedI 的表征能力
与其相比,UniMedI 在 CheXpert 数据集上的 AUROC 分别提高 +0.6%、+0.6% 和 +0.8%;在 RSNA 数据集上的 AUROC 分别提高 +0.9%、+0.5% 和 +0.7%;在 COVID 数据集上的 AUROC 分别提高 +5.5%、+7.6% 和 +2.3%。实验结果表明了所提算法的有效性。
1%、10%、100% 训练数据下在 CheXpert、RSNA 和 COVID 数据集上的 2D 线性分类结果
在 3D 医学图像分类实验结果中(如下图示),与最先进的 UniMiss 对比,UniMedI 在 CC-CCII 数据集上分别提升了 +22.6%、+2.0% 和 +0.8% 的 ACC 增益。这些数据均验证了 UniMedI 的数据效率和有效性。
1%、10%、100% 训练数据下在 CC-CCII 上的 3D 线性分类结果
与此同时,当使用完整的训练数据对全视觉编码器进行微调时,UniMedI 在 CC-CCII 和 LUNA 多个 3D 医学图像数据集上的性能均优于其他方法。
如下图所示,UniMedI 在 CC-CCII 数据集上的 ACC 值为 93.8%,在 LUNA2016-v2 数据集上的 ACC 值为 95.9%。这显示其在 2D 和 3D 医学图像分类任务上的显著泛化能力,表明该框架具有提取三维 CT 图像通用特征的能力。
完整训练数据的 CC-CCII 和 RICORD 数据集上的 3D 微调结果
医学语义分割实验结果显示,在 2D 医学语义分割结果中,UniMedI 明显优于当前最先进的 MGCA 算法,当使用 1% 的训练数据时,UniMedI 达到了 67.8% 的 Dice。在 3D 医学语义分割结果中,UniMedI 在 BCV 数据集上与 UniMiss 相比,当有限标签可用性为 40% 和 100% 时,精度分别比 UniMiss 提高了 0.6% 和 0.4%,如下图所示。
* 医学语义分割实验:用来评估分割性能,使用 RSNA 肺炎正视图胸片,和 BCV 数据集(包括 50 张 CT 扫描)。
这些结果验证了 UniMedI 在提取有意义的特征和有效利用有限注释数据方面具有强大的优越性,证明了其在利用局部表示进行语义分割任务时具备更高的熟练程度。
科技助力,加深 VLP 与医学图像的羁绊
视觉语言预训练模型正在成为连接计算机视觉和自然语言处理的重要桥梁,尤其是在医学图像领域,通过大规模的视觉和语言数据进行预训练,它们能够轻松捕获复杂医学图像与文本之间的复杂关系,进而辅助医生进行图像诊断,帮助企业进行药物研发,亦或者实现智能的医学图像管理。
本次研究得以入选国际顶会,也从另一方面再次佐证了在人工智能与医学图像的交叉领域,VLP 所蕴藏着的巨大想象空间。实际上,除了浙江大学与微软亚洲研究院两大团队本次强强联手外,已经有不少实验室针对这一领域进行了攻坚。
比如上述研究中所提到的先进方法之一的 UniMiss,其相关成果早在 2022 年就由澳大利亚阿德莱德大学和西北工业大学计算机学院的团队,以题为「UniMiss : Universal Medical Self-Supervised Learning via Breaking Dimensionality Barrier」刊登在当年的 ECCV。
论文地址:
https://dl.acm.org/doi/abs/10.1007/978-3-031-19803-8_33
在这篇研究中,作者主张利用大量的 2D 图像来弥补 3D 数据的不足,旨在建立一个通用的医学自我监督表达学习框架,并命名为 UniMiss。实验结果显示,相比 ImageNet 预训练和其他高级 SSL (self-Supervised learning) 对手,UniMiss 展示出了极大的优势 ,在 2D/3D 医学图像分析任务中,无论是分割还是分类,其结果均能够人满意。
不仅如此,在今年的 7 月份,该团队又针对 UniMiss 进行了新一轮研究,并提出了 UniMiss+。目前,相关成果以题为「UniMiSS+: Universal Medical Self-Supervised Learning From Cross-Dimensional Unpaired Data」收录于知名国际期刊 IEEE Transactions on Pattern Analysis and Machine Intelligence 上。
论文地址:
https://ieeexplore.ieee.org/document/10617802
在最新的研究中,该团队在 UniMiss+ 中引入了数字重建 X 光片技术,用来模拟 CT 扫描的 X 光片图像,以便访问配对的 CT 和 X 光图像数据。相比上一代 UniMiss 得到了巨大改进。
总而言之,融合人工智能与医学图像的相关科研仍旧在火热进行,假以时日,这些成果也必将转化为应用,落地到真实的医学场景中去,成为造福医务工作人员、患者、企业的新工具。
往期回顾
Z动态|一家融了15轮的成都公司,要被浙大系半导体并购了
Z科技|有望用于清洁能源气体储存,浙大团队提出高孔性金属有机框架合成新策略
Z科技|入选NeurIPS 24,浙大团队提出全新去噪蛋白质语言模型DePLM,突变效应预测优于SOTA模型
—— Angel For ZJU