AD-DAE: Alzheimer's Disease Progression Modeling with Unpaired Longitudinal MRI using Diffusion Auto-Encoders

기술 요약: AD-DAE – 디퓨전 오토인코더를 이용한 비쌍생(Unpaired) 종단적 MRI 기반 알츠하이머병 진행 모델링

문제 정의
알츠하이머와 같은 신경 퇴행성 질환의 진행을 모델링하려면 종단적 시간 경과에 따른 고차원 이미지 분포를 포착해야 한다. 이전 시점의 MRI로부터 후속 이미지를 예측하기 위해 생성 모델(예: VAEs, GANs)이 적용되어 왔으나, 기존 방법론들은 상당한 한계에 직면해 있다. 많은 접근 방식은 피험자별로 쌍을 이룬(paired) 종단적 데이터에 의존하는데, 이러한 데이터는 흔히 부족하다. VAE를 사용하는 방식은 피험자의 정체성(identity)을 질병 관련 진행과 분리하는 데 어려움을 겪으며, 결과적으로 정체성을 보존하지 못하고 흐릿한 재구성 이미지를 생성한다. 반면 GAN은 명시적인 쌍생 가이드 없이는 진행의 방향과 크기를 정밀하게 제어하기 어렵다. 또한, 잠재 디퓨전 모델(Latent Diffusion Models)은 별도의 인코더-디코더 파이프라인에 의존하여 이미지 생성 중의 제어력을 제한한다. 핵심 과제는 쌍생 종단적 감독 없이도 단일 시점의 이미지만을 사용하여, 피험자의 정체성을 보존하면서도 해부학적으로 타당하고 질병 관련 영역에 국한된 변화를 보이는 제어 가능한 후속 이미지를 생성하는 것이다.

방법론: AD-DAE 프레임워크
저자들은 압축되고 제어 가능한 잠재 공간에서 질병 진행을 모델링하기 위해 디퓨전 오토인코더(DAE)를 활용하는 AD-DAE(Alzheimer's Disease Diffusion Auto-Encoder) 프레임워크를 제안한다. 이 아키텍처는 세 가지 주요 구성 요소로 이루어져 있다:

1. 디퓨전 오토인코더 (DAE):

   • 인코더 (E): 입력 3D MRI 이미지를 잠재 표현($z$)으로 매핑한다. 일반적인 오토인코더와 달리, DAE 인코더는 고수준의 의미론적 구조를 포착하는 데 집중하며, 디노이징 디코더는 고주파 세부 사항을 처리한다.
   • 디코더 (D): 인코더로부터 얻은 잠재 표현 $z$에 조건화되어, 가우시안 노이즈를 타겟 이미지로 반복적으로 변환하는 디노이징 디퓨전 모델(DDIM 기반)이다. 이 설계는 정보의 분리를 용이하게 하는 명시적이고 압축적인 잠재 표현을 가능하게 한다.

2. 잠재 이동 추정 모듈 (A):

   • 이 모듈은 진행 양상을 시뮬레이션하기 위해 잠재 공간 내에서 제어된 움직임을 유도한다. 이 모듈은 진행 속성, 즉 연령 차이($v_a$)와 인지 상태($v_d$, 예: CN, MCI, AD)를 입력으로 받는다.
   • 이 모듈은 베이스라인 이미지로부터 유도된 잠재 벡터의 특정 부분 집합($z_b$)에 더해질 이동 벡터($z'$)를 생성한다.
   • 결정적으로, 이 이동은 $d$-차원 잠재 공간의 처음 $m$개 차원에만 적용된다($z'_f = z_b + [z'; 0]$). 이 메커니즘은 진행 관련 요인(첫 $m$ 차원에 인코딩됨)과 피험자 정체성 특징(나머지 $d-m$ 차원에 인코딩됨)을 분리하도록 설계되었다.

3. 일관성 모듈 (R):

   • 쌍생 그라운드 트루스(ground truth) 없이도 생성된 진행이 의미 있고 해부학적으로 제약되도록 보장하기 위해, 일관성 모듈이 회귀 모델로서 작용한다.
   • 이 모듈은 베이스라인 이미지($x_b$), 생성된 후속 이미지($\hat{x}_f$), 그리고 잔차($x_b - \hat{x}_f$)를 입력으로 받아 진행 속성($v'_d, v'_a$)을 예측한다.
   • 이 모듈은 잔차를 마스킹함으로써 질병 특이적 영역(해마, 편도체, 측뇌실)에 집중한다. 이는 잠재 이동 모듈이 진행 속성과 상관관계를 갖고 관련 해부학적 영역에 공간적으로 국한된 변화를 생성하도록 암묵적으로 가이드한다.

훈련 및 추론

• 훈련: 모델은 두 가지 모드로 훈련된다. 초기에 이미지 특징을 학습하기 위해 오토인코딩 모드($\lambda_1=0$)로 작동한다. 이후 진행 생성 모드($\lambda_1=1$)로 전환되어, 이미지 재구성을 위한 평균 제곱 오차(MSE)와 일관성 모듈의 속성 회귀를 위한 교차 엔트로피(CE) 손실을 포함한 결합 손실 함수를 최적화한다.
• 추론: 베이스라인 이미지와 타겟 속성(연령 차이, 인지 상태)이 주어지면, 인코더가 $z_b$를 생성한다. 잠재 이동 모듈은 $z_b$에 더해질 $z'$를 계산한다. 수정된 잠재 벡터는 디퓨전 디코더가 $\hat{x}_f$를 반복적으로 생성하도록 가이드한다.

주요 기여

1. AD-DAE 프레임워크: 쌍생 종단적 데이터 없이도 진행 관련 요인을 피험자 정체성으로부터 효과적으로 분리하며 제어 가능한 잠재 이동을 가능하게 하는 디퓨전 오토인코더 프레임워크를 도입하였다.
2. 비쌍생 진행 모델링: 진행 속성과 상관관계를 갖고 알츠하이머 관련 해부학적 영역에 공간적으로 국한된 변화를 강제함으로써, 피험자별 종단적 쌍 없이도 질병 진행을 모델링하는 메커니즘을 제시하였다.
3. 포괄적 검증: ADNI 및 OASIS 데이터셋의 피험자 1,016명을 대상으로 검증을 수행하였으며, 이미지 품질 지표, 체적 진행 분석, 다운스트림 질병 분류 작업을 통해 효과성을 입증하였다.

결과

• 이미지 품질: 테스트 세트에서 AD-DAE는 베이스라인 모델들(CAAE, IPGAN, UVCGAN, BrLP, DE-CVAE, SITGAN)보다 우수한 정량적 성능을 달성하였다. 가장 높은 PSNR(CN의 경우 30.10 dB, MCI/AD의 경우 29.43 dB)과 SSIM(0.94)을 기록했으며, 가장 낮은 MSE를 유지하였다.
• 체적 분석: 해마, 편도체, 측뇌실의 체적 변화에 대한 영역별 평균 절대 오차(MAE)에서 AD-DAE는 모든 베이스라인(예: SITGAN의 해마 MAE 0.1161 대비 AD-DAE 0.0282)을 크게 앞질렀다.
• 변형 분석: 야코비안 행렬식(Jacobian determinant) 분석을 통해 AD-DAE가 다른 방법들에 비해 비선형적 해부학적 변형(위축 및 뇌실 확장)을 더 정확하게 포착함을 확인하였으며, 그 오차가 그라운드 트루스에 가장 근접하였다.
• 다운스트림 유용성: 질병 분류 작업에서 AD-DAE로 생성된 데이터를 실제 훈련 데이터에 증강했을 때, 실제 데이터 100%를 사용하는 것과 유사한 수준으로 분류 정확도를 효과적으로 향상시켰다.
• 잠재 공간 조직화: UMAP 투영 및 잠재 스왑 실험을 통해 첫 $m$개의 잠재 차원이 연령과 인지 상태를 성공적으로 인코딩함을 보여주었으며, 이를 통해 피험자의 정체성을 보존하면서도 진행 변화(예: 뇌실 수축/확장)를 생성할 수 있음을 입증하였다.

의의 및 주장
본 논문은 AD-DAE가 이미지 공간의 디퓨전과 오토인코딩 형식의 제어 가능성을 결합함으로써 종단적 질병 진행 모델링의 중요한 진전을 이루었다고 주장한다. 저자들은 이 접근 방식이 피험자 정체성으로부터 진행 관련 요인을 분리하도록 잠재 공간을 성공적으로 조직하여, 쌍생 종단적 가이드 없이도 그럴듯한 후속 이미지를 생성할 수 있게 한다고 강조한다. 결과는 이미지 공간에서의 비선형적 해부학적 변형이 DAE의 잠재 공간 내에서의 선형 변환에 의해 효과적으로 근사될 수 있음을 시사한다. 이 프레임워크는 피험자의 정체성을 유지하면서 질병 특이적인 구조적 변화를 정확하게 반영하는 데이터를 생성할 수 있는 견고한 도구로서, 쌍생 종단적 데이터가 없는 상황에서도 다운스트림 임상 작업 및 연구를 지원할 수 있는 능력을 갖추고 있다.