[论文解读] UniBrain: Universal Brain MRI Diagnosis with Hierarchical Knowledge-enhanced Pre-training

Jiayu Lei, Lisong Dai+10 more|arXiv (Cornell University)|Sep 13, 2023

Radiomics and Machine Learning in Medical Imaging被引用 5

一句话总结

UniBrain 提出了一种分层知识增强的预训练框架，用于通用脑部 MRI 诊断，基于 24,770 对影像-报告数据。该方法利用细粒度、模态特定的报告分解与耦合视觉-语言感知机制，提升跨模态对齐效果，实现更优的诊断性能，并在某些疾病上达到人类放射科专家水平的准确率。

ABSTRACT

Magnetic resonance imaging~(MRI) have played a crucial role in brain disease diagnosis, with which a range of computer-aided artificial intelligence methods have been proposed. However, the early explorations usually focus on the limited types of brain diseases in one study and train the model on the data in a small scale, yielding the bottleneck of generalization. Towards a more effective and scalable paradigm, we propose a hierarchical knowledge-enhanced pre-training framework for the universal brain MRI diagnosis, termed as UniBrain. Specifically, UniBrain leverages a large-scale dataset of 24,770 imaging-report pairs from routine diagnostics. Different from previous pre-training techniques for the unitary vision or textual feature, or with the brute-force alignment between vision and language information, we leverage the unique characteristic of report information in different granularity to build a hierarchical alignment mechanism, which strengthens the efficiency in feature learning. Our UniBrain is validated on three real world datasets with severe class imbalance and the public BraTS2019 dataset. It not only consistently outperforms all state-of-the-art diagnostic methods by a large margin and provides a superior grounding performance but also shows comparable performance compared to expert radiologists on certain disease types.

研究动机与目标

• 解决现有 MRI 诊断模型局限于狭窄疾病类别和小规模数据导致的泛化能力有限的问题。
• 克服多模态脑部 MRI 预训练中数据稀缺、领域特定知识整合与可解释性方面的挑战。
• 开发一种可扩展的通用框架，能够对多种脑部疾病实现强大的零样本和少样本泛化能力。
• 通过将图像特征与细粒度报告组件对齐，提升模型的可解释性。
• 通过查询驱动的感知模块实现灵活诊断，支持可变长度的疾病描述。

提出的方法

• 提出一种分层视觉-语言对齐机制，将放射科报告分解为模态特定与全局组件，实现细粒度的跨模态学习。
• 设计一个自动报告分解模块，从自由文本报告中提取与疾病相关的内容片段，实现与对应 MRI 模态的针对性对齐。
• 引入耦合视觉-语言感知（CVP）模块，利用疾病描述作为查询，关注全局图像特征，并生成疾病特定的分类与定位预测。
• 实施多级对比损失，结合全局报告-图像对齐与模态特定对齐，以在不同粒度上增强特征表示。
• 在包含 24,770 对脑部 MRI 影像-报告的大型数据集上进行模型训练，覆盖多种疾病和 MRI 序列（T1WI、T2WI、FLAIR、DWI）。
• 采用双分支主干网络（视觉与文本编码器），通过交叉注意力机制在多个层次融合视觉与语言特征。

实验结果

研究问题

• RQ1 分层知识增强的预训练框架是否能提升 MRI 诊断中对多种脑部疾病的泛化能力？
• RQ2 与全局对齐相比，细粒度的模态特定报告对齐是否能更高效、更准确地提升视觉-语言表征学习？
• RQ3 在 CVP 模块中使用疾病描述作为查询，在多大程度上能提升诊断性能与可解释性？
• RQ4 与当前最先进方法相比，UniBrain 在开放词汇和领域分布偏移设置下的表现如何？
• RQ5 该模型是否能在特定疾病类型上实现与专家放射科医生相当的诊断性能？

主要发现

• UniBrain 在 BraTS2019 数据集上取得 63.27% 的宏平均 F1 分数，优于所有 SOTA 方法，在多项指标上 F1 和 AUC 提升 1.5–2.5%。
• 在内部数据集上，模型取得 62.27% 的宏平均 F1 和 90.71% 的 AUC，即使在严重类别不平衡的情况下也表现出强大的鲁棒性与泛化能力。
• 移除 CVP 模块后，F1 分数下降超过 10 个百分点（从 63.27% 降至 51.73%），证实其在疾病特异性特征学习中的关键作用。
• 在 CVP 模块中使用疾病描述作为查询，相比仅使用疾病名称，性能显著提升，F1 分数提高约 11.5 个百分点。
• 消融实验表明，若同时移除全局与模态特定的影像-报告对齐机制，F1 分数从 63.27% 降至 60.61%，证实分层知识增强的重要性。
• 在某些疾病类型上，UniBrain 的诊断性能与专家放射科医生相当，尤其在胶质母细胞瘤和脑梗死分类方面表现突出。