BrainAnytime: Anatomy-Aware Cross-Modal Pretraining for Brain Image Analysis with Arbitrary Modality Availability

"임의의 조영제 가용성을 위한 뇌 영상 분석을 위한 해부학적 인식 교차-모달 사전 학습 BrainAnytime"의 기술적 요약

문제 제기

신경퇴행성 질환, 특히 알츠하이머병 (AD) 에 대한 임상 진단 워크플로는 "조영제 강화 경로 (modality escalation pathway)"를 따릅니다. 임상가는 일반적으로 T1 강조 MRI 와 같은 일상적인 구조적 영상으로 시작하여, 감별 진단을 정교화하기 위해 FLAIR, T2 와 같은 시퀀스를 선택적으로 추가하고, 불확실한 사례를 위해 아밀로이드-PET 과 같은 분자 영상을 보류합니다. 결과적으로 실제 세계 코호트의 환자 데이터는 이질적이고 종종 불완전한 조영제 부분집합을 특징으로 합니다.

현재의 AI 모델들은 일반적으로 고정된 입력 조영제를 가정하여, 데이터 가용성이 변경될 때 아키텍처 전환이나 재학습을 요구합니다. 또한, 기존 대규모 뇌 영상 사전 학습 모델들은 이러한 맥락에서 세 가지 한계를 겪는 경우가 많습니다:

1. 대부분 MRI 만을 대상으로 하여, 중요한 생체 표지자인 아밀로이드-PET 과의 명시적 연결 고리를 학습하지 못합니다.
2. 고정된 조영제 조합을 가정하여, 실제 세계 데이터의 단계 의존적 결손과 불일치합니다.
3. 표준 마스킹 오토인코더 (MAE) 사전 학습은 공간적으로 균일한 마스킹에 의존하여, AD 가 균일한 병리가 아닌 해마와 같은 선택적 지역적 취약성을 보인다는 사실을 간과합니다.

방법론

저자들은 임의의 조영제 가용성 하에서 뇌 영상 분석을 지원하도록 설계된 통합 사전 학습 프레임워크인 BrainAnytime을 제안합니다. 이 프레임워크는 다중 시퀀스 MRI 와 아밀로이드-PET 을 아우르는 5 개의 데이터셋 (A4, DOD-ADNI, AIBL, BraTS23, NACC) 의 34,899 개 3D 뇌 스캔으로 사전 학습됩니다. 공유된 3D 마스킹 오토인코더 아키텍처 내에서 세 가지 핵심 구성 요소를 통합합니다:

1. Multi-MAE3D (3D 다중-모달 마스킹 오토인코더)

• 아키텍처: 공유된 Transformer 인코더가 사용 가능한 조영자의 임의 부분집합 (T1, T2, FLAIR, PET) 을 처리합니다.
• 토큰화: 각 조영자는 조영자별 3D 합성곱 패치 임베딩을 사용합니다.
• 결손 데이터 처리: 누락된 조영자는 완전히 마스킹됩니다. 입력 가용성과 무관하게 일정한 토큰 수를 유지하기 위해, 가시 토큰 예산은 디리클레 분포 (Dir($\alpha$)) 를 사용하여 관측된 조영자들 간에 재분배됩니다.
• 디코딩: 각 조영자는 모든 패치를 재구성하는 경량 디코더를 갖습니다. 훈련 손실 ($L_{MAE}$) 은 패치별 정규화를 사용한 평균 제곱 오차로 마스킹된 패치에 대해서만 계산됩니다.

2. 상호 교차-모달 증류 (RCMD)

• 목적: 구조적 MRI 와 분자적 PET 간의 대응 관계를 학습하기 위함입니다.
• 메커니즘: 지수 이동 평균 (EMA) 교사 네트워크가 인코더 아키텍처를 공유합니다. MRI-PET 쌍이 있는 샘플의 경우, 그룹 수준 표현 ($z_{MRI}$ 및 $z_{PET}$) 이 평균 풀링을 통해 계산됩니다.
• 손실: 두 개의 대칭적 MLP 예측기 ($g_{M\to P}$ 및 $g_{P\to M}$) 가 교차-모달 교사 특징을 예측하도록 훈련됩니다. 손실 ($L_{RCMD}$) 은 반대 조영자 공간에서 학생의 예측된 특징과 교사의 특징 간의 코사인 유사도를 최대화합니다.

3. 병리 인식 커리큘럼 마스킹 (PACM)

• 목적: 균일하게 마스킹하는 대신 질병에 취약한 해부학적 구조를 재구성하는 데 우선순위를 두기 위함입니다.
• 중요도 점수화: AAL116 아틀라스를 매핑한 신경병리학적 사전 지식 (Braak 단계) 에서 정적 점수가 도출됩니다. 동적 점수는 EMA 해부학 교사로부터 주의 맵을 집계하여 얻습니다. 이러한 점수들을 결합하여 패치 중요도 점수를 생성합니다.
• 커리큘럼 일정: 세 단계 전략이 사용됩니다:
  1. 초기 단계: 균일한 무작위 마스킹.
  2. 중기 단계: 온도 매개변수 ($\tau$) 가 코사인 감쇠를 통해 어닐링되어, 마스킹할 고중요도 영역을 선호하도록 마스크 분포를 점진적으로 날카롭게 만듭니다 (즉, 모델이 중요한 영역을 재구성하도록 강제함).
  3. 후기 단계: 고정된 온도.
• 선택: 패치들은 Gumbel-top-k 트릭을 사용하여 선택되어, 고중요도 패치가 마스킹되고 따라서 재구성될 가능성이 더 높도록 보장합니다.

훈련 및 미세 조정:
총 손실은 $L = L_{MAE} + \lambda \cdot \mathbb{1}_{[paired]} \cdot L_{RCMD}$입니다. 하류 작업의 경우, 경량 작업 헤드가 있는 학생 인코더만 사용됩니다. 추론 중에는 누락된 조영자가 0 으로 설정되고 어텐션이 차단되어, 모델이 아키텍처 변경 없이 임의의 부분집합을 처리하도록 보장합니다.

주요 기여

1. 통합 프레임워크: 재학습이나 아키텍처 전환 없이 단일 모델이 (단일 T1 스캔부터 완전한 다중 모달 검사까지) 조영자의 임의 부분집합을 수용할 수 있습니다.
2. 교차-모달 학습: 상호 증류를 통해 MRI 와 PET 간의 구조적 - 분자적 대응 관계를 명시적으로 학습하여, 구조적 증거와 병리학적 증거 간의 간극을 해소합니다.
3. 해부학적 인식 사전 학습: 표준 MAE 의 균일한 마스킹 한계를 해결하기 위해, 모델의 재구성 노력을 AD 와 관련된 신경해부학으로 유도하는 커리큘럼 마스킹 전략을 도입합니다.
4. 확장 가능한 사전 학습: 임상적 결손 패턴을 모방한 대규모 이질적 데이터셋 (34,899 개 스캔) 에서 성공적으로 사전 학습되었습니다.

실험 결과

이 모델은 5 개의 임상적으로 동기화된 조영자 설정 (T1-only 에서 T1+T2+FLAIR+PET 까지) 에서 4 개의 하류 작업 (CN 대 AD, CN 대 MCI, MMSE 회귀, 나이 예측) 에 대해 평가되었습니다.

• 성능: BrainAnytime 은 조영자별 모델 (예: 3D ResNet, M3T), 이중 모달 융합 방법, 대규모 뇌 기반 모델 (BrainIAC, SAM-Brain3D), 그리고 불완전 다중 모달 학습 접근법 (Flex-MoE, FuseMoE) 보다 우수한 성능을 발휘했습니다.
• 주요 지표:
  • CN 대 AD: 평균 정확도 **89.4%**를 달성하여 (다음으로 가장 좋은 기준인 Flex-MoE 의 84.2% 대비), 가장 강력한 결손 조영자 기준 대비 **6.2%**의 상대적 개선을 보였습니다.
  • CN 대 MCI: 평균 정확도 **67.2%**를 달성하여 (MoE-retriever 의 60.6% 대비), **7.0%**의 상대적 개선을 보였습니다.
  • 회귀: 모든 조영자 설정에서 MMSE 및 나이 예측에서 우수한 성능을 입증했습니다.
• 강건성: 동일 피험자 내 강건성 연구는 더 많은 조영자가 추가됨에 따라 성능이 일관되게 향상됨을 확인했습니다 (예: T1-only 의 65.0% 에서 완전 입력의 80.4% 로 CN 대 AD F1 증가), 이는 모델이 사용 가능한 데이터를 효과적으로 활용함을 보여줍니다.
• 어블레이션: RCMD 나 PACM 을 개별적으로 제거하면 성능이 저하되어, 교차-모달 증류와 해부학적 인식 마스킹의 필요성을 확인했습니다.

의의

이 논문은 BrainAnytime 이 임상적 현실 (단계별, 불완전한 데이터 획득) 과 현재 AI 패러다임 (고정 입력 모델) 사이의 중요한 간극을 해소한다고 주장합니다. 단일 모델이 사용 가능한 모든 영상을 유연하게 수용할 수 있게 함으로써, 이는 임상적 조영자 강화 경로와 부합합니다. 이 프레임워크는 교차-모달 증류와 병리 인식 마스킹을 통한 사전 학습이 기존 기준, 특히 데이터가 불완전하거나 조영자가 누락된 시나리오에서 크게 우수한 강건한 기반 모델을 산출함을 보여줍니다. 이 접근법은 모든 환자에게 완전한 다중 모달 데이터가 거의 제공되지 않는 실제 임상 환경에서 AI 를 배포하기 위한 실용적인 해결책을 제공합니다.