기술 요약: 뇌 MRI를 위한 MultiMAE

문제 정의

의료 영상, 특히 뇌 MRI 분야에서 딥러닝 모델은 입력 시퀀스 누락(missing input sequences) 문제에 직면하곤 합니다. 임상 환경에서는 획득 문제나 환자의 제약으로 인해 완전한 다중 모달 데이터(예: T1, T1c, T2 또는 FLAIR 시퀀스 누락)가 부족한 경우가 빈번합니다. 일반적인 접근 방식은 모든 가용 모달리티를 채널 차원(RGB 채널과 유사)으로 쌓는 초기 융합(early-fusion) 전략을 사용합니다. 이 방식은 단순하지만, 모든 모달리티가 학습과 추론 단계 모두에서 존재해야 한다는 전제가 필요합니다. 결과적으로, 초기 융합 방식으로 학습된 모델은 입력이 누락될 경우 성능이 급격히 저하되며, 교차 모달 추론(cross-modal reasoning)에 활용할 수 있는 뚜렷한 모달리티별 표현을 학습하는 데 실패합니다.

방법론

저자들은 3D 뇌 MRI에 적합하도록 조정된 다중 모달 마스크드 오토인코더(Multi-Modal Masked Autoencoder, MultiMAE) 프레임을 제안합니다. 자연 이미지를 위한 MultiMAE 아키텍처에서 영감을 얻은 이 방법은 초기 융합에서 벗어나, 트랜스포머 인코더를 통한 **지연 융합(late-fusion)**과 결합된 모달리티별 인코딩(modality-specific encoding) 전략으로 전환합니다.

아키텍처 및 사전 학습

• 백본(Backbone): 모델은 수정된 Vision Transformer (ViT-B/16)를 인코더로 활용합니다.
• 모달리티별 어댑터(Modality-Specific Adapters): 채널을 쌓는 대신, 각 MRI 시퀀스(t1, t1c, t2, fla)는 고유한 선형 어댑터를 통해 통과되어 중첩되지 않는 $16 \times 16 \times 16$ 패치를 공유 토큰 공간으로 투영합니다.
• 마스킹 전략: 전역 마스킹 비율 75%가 적용됩니다. 결정적으로, 마스킹은 디리클레 분포($\alpha=1$)로부터 모달리티별로 샘지되어, 특정 모달리티의 최대 100%까지 마스킹하면서도 전체 비율을 유지할 수 있습니다. 이는 모델이 교차 모달 의존성을 학습하도록 강제합니다.
• 인코더-디코더 설계:
  • 인코더: 표준 ViT-B/16 (12개 레이어)이 [cls] 토큰과 함께 모든 모달리티의 마스킹되지 않은 토큰을 처리합니다.
  • 디코더: 4개의 경량화된 모달리티별 디코더(각 3개 레이어)가 사용됩니다. 이들은 쿼리(누락된 패치를 위한 학습 가능한 [mask] 토큰)가 인코더가 제공하는 공유 컨텍스트를 참조(attend)하는 교차 주의(cross-attention) 메커니즘을 활용합니다.
• 목적 함수: 모델은 MSE 손실을 사용하여 마스킹된 패치를 재구성하도록 사전 학습되며, 이를 통해 가용한 시퀀스로부터 누락된 시퀀스를 추론할 수 있는 풍부한 표현을 학습하도록 유도합니다.

다운스트림 적응

사전 학습된 인코더는 두 가지 다운스트림 태스크에 적응됩니다:

1. 세그멘테이션(Segmentation): UNETR 아키텍처를 사용합니다. 모달리티별 개별 토큰화로 인한 가변적인 시퀀스 길이를 처리하기 위해, 각 공간 위치의 토큰들을 평균하여 일관된 시퀀스 길이를 유지합니다.
2. 분류(Classification): 단일 선형 레이어가 디코더를 대체합니다. 패치 토큰과 [cls] 토큰을 평균하여 신경교종(glioma) 아형(GBM, Astro, Oligo) 분류를 위한 특징 벡터를 형성합니다.

주요 기여

1. 3D MRI를 위한 MultiMAE 사전 학습: 본 연구는 3D 뇌 MRI에 MultiMAE를 적용하여 모달리티별 인코딩과 다중 시퀀스 간의 마스크드 모델링을 가능하게 하는 방안을 조사했습니다.
2. 누락된 모달리티에 대한 강건성: 저자들은 이 사전 학습 전략이 세그멘테이션(신경교종 및 전이암)과 분류 태스크 모두에서 입력 모달리티가 누락되었을 때 성능과 강건성을 크게 향상시키며, 표준 MAE-ViT 베이스라인보다 우수함을 입증했습니다.
3. 누락된 모달리티의 합성: 이 프레임워크는 사전 학습 과정에서 학습된 교차 모달 관계를 활용하여 추론 시점에 완전히 누락된 MRI 시퀀스를 합성할 수 있는 능력을 보여줍니다.

실험 결과

방법론은 대규모 결합 데이터셋(BraTS 2021/2023, UPENN-GBM, UCSF-PGDM 등)에서 평가되었으며, 내부 및 외부 데이터셋에서 테스트되었습니다.

• 누락된 모달리티 생성: 모델은 높은 충실도로 누락된 모달리티를 성공적으로 재구성했습니다 (PSNR ~20–21, SSIM ~0.71–0.76). 이는 표준 MAE-ViT 베이스라인(PSNR ~7–10, SSIM ~0.52–0.61)을 크게 상회하는 수치입니다.
• 세그멘테이션 성능:
  • 사전 학습된 MultiMAE는 내부 및 외부 테스트 세트 모두에서 신경교종 및 전이암 세그멘테이션에 대해 MAE-ViT 베이스라인보다 일관되게 우수한 성능을 보였습니다.
  • 누락된 모달리티 시나리오에서 MultiMAE는 탁월한 회복력을 보여주었습니다. 예를 들어, T1c 모달리티가 누락되었을 때, MultiMAE는 종양 핵(tumor core)에 대해 약 51.55의 Dice 점수(인코더 동결 시)를 유지한 반면, ViT 베이스라인은 약 47.91로 떨어졌습니다.
  • 이 모델은 입력이 누락된 상황에서 베이스라인 대비 전체 Dice 점수에서 10.1, MCC에서 0.46의 절대적 개선을 달이트했습니다.
• 분류 성능:
  • 모델은 신경교종 아형 분류를 위해 강력한 다중 클래스 정확도와 MCC(최대 0.70)를 달성했습니다.
  • 세그멘테이션 태스크와 달리, 분류에서는 인코더를 동결하는 것이 덜 효과적이었으며, 전체 파인튜닝(full fine-tuning)이 가장 좋은 결과를 냈습니다.
  • 누락된 모달리티(예: T1c 없음)가 있는 경우에도 MultiMAE는 0.89의 높은 정확도를 유지하여, ViT 베이스라인의 0.41과 비교해 효과적인 교차 모달 추론 능력을 입증했습니다.

의의 및 주장

본 논문은 MultiMAE 패러다임이 불완전한 데이터라는 핵심적인 임상적 과제를 해결하는, 뇌 MRI를 위한 유연하고 일반화 가능한 인코더를 제공한다고 주장합니다. 각 시퀀스를 별도의 모달리티로 취급하고 교차 모달 재구성을 강제함으로써, 모델은 가용한 입력으로부터 누락된 시퀀스를 추론하는 법을 배웁니다.

저자들은 이 접근 방식이 ViT 백본에 미세한 변화를 도입하면서도 강건성 측면에서 상당한 이득을 가져온다는 점을 강조합니다. 이러한 사전 학습 전략은 "보지 못한 입력 구성(unseen input configurations)"을 재학습 없이도 처리할 수 있게 하여, 데이터 완결성을 보장할 수 없는 자원 제한적인 임상 환경에 적합하게 만듭니다. 다만, 저자들은 다음과 같은 점을 겸허히 언급했습니다:

• 픽셀 단위 MSE 손실의 한계로 인해 재구성 결과에 일부 흐릿함(blurring)이 나타납니다.
• 내부 데이터셋과 외부 데이터셋 사이에 도메인 격차가 존재하며, 이는 광범위한 사이트 적응(cross-site adaptation)의 과제를 반영합니다.
• 교차 모달 통합의 복잡성으로 인해 상당한 사전 학습 연산량과 데이터가 수렴에 필요하며, 무작위 초기화된 버전의 모델은 일부 경우에서 오히려 성능이 낮았습니다.

궁극적으로, 이 연구는 다중 모달 뇌 MRI 데이터를 더욱 우아하게 처리할 수 있도록 MAE 사전 학습 패러다임을 재구성하였으며, 누락된 입력에 대한 강건성을 우선시하였습니다.