UniBrain: Universal Brain MRI Diagnosis with Hierarchical Knowledge-enhanced Pre-training

TL;DR

UniBrain is a hierarchical knowledge-enhanced pre-training framework that improves universal brain MRI diagnosis by leveraging large-scale data and report-based hierarchical alignment, outperforming state-of-the-art methods and rivaling expert radiologists.

Contribution

Proposes a novel hierarchical pre-training framework for universal brain MRI diagnosis utilizing report-based alignment and large-scale data, enhancing generalization and diagnostic accuracy.

Findings

Outperforms all state-of-the-art diagnostic methods.

Provides superior grounding performance.

Achieves comparable results to expert radiologists.

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.