A continually expandable foundation model for brain MRI

기술 요약: 뇌 MRI를 위한 지속적으로 확장 가능한 파운데이션 모델

문제 정의
현재의 뇌 자기공명영상(MRI)용 계산 모델은 매우 파편화되어 있으며, 일반적으로 특정 질병, 인구 집단 또는 영상 프로토콜에 맞춰 개발됩니다. 파운데이션 모델은 대규모 데이터셋으로부터 재사용 가능한 표현(representation)을 학습함으로써 보다 일반화된 시스템으로 가는 경로를 제시하지만, 대개 "사전 학습 후 적응(train-then-adapt)"이라는 정적인 패러 Paradigm을 따릅니다. 이러한 접근 방식의 결정적인 한계는 새로운 데이터(예: 새로운 질병 코호트 또는 영상 양식)로 사전 학습된 모델을 업데이트할 때 이전에 습득한 능력이 저하되는 "파괴적 망각(catastrophic forgetting)"입니다. Elastic Weight Consolidation(EWC)과 같은 정규화 기반 접근 방식이나 마스킹 전략과 같은 기존의 지속 학습(continual-learning) 방법들은 이질적인 임상 영역 전반에 걸쳐 대규모 3차원 볼륨 모델을 확장하는 동안 이전의 계산 경로를 보존하는 데 어려움을 겪습니다.

방법론: Alcmaeon 프레임워크
저자들은 대규모 자기지도 사전 학습 및 순차적 도메인 확장을 위해 설계된 3차원 뇌 MRI 파운데이션 모델인 Alcmaeon을 소개합니다. 이 프레임워크는 네 가지 주요 구성 요소로 이루어져 있습니다:

1. 아키텍처:

   • 3D-SWIN: 공유된 잠재 표현(latent representation)을 통해 구조적 MRI, 확산 유도 영상 및 정량적 지도를 재구성하는 계층적 시프트 윈도우 Swin Transformer(SwinUNETR에서 변형됨) 기반의 볼륨 인코더-디코더입니다.
   • 3D-DiT: 압축된 병목(bottleneck) 내의 잠재 표현 분포를 모델링하여 이미지 합성 및 교차 모달 변환을 지원하는 볼륨 확산 트랜스포머(Diffusion Transformer)입니다.
   • 태스크 헤드(Task Heads): 학습된 표현을 진단 및 예후 엔드포인트에 연결합니다.

2. 지속 학습 메커즘 (그래프-블루프린트 프루닝 - GBP):

   • 파괴적 망각을 해결하기 위해 저자들은 구조적 메모리 메커즘인 GBP를 제안합니다.
   • 돌출도 순위 지정(Salience Ranking): 각 새로운 도메인에 대해 모듈(어텐션, 피드포워드, 정규화, 적응형 변조)을 활성화-그래디언트 돌출도를 기준으로 순위를 매깁니다.
   • 블루프린트 구축: 높은 돌출도를 가진 모듈은 누적된 "블루프린트"에 추가되고 이후의 적응 과정 동안 동결(보호)됩니다. 나머지 모듈은 학습 가능한 상태로 남습니다.
   • 검사 가능성(Inspectability): 인코더-디코더와 잠재 생성기를 위한 별도의 블루프린트가 유지되어, 어떤 구조적 구성 요소가 보존되고 어떤 것이 새로운 학습을 위해 사용 가능한지에 대한 기록을 제공합니다.

3. 학습 시퀀스:

   • 모델은 UK Biobank(건강한 인구 집단)로부터 425,000개 이상의 볼륨을 사용하여 사전 학습되었습니다.
   • 그 후 네 가지 뚜렷한 임상 도메인에 걸쳐 순차적으로 확장되었습니다:
     • D1: 건강한 노화 (HABS-HD).
     • D2: 신경 퇴행성 질환 (PPMI, ADNI, PREVENT-AD).
     • D3: 발달 및 정신 의학 영상 (ABCD, RBC).
     • D4: 성인 뇌종양 (UPenn-GBM, UCSF-PDGM, Erasmus Glioma Database).
   • 실험은 다양한 모델 규모(Small에서 Extra-Large까지), 잠재 설계(결정론적 vs 변이형), 그리고 입력 구성(단일 채널 vs 이중 채널)을 평가했습니다.

주요 결과

• 재구성 및 스케일링:

  • 모델 파라미터 수를 증가시키는 것(약 33M에서 ~676M까지)은 평가된 조건 하에서 재구성 충실도(reconstruction fidelity)에 있어 완만한 이득만을 가져왔습니다.
  • 잠재 설계: 결정론적 잠재 표현이 구조적 및 확산 MRI에 대해 변이형 오토인코더(VAE)보다 우수한 성능을 보였습니다. 그러나 변이형 병목은 NODDI 유도 미세구조 지도에 더 우수했습니다.
  • 생성적 영향: 3D-DiT 생성기를 추가하는 것은 재구성 충실도에는 최소한의 영향을 미쳤으나 다운스트림 합성 작업을 가능하게 했습니다.

• 지속 학습 성능:

  • 망각 완화: GBP는 새로운 도메인에 적응할 때 이전 도메인의 성능을 보존하는 데 있어 순차적 적응(SEQ) 및 EWC보다 크게 뛰어난 성능을 보였으며, 특히 종양 영상(D4)으로 전환된 이후에 두드러졌습니다. (예: 단일 채널 설정에서 D1에 대한 PSNR이 EWC의 10.1 dB 대비 GBP는 25.6 dB를 유지함).
  • 지표: GBP는 단일 채널, 이중 채널 및 하이브리드 입력 구성 전반에서 PSNR, MSE, RMSE, MAE 지표에 대해 더 낮은 피크 망각(peak forgetting)을 보여주었습니다.
  • 트레이드오프: GBP는 유지 능력 면에서는 탁월했지만, 제약 없는 방법들과 비교했을 때 새로운 도메인에 대한 학습 정확도에서 약간의 비용이 발생했으며, 구조적 제약으로 인해 SSIM 점수가 때때로 더 낮았습니다.

• 다운스트림 태스크 적응:

  • 교차 모달 합성: 3D-SWIN 인코더로부터 얻은 표현은 고충실도의 교차 모달 종양 MRI 변환(예: T1에서 T2, T2에서 FLAIR)을 지원했습니다.
  • 생존 모델링: GBP-적응 모델의 동결된 표현은 예후 정보를 유지하여 EWC(0.601)와 대등한 일치 지수(concordance indices)를 달달성하였고, 여러 베이스라인 인코더보다 우수한 성능을 보였습니다.
  • 질병 분류: GBP-적응 모델의 중간 인코더 특징(S8-Lv2)은 알츠하이머병 분류에서 SEQ 및 EWC 변형 모델보다 가장 높은 평균 정확도를 달-달성했습니다.
  • 수술 후 예측: GBP는 특히 제한된 라벨과 상당한 해부학적 변화(종양 질량에서 절제 구멍으로의 변화)를 다룰 때, 퓨샷(few-shot) 수술 후 결과 예측에서 가장 명확한 이점을 보였습니다.

의의 및 주장
본 논문은 Alcmaeon이 모든 질병이나 프로토콜에 대해 별도의 사전 학습된 모델을 요구하지 않고도, 새로운 데이터와 함께 성장할 수 있는 뇌 MRI 파운데이션 모델을 향한 실행 가능한 경로를 제공한다고 주장합니다. 핵심 기여는 단일한 보편적 표현이 아니라, 확장 가능한 태스크 관련 표현 세트와 **검사 가능한 메커니즘(GBP)**을 결$합하여 이전의 지식을 보호하는 것입니다.

저자들은 다음을 강조합니다:

1. 구조적 지속 학습: 블루프린트를 통해 특정 계산 모듈을 보호함으로써, 모델이 이전의 뚜렷한 코호트(예: 건강한 노화)로부터 학습된 능력을 유지하면서도 새롭고 병리학적으로 구별되는 도메인(예: 종양)을 학습할 수 있습니다.
2. 표현 특이성: 모델 아키텍처의 서로 다른 수준은 서로 다른 목적을 수행합니다. 중간 인코더 특징은 분류에 최적이며, 결합된 멀티모달 표현은 위험 순위 지정 및 합성을 지원합니다.
3. 실용적 유용성: 이 프레임워크는 다양한 인구 집단과 스캐너에 걸쳐 영상 데이터가 점진적으로 축적되는 실제 상황을 다루며, 파괴적 망포 없이 새로운 데이터를 통합하는 방법을 제공합니다.

저자들은 임상 배포와 관련하여 신중한 입장을 유지하며, 수술 후 평가가 소규모 코호트(49명 환자)를 대상으로 이루어졌음을 언급하고, 모델이 환자 수준의 결정을 내리는 데 기여하기 위해서는 전향적이고 다기관적인 검증이 필요하다고 밝혔습니다. 또한, 블루프린트에서 파생된 해부학적 투영은 모델 민감도에 대한 탐색적 시각화이지 확립된 인구 수준의 생물학적 패턴이 아님을 강조했습니다.