[论文解读] Self-supervised Image-text Pre-training With Mixed Data In Chest X-rays

Xiaosong Wang, Ziyue Xu+3 more|arXiv (Cornell University)|Mar 30, 2021

COVID-19 diagnosis using AI参考文献 44被引用 14

一句话总结

该论文提出了一种自监督的、基于Transformer的图像-文本预训练框架UWOX，利用混合数据输入（包括成对、非成对以及跨机构的图像-文本数据）在胸部X光片中进行联合表征学习。通过引入多尺度掩码视觉建模和交叉相关模块，UWOX在MIMIC-CXR、NIH14-CXR和OpenI-CXR数据集上的分类、检索和图像重建任务中均实现了最先进性能，即使在异质数据源下也展现出强大的鲁棒性和泛化能力。

ABSTRACT

Pre-trained models, e.g., from ImageNet, have proven to be effective in boosting the performance of many downstream applications. It is too demanding to acquire large-scale annotations to build such models for medical imaging. Meanwhile, there are numerous clinical data (in the form of images and text reports) stored in the hospital information systems. The paired image-text data from the same patient study could be utilized for the pre-training task in a weakly supervised manner. However, the integrity, accessibility, and amount of such raw data vary across different institutes, e.g., paired vs. unpaired (image-only or text-only). In this work, we introduce an image-text pre-training framework that can learn from these raw data with mixed data inputs, i.e., paired image-text data, a mixture of paired and unpaired data. The unpaired data can be sourced from one or multiple institutes (e.g., images from one institute coupled with texts from another). Specifically, we propose a transformer-based training framework for jointly learning the representation of both the image and text data. In addition to the existing masked language modeling, multi-scale masked vision modeling is introduced as a self-supervised training task for image patch regeneration. We not only demonstrate the feasibility of pre-training across mixed data inputs but also illustrate the benefits of adopting such pre-trained models in 3 chest X-ray applications, i.e., classification, retrieval, and image regeneration. Superior results are reported in comparison to prior art using MIMIC-CXR, NIH14-CXR, and OpenI-CXR datasets.

研究动机与目标

• 为应对大规模标注医学影像数据有限的挑战，利用医院数据库中未被充分利用的原始临床图像-文本对。
• 开发一种统一的预训练框架，能够处理混合数据输入，包括成对、非成对以及跨机构的图像-文本数据。
• 通过自监督预训练提升下游胸部X光应用（如分类、检索和图像重建）的性能。
• 证明基于Transformer的架构结合多尺度视觉建模，在数据充足时可超越或匹配基于CNN的模型在医学影像中的表现。

提出的方法

• 提出一种统一的基于Transformer的架构，通过自监督预训练实现图像与文本表征的联合学习。
• 引入多尺度掩码视觉建模作为自监督任务，用于重建图像块，提升空间特征学习能力。
• 采用交叉相关模块，即使在数据非成对或来自不同机构的情况下，也能建模图像与文本之间的关系。
• 使用成对匹配损失，增强在成对与非成对设置下图像与文本表征之间的对齐。
• 对文本采用掩码语言建模，并与掩码图像块重建相结合，实现联合预训练。
• 支持在混合数据输入上进行训练，包括跨机构数据（例如，来自一个来源的图像与来自另一个来源的报告），以最大化数据利用率。

实验结果

研究问题

• RQ1 自监督图像-文本预训练框架能否有效利用来自多个机构的混合数据输入，包括成对与非成对的图像-文本数据？
• RQ2 与单尺度方法相比，多尺度掩码视觉建模在医学影像中如何提升图像表征学习能力？
• RQ3 当数据非成对或来自不同来源时，所提出的交叉相关模块在多大程度上增强了图像与文本特征之间的对齐？
• RQ4 所提出的UWOX框架是否在下游胸部X光任务（如分类、检索和图像重建）中优于现有方法？
• RQ5 在数据充足的情况下，纯Transformer架构能否在医学图像分析中实现与CNN基线相当或更优的性能？

主要发现

• UWOX在MIMIC-CXR数据集上达到最先进性能，在图像和文本检索任务中均表现优异，图像查询和文本查询的P@10分别达到0.621和0.603。
• 采用块大小B=8的多尺度图像编码器与解码器在图像重建质量上表现最佳，平均SSIM为0.872，PSNR为356.5。
• 成对匹配损失显著提升了仅基于图像的检索性能，AUC从0.727提升至0.732。
• UWOX-mixup（利用NIH14-CXR的额外图像）在检索精度上达到最高，文本查询的P@1为0.763，图像+文本查询的P@10为0.657。
• 该模型展现出强大的泛化能力，在NIH14-CXR数据集的混合场景中取得最佳整体结果（P100%+uP），表明其对数据异质性的鲁棒性。
• 当训练数据充足时，基于Transformer的图像建模方法在分类与重建任务中可匹配或超越基于CNN的基线模型。