Modality-Aware and Anatomical Vector-Quantized Autoencoding for Multimodal Brain MRI

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 뇌 질환 진단을 위해서는 T1, T2, FLAIR 등 다양한 조영 대조 (Contrast) 를 가진 다중 모달리티 (Multi-modal) MRI 가 필수적입니다. 그러나 운동 아티팩트, 시간 및 비용 제약으로 인해 모든 피험자에게 고품질의 다중 모달리티 데이터를 획득하는 것은 어렵습니다. 이를 해결하기 위해 한 모달리티에서 다른 모달리티로의 MRI 합성 (Synthesis) 기술이 주목받고 있습니다.
• 기존 방법의 한계:
  • 단일 모달리티 집중: 기존 뇌 VAE(변분 오토인코더) 들은 주로 T1 위주 단일 모달리티 데이터에 집중하여, T2 등 다른 모달리티의 보완적 진단 가치를 간과합니다.
  • 데이터 의존성: 많은 방법들이 특정 데이터셋에 맞춰 VAE 를 처음부터 훈련하거나 (cDPM, BrLP 등), 자연 이미지로 사전 훈련된 2D VAE 를 슬라이스 단위로 적용 (SD-VAE) 합니다.
  • 구조적 불일치: 2D 기반 모델은 3D 뇌 구조의 연속성을 유지하지 못해 슬라이스 간 경계가 불연속적이거나 아티팩트가 발생합니다. 반면, 범용 의료 VAE(MedVAE 등) 는 모달리티 간 차이를 명시적으로 모델링하지 않아 뇌 해부학적 구조와 모달리티 특유의 외관 (Appearance) 을 효과적으로 분리하지 못합니다.
• 목표: 다양한 모달리티 간에 공유되는 **해부학적 구조 (Anatomical Structure)**와 모달리티에 의존하는 **외관 (Appearance)**을 명확히 분리 (Disentangle) 하면서도, 3D 공간 일관성을 유지하는 강력한 VAE 를 개발하는 것.

2. 제안 방법론: NeuroQuant (Methodology)

논문에서 제안한 NeuroQuant는 다중 모달리티 뇌 MRI 를 위한 통합된 3D 벡터 양자화 오토인코더 (3D VQ-VAE) 입니다. 주요 구성 요소는 다음과 같습니다.

가. 이중 스트림 아키텍처 (Dual-Stream Architecture)

• 공유 인코더: 입력 MRI 볼륨을 처리하는 단일 3D 인코더를 사용하여, 모달리티 불변의 해부학적 정보와 모달리티 의존적 외관 정보를 분리합니다.
  • 해부학적 스트림 ($z_{anat}$): 모든 모달리티에 공통적인 뇌 구조 정보를 인코딩합니다.
  • 모달리티 스트림 ($z_{mod}$): 각 MRI 모달리티 (T1, T2 등) 고유의 대비 (Contrast) 및 강도 분포 정보를 인코딩합니다.

나. 분해된 다축 어텐션 (Factorized Multi-Axis Attention)

• 문제: 전체 3D 셀프 어텐션은 계산 비용이 너무 높음 ($O((DHW)^2)$).
• 해결: 축 (Axial, Coronal, Sagittal) 단위로 어텐션을 순차적으로 수행하여 계산 복잡도를 $O(DHW(D+H+W))$로 줄이면서도 장기적인 뇌 영역 간의 의존성을 포착합니다. 이는 슬라이스 간 및 모달리티 간 맥락 정보를 효과적으로 교환하게 합니다.

다. 벡터 양자화 및 공유 코드북 (Vector Quantization & Shared Codebook)

• 해부학적 잠재 변수 ($z_{anat}$) 는 공유된 벡터 양자화 코드북을 통해 이산화된 토큰 (Anatomical Tokens) 으로 변환됩니다.
• 이는 해부학적 구조를 컴팩트한 사전 (Vocabulary) 으로 표현하여, 다른 모달리티에서도 일관된 구조 정보를 유지하도록 합니다.

라. 모달리티-FiLM 디코더 (Modality-FiLM Decoder)

• 디코딩 단계에서 양자화된 해부학적 토큰을 기반으로 3D 볼륨을 재구성합니다.
• FiLM (Feature-wise Linear Modulation): 모달리티 스트림 ($z_{mod}$) 에서 추출된 특징을 사용하여 디코더의 각 레이어에 **FiLM 파라미터 ($\gamma, \beta$)**를 동적으로 적용합니다. 이를 통해 구조는 유지하면서 모달리티 특유의 대비와 강도 변화를 적응적으로 재현합니다.

마. 2D/3D 결합 훈련 전략 (Joint 2D/3D Training)

• 3D MRI 데이터의 슬라이스 기반 획득 특성을 고려하여, 전체 3D 볼륨과 무작위 2D 슬라이스를 혼합하여 훈련합니다.
• 3D 훈련은 볼륨 일관성을, 2D 훈련은 고주파수 디테일과 슬라이스 수준의 정확도를 향상시킵니다.

바. 손실 함수 (Loss Functions)

• 재구성 손실: 픽셀 단위 (L1) 와 구조적 유사성 (3D SSIM) 을 함께 최적화합니다.
• 교차 모달 재구성 (Cross-modal Reconstruction): 한 모달리티의 해부학적 코드와 다른 모달리티의 FiLM 파라미터를 결합하여 재구성함으로써, 해부학 코드가 모달리티 불변임을 학습시킵니다.
• 모달리티 적대적 손실 (Modality-Adversarial Loss): 해부학적 잠재 공간에서 모달리티 정보를 제거하기 위해 적대적 학습 (Gradient Reversal) 을 적용하여 해부학과 외관의 분리를 강화합니다.

3. 주요 기여 (Key Contributions)

1. NeuroQuant 프레임워크: 다중 모달리티 뇌 MRI 를 위해 해부학적 구조와 모달리티 특유 외관을 공유 잠재 공간에서 분리하여 학습하는 최초의 통합 3D VQ-VAE 를 제안했습니다.
2. 새로운 아키텍처 설계: 분해된 다축 어텐션, VQ 기반 해부학적 토큰, 모달리티-FiLM 디코더를 결합하여 슬라이스 간 일관성과 모달리티 적응성을 동시에 달성했습니다.
3. 성능 입증: NCANDA 와 ABCD 두 개의 대규모 데이터셋에서 기존 SOTA(VQGAN, SD-VAE, MediTok, MedVAE 등) 를 모든 평가 지표에서 압도적으로 능가함을 입증했습니다.

4. 실험 결과 (Results)

• 데이터셋: NCANDA (청소년 뇌 발달) 와 ABCD (장기 소아 뇌 영상) 데이터셋의 T1/T2 쌍을 사용했습니다.
• 정량적 평가:
  • 재구성 품질: PSNR 및 SSIM 지표에서 모든 베이스라인 모델을 상회했습니다 (예: NCANDA T1 기준 PSNR 28.89, SSIM 95.27).
  • 해부학적 일관성: SynthSeg 를 이용한 뇌 영역 분할 (Dice Score) 에서 가장 높은 점수를 기록하여, 재구성된 이미지가 실제 해부학적 구조를 정확하게 반영함을 증명했습니다.
  • 잠재 공간 유용성: 학습된 잠재 벡터를 이용한 성별 분류 (Sex Classification) 정확도에서 가장 높은 성능을 보여, 잠재 공간이 의미 있는 정보를 잘 포착하고 있음을 입증했습니다.
• 정성적 평가:
  • 기존 2D 모델들은 슬라이스 간 경계에서 아티팩트 (계단 현상, 불연속) 가 발생했으나, NeuroQuant 는 세 가지 방향 (축, 관상, 시상) 에서 모두 선명한 에지와 매끄러운 조직 경계를 유지했습니다.
  • MedVAE 와 같은 3D 모델은 과도한 평활화 (Oversmoothing) 로 인해 미세한 구조가 흐려지는 반면, NeuroQuant 는 미세한 피질 주름과 뇌실 경계를 선명하게 복원했습니다.
• Ablation Study:
  • 2D/3D 결합 훈련, 분해된 어텐션 (FMA), FiLM 디코더, 적대적 손실 등 각 구성 요소가 제거될 때 성능이 저하됨을 확인하여, 모든 요소가 필수적임을 입증했습니다.

5. 의의 및 결론 (Significance)

• 확장 가능한 기반 모델: NeuroQuant 는 단일 모달리티에 국한되지 않고 다양한 MRI 조영 대조를 통합적으로 처리할 수 있는 확장 가능한 기반 모델 (Foundation Model) 의 토대를 마련했습니다.
• 하류 작업 지원: 고품질의 합성 MRI 생성, 데이터 증강, 인과 분석 (Counterfactual generation), 그리고 질병 메커니즘 탐구 등 다양한 하류 생성 및 분석 작업에 강력한 잠재 표현을 제공합니다.
• 의료 영상 AI 의 발전: 자연 이미지 기반 모델의 한계를 넘어, 의료 영상 특유의 3D 구조적 일관성과 모달리티 간 상호 보완적 관계를 명시적으로 모델링한 새로운 패러다임을 제시했습니다.

이 연구는 뇌 MRI 분석에서 **구조 (Structure)**와 **외관 (Appearance)**을 효과적으로 분리하면서도 3D 공간의 일관성을 유지하는 새로운 표준을 제시했다는 점에서 의의가 큽니다.