[논문 리뷰] UniBrain: Universal Brain MRI Diagnosis with Hierarchical Knowledge-enhanced Pre-training

Jiayu Lei, Lisong Dai외 10명|arXiv (Cornell University)|2023. 09. 13.

Radiomics and Machine Learning in Medical Imaging인용 수 5

한 줄 요약

UniBrain는 24,770개의 영상-보고서 쌍을 사용하여 보편적인 뇌 MRI 진료를 위한 계층적 지식 강화 사전학습 프레임워크를 제안한다. 이는 세분화된, 모odal별로 분해된 보고서와 결합된 시각-언어 인식을 활용하여 교차 모odal 정렬을 향상시키며, 특정 질환에서 인간 전문가 수준의 정확도를 달성하는 뛰어난 진단 성능을 보인다.

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) 모듈을 도입한다.
• 글로벌 보고서-이미지 정렬과 모달별 정렬을 통합한 다중 수준의 대비 손실을 구현하여 다양한 정밀도 수준에서 특징 표현을 강화한다.
• 다양한 질환과 MRI 시퀀스(T1WI, T2WI, FLAIR, DWI)를 포함하는 24,770개의 뇌 MRI 영상-보고서 쌍으로 구성된 대규모 데이터셋에서 모델을 훈련한다.
• 다중 수준에서 시각 및 언어 특징을 융합하기 위해 크로스 어텐션 메커니즘을 활용한 이중 브랜치 백본(시각 및 텍스트 인코더)을 사용한다.

실험 결과

연구 질문

• RQ1 계층적 지식 강화 사전학습 프레임워크는 뇌 MRI 진료에서 다양한 질환에 걸쳐 일반화 능력을 향상시킬 수 있는가?
• RQ2 글로벌 정렬에 비해 세분화된 모달별 보고서 정렬은 시각-언어 표현 학습의 효율성과 정확도를 향상시키는가?
• RQ3 CVP 모듈에서 질환 기술을 질의로 사용할 경우 진단 성능과 해석 가능성은 어느 정도 향상되는가?
• RQ4 기존 최신 기술 대비 오픈 어휘 및 도메인 이동 설정에서 UniBrain의 성능은 어떠한가?
• RQ5 특정 질환 유형에서 모델은 전문 방사선의사 수준의 진단 성능을 달성할 수 있는가?

주요 결과

• UniBrain는 BraTS2019 데이터셋에서 매크로 F1 스코어 63.27%를 기록하여 모든 최신 기술을 능가하며, 여러 지표에서 F1과 AUC 모두 1.5–2.5% 향상되었다.
• 내부 데이터셋에서 심한 클래스 불균형이 존재하는 상황에서도 매크로 F1 62.27%와 AUC 90.71%를 달성하여 강력한 내성과 일반화 능력을 입증했다.
• CVP 모듈을 제거하면 F1이 63.27%에서 51.73%로 10个百分点 이상 감소하여, 질환 특화 특징 학습에서 CVP의 핵심적 역할을 확인했다.
• CVP에서 질환 기술을 질의로 사용할 경우 질환 이름만 사용하는 것보다 성능이 유의미하게 향상되었으며, F1이 약 11.5个百分点 향상되었다.
• 절단 실험에서 글로벌 및 모달별 영상-보고서 정렬을 모두 제거하면 F1이 63.27%에서 60.61%로 감소하여 계층적 지식 강화의 중요성을 확인했다.
• 특정 질환 유형에서 UniBrain의 진단 성능는 전문 방사선의사와 유사했으며, 특히 겨드랑이 종양 및 뇌경색 분류에서 뛰어난 성능를 보였다.