CLIMB: Controllable Longitudinal Brain Image Generation using Mamba-based Latent Diffusion Model and Gaussian-aligned Autoencoder

1. 문제 정의 (Problem)

• 배경: 신경퇴행성 질환 (예: 알츠하이머) 의 조기 진단 및 치료 계획 수립을 위해 환자의 뇌 구조가 시간에 따라 어떻게 변화하는지 (종단적 변화) 를 예측하는 것은 매우 중요합니다. 이를 위해 기준선 MRI 와 함께 미래의 뇌 MRI 를 합성하는 연구가 활발합니다.
• 기존 기술의 한계:
  • 계산 비용: 기존 확산 모델 (Diffusion Models) 은 픽셀 공간에서 반복적인 탈노이즈 과정을 거치므로 계산 비용이 매우 높고 추론 속도가 느립니다.
  • 잠재 공간의 문제: 잠재 확산 모델 (LDM) 은 계산 효율성을 높이기 위해 잠재 공간 (Latent Space) 을 사용하지만, 기존 방식인 변분 오토인코더 (VAE) 는 잠재 특징을 확률적으로 샘플링할 때 노이즈가 발생하여 재구성된 이미지의 선명도가 떨어지는 (blurry) 문제가 있습니다.
  • 아키텍처의 비효율성: 기존 LDM 은 주로 트랜스포머 (Transformer) 기반의 셀프 어텐션 (Self-Attention) 을 사용하여 전역 문맥을 포착하지만, 이는 3 차원 의료 영상 처리 시 메모리 사용량이 많고 계산 복잡도가 높습니다.
  • 통제 가능성 부족: 많은 기존 모델이 시간적 진화를 효과적으로 모델링하거나 나이, 성별, 유전 정보, 뇌 구조 부피 등 다양한 임상 변수를 조건으로 활용하여 합성하는 데 한계가 있습니다.

2. 제안 방법론 (Methodology)

저자들은 CLIMB (Controllable Longitudinal brain Image generation via MamBa-based latent diffusion model) 라는 새로운 프레임워크를 제안했습니다. 이는 크게 두 단계로 구성됩니다.

A. 가우스 정렬 오토인코더 (Gaussian-aligned Autoencoder, GATE)

• 목적: 확산 모델이 효율적으로 작동할 수 있도록 잠재 공간이 가우스 분포를 따르도록 정렬하면서도, VAE 의 샘플링 노이즈로 인한 이미지 흐림을 방지합니다.
• 핵심 기법:
  • 슬라이스드 워셔스타인 거리 (Sliced Wasserstein Distance): 잠재 특징의 분포를 사전 분포 (가우스) 에 정렬하기 위해 사용합니다. 이는 확률적 샘플링 대신 결정론적 (deterministic) 정렬 방식을 사용하여 재구성 이미지의 선명도와 충실도를 유지합니다.
  • 손실 함수: 재구성 손실, 슬라이스드 워셔스타인 거리 손실, 지각적 손실 (Perceptual Loss), 그리고 적대적 손실 (Adversarial Loss) 을 결합하여 고품질의 잠재 표현을 학습합니다.

B. Mamba 기반 잠재 확산 모델

• 아키텍처: U-Net 스타일의 인코더 - 디코더 구조를 사용하며, 기존 트랜스포머의 셀프 어텐션 대신 Mamba (State Space Model, SSM) 아키텍처를 백본으로 채택했습니다.
• 장점: Mamba 는 선형 시간 복잡도 (linear-time complexity) 를 가지며 전역 문맥을 효율적으로 포착하여, 트랜스포머 대비 계산 오버헤드를 크게 줄이고 메모리 효율성을 높입니다.
• 조건부 입력 (Controllability):
  • 입력 데이터: 기준선 MRI 스캔, 획득 시점의 나이.
  • 조건 변수: 예측된 미래 나이, 성별, 질병 상태, 유전 정보, 뇌 구조 부피 (해마, 회백질, 백질, 뇌실, 편도체 등).
  • 예측 모델: IRLSTM 모델을 활용하여 미래 시점의 질병 상태와 뇌 부피를 예측한 후, 이를 확산 모델의 조건으로 주입합니다.
• 학습 전략: 두 단계 학습을 통해 안정성을 확보합니다. 1 단계에서는 예측 변수만으로 확산 모델을 학습하고, 2 단계에서는 이미지 특징 처리 네트워크를 추가하여 미세 조정합니다.

C. 추론 과정

• 기준선 MRI 를 GATE 인코더로 인코딩하고, IRLSTM 으로 미래 시점의 임상 변수를 예측합니다.
• Mamba 기반 탈노이즈 네트워크를 통해 잠재 공간에서 미래 MRI 의 잠재 표현을 샘플링합니다.
• GATE 디코더를 통해 최종 미래 MRI 이미지를 생성합니다. (DDIM 을 사용하여 25 단계로 빠른 샘플링 구현)

3. 주요 기여 (Key Contributions)

1. CLIMB 모델 제안: 기준선 MRI 와 다양한 임상 변수 (나이, 성별, 질병 상태, 뇌 부피 등) 를 활용하여 시간 경과에 따른 뇌 구조 변화를 제어 가능하게 합성하는 새로운 프레임워크.
2. GATE (Gaussian-aligned Autoencoder) 도입: VAE 의 샘플링 노이즈를 제거하고 슬라이스드 워셔스타인 거리를 사용하여 잠재 공간을 가우스 분포에 정렬함으로써, 고선명도의 재구성 이미지를 보장.
3. Mamba 아키텍처 적용: 트랜스포머 기반의 고비용 어텐션 메커니즘을 Mamba 로 대체하여 계산 복잡도를 줄이면서도 전역 문맥을 효과적으로 포착하는 효율적인 확산 모델 구현.

4. 실험 결과 (Results)

• 데이터셋: 알츠하이머병 신경영상 이니셔티브 (ADNI) 의 1,390 명 참가자, 총 6,306 건의 MRI 스캔 사용.
• 비교 대상: SADM, BrLP 등 기존 최첨단 (SOTA) 모델.
• 정량적 평가:
  • SSIM (구조적 유사성): 0.9433 (기존 모델 대비 가장 높음).
  • PSNR: 27.82.
  • LPIPS (지각적 유사성): 0.0587 (가장 낮음, 즉 실제 이미지와 가장 유사).
  • MSE: 2.01e-3 (가장 낮음).
  • 효율성: 이미지당 2.92 초의 빠른 추론 속도와 5,699 MB 의 낮은 메모리 사용량을 기록하여 기존 모델보다 효율적임.
• 정성적 평가: 생성된 이미지가 실제 MRI 와 해부학적으로 일치하며, 특히 뇌 중심부의 구조적 오류나 아티팩트가 적고 시간 경과에 따른 변화가 자연스럽게 표현됨.

5. 의의 및 결론 (Significance)

• 임상적 가치: CLIMB 는 환자의 미래 뇌 MRI 를 예측하여 질병 진행을 모델링하고, 결측 데이터를 보완하며, 임상 의사결정을 지원하는 데 기여할 수 있습니다. 이는 조기 진단 및 치료 계획 수립에 유용합니다.
• 기술적 혁신: 의료 영상 생성 분야에서 Mamba 아키텍처와 가우스 정렬 오토인코더를 결합하여, 고품질 생성과 계산 효율성을 동시에 달성한 선구적인 연구입니다.
• 향후 과제: SSIM/MSE 수치적 성과는 우수하나, 실제 임상적/해부학적 분석을 위해서는 방사선사나 신경과 전문의에 의한 검증이 필요함을 강조했습니다.

요약하자면, 이 논문은 계산 효율성을 극대화하면서도 (Mamba 사용), 이미지 선명도를 유지 (GATE 사용) 하고, 다양한 임상 변수를 통해 미래 뇌 영상을 정밀하게 예측 (CLIMB) 할 수 있는 새로운 의료 영상 생성 프레임워크를 제시했습니다.