기술 요약: ProgFormer

문제 정의

미래의 구조적 뇌 MRI(sMRI)를 예측하는 것은 알츠하이머병과 같은 신경 퇴행성 질환을 분석하고 예측하는 데 있어 매우 중요한 과제입니다. 그러나 이는 독특한 도전 과제를 안겨줍니다. 종단적 변화(longitudinal changes)는 대개 미세하며 특정 해부학적 영역(예: 해마, 측뇌실)에 국한되어 나타나는 반면, 피험자의 뇌 구조 대부분은 시간이 지나도 안정적인 상태를 유지하기 때문입니다.

기존 방법들은 일반적으로 다음 두 가지 범주로 나뉘며, 각각 한계점을 가집니다:

1. 잠재 공간 기반 방법(Latent-space-based methods): 스캔 데이터를 저차원 표현으로 인코딩한 후 예측합니다. 계산 효율성은 높지만, 별도로 학습된 오토인코더에 의존하므로 압축 및 재구성 과정에서 정보 손실이 발생할 수 있으며, 이는 미세한 구조적 세부 사항의 충실도를 저하시킬 수 있습니다.
2. 직접 복셀 공간 방법(Direct voxel-space methods): 잠재 재구성을 피하기 위해 복셀 공간에서 직접 출력을 예측합니다. 그러나 이들은 안정적인 전역 뇌 구조와 미세한 진행 관련 변화를 모두 모델링하기 위해 단일화된 예측 경로를 사용하는 경향이 있습니다. 결과적으로, 지배적인 안정적 콘텐츠가 미세한 국소 변화를 압도하게 되어, 정밀한 질병 진행을 해결하는 데 어려움을 겪습니다.

방법론: ProgFormer

저자들은 별도의 이미지 오토인코더에 의존하지 않고 이러한 문제를 해결하기 위해 설계된 계층적 복셀 공간 확산 트랜스포머(hierarchical voxel-space Diffusion Transformer)인 ProgFormer를 제안합니다. 이 모델은 조건부 흐름 매칭(Conditional Flow Matching) 프레임워크 내에서 작동하며, 타겟 미래 MRI 분포로 가우시안 노이즈를 운송하기 위한 속도장(velocity field)을 학습합니다.

핵심 아키텍처

ProgFormer는 예측 과정을 두 개의 구별되지만 결합된 경로로 구조화합니다:

1. 조대 경로 (Coarse Pathway, 패치 수준 추론):

   • 입력: 관찰된 스캔과 현재의 흐름 상태를 겹치지 않는 3D 패치 토큰($p \times p \times p$)으로 처리합니다.
   • 메커즘: 각 스캔 내에서의 **공간 다두 자가 주의 집중(Spatial Multi-Head Self-Attention, MHSA)**과 관찰된 스캔들 간의 **인과적 시간적 MHSA(Causal Temporal MHSA)**를 활용합니다. 이를 통해 모델은 전역적인 뇌 구조와 종단적 문맥을 포착할 수 있습니다.
   • 출력: 일차적인 볼륨 속도($\nu_{coarse}$)와 전체 뇌 구조 및 시간적 문맥을 인코딩하는 조대 표현($C$)을 생성합니다.

2. 미세 경로 (Fine Pathway, 복셀 수준 정밀화):

   • 입력: 각 3D 패치 내의 개별 복셀 토큰을 대상으로 작동합니다.
   • 메커니즘: 패치 내에서 복셀 MHSA를 적용하며, 결정적으로 **조대-미세 다두 교차 주의 집중(Coarse-to-Fine Multi-Head Cross-Attention, MHCA)**을 적용합니다. 이 교차 주의 집중에서 복셀 토근은 쿼리(query) 역할을 하고, 해당 패치의 조대 표현($C$)이 키(key)와 값(value) 역할을 합니다.
   • 기능: 이 메커니즘은 조대 경로의 출력을 "시공간적 그라운딩(spatio-temporal grounding)"으로 사용하여 국소적인 구조적 세부 사항과 강도 변화의 정밀화를 유도합니다.
   • 출력: 정밀화된 속도($\nu_{fine}$)를 생성합니다.

통합 및 추론

• 속도 추정: 최종 예측된 속도장은 조대 속도와 미세 속도의 합입니다: $\nu_\theta = \nu_{coarse} + \nu_{fine}$.
• 생성: 추론 과정에서 모델은 가우시안 노이즈에서 시작하여 추정된 속도장을 일련의 오일러 단계(Euler steps)에 걸쳐 적분하여 예측된 미래 MRI를 생성합니다.
• 변형 모델: 저자들은 두 가지 구성을 평가합니다:
  • 순차적 방식 (Sequential, ProgFormer-S): 모든 조대 블록이 완료된 후에 모든 미세 블록이 시작됩니다.
  • 인터리브 방식 (Interleaved, ProgFormer-I): 조대 블록과 미세 블록이 단계별로 교차 배치되어, 미세 블록이 점진적으로 정밀해지는 조대 특징을 받을 수 있도록 합니다.

조건화 및 학습

• 조건화: 모델은 AdaLN-Zero(제로 초기화가 적용된 적응형 레이어 정규화)를 통해 흐름 시간, 목표 연령, 공변량(진단, 성별)에 따라 조건화되며, 이는 어텐션 및 MLP 모듈의 스케일, 시프트, 게이트 파라미터를 조절합니다.
• 손실 함수: 가중치가 적용된 흐름 매칭 목적 함수가 사용됩니다. SynthSeg에서 유도된 해부학적 라벨 맵은 진행 관련 영역(예: 해마, 뇌실)에는 더 높은 가중치를, 안정적인 조직/배경에는 더 낮은 가중치를 할당하여 모델이 질병 관련 변화를 우선시하도록 보장합니다.

주요 기여

1. 엔드 투 엔드 계층적 아키텍처: ProgFormer는 별도의 오토인코더 없이 전역 패치 수준의 시공간적 추론과 미세한 패치 내 복셀 모델링을 공동으로 수행하는 최초의 통합 아키텍처입니다.
2. 조대-미세 그라운딩: 이 설계는 전역적 안정성(조대 경로)의 모델링과 국소적 진행(미세 경로)의 모델링을 명시적으로 분리하며, 전자를 후자를 가이드하는 용도로 사용하여 미세한 변화가 안정적인 구조에 의해 가려지는 것을 방지합니다.
3. 최첨단 성능: 세 가지 벤치마크(ADNI, AIBL, OASIS)에 대한 광범한 실험을 통해 기존의 잠재 공간 및 직접 볼륨 방식에 비해 우수한 쌍별(pairwise) 및 궤적(trajectory) 예측 성능을 입증했습니다.

실험 결과

모델은 ADNI, AIBL, OASIS 데이터셋에 대해 PSNR, SSIM, 영역별 평균 절대 오차(R-MAE)를 사용하여 평가되었습니다.

• 쌍별 예측 (Pairwise Prediction): ProgFormer-S는 세 코호트 모두에서 가장 높은 PSNR을 달 기록했습니다 (예: ADNI에서 26.39로, 차순위 모델보다 0.20 높음). ProgFormer-I는 ADNI에서 가장 높은 지역 정확도(최저 R-MAE)와 가장 높은 SSIM을 달성했습니다.
• 궤적 예측 (Trajectory Prediction): 관찰된 이력(다수의 스캔)을 사용하는 것이 단일 입력 베이스라인보다 성능을 크게 향상시켰습니다. 이력을 사용한 ProgFormer-S가 모든 코호트에서 가장 높은 PSNR을 기록했습니다.
• 절제 연구 (Ablation Studies):
  • 조대 경로를 제거하면 예측이 붕괴되어 PSNR이 약 7로 떨어졌으며, 이는 전역 구조를 구축하는 데 있어 조대 경로의 필요성을 확인시켜 줍니다.
  • 미세 경로 또는 교차 주의 집중을 제거하면 이미지 품질과 지역적 정확도가 모두 크게 저하되었으며, 이는 조대 문맥에 의해 가이드되는 복셀 수준 정밀화의 중요성을 강조합니다.
  • 지역적 가중치는 결정적인 것으로 나타났습니다. 가중치를 제거하면 성능이 크게 하락하여, 진행 관련 영역을 강조하는 것이 전반적인 충실도를 높인다는 것을 증명했습니다.
• 시각적 분석: 정성적 결과에 따르면, ProgFormer는 베이스라인 모델들보다 뇌실 구성과 피질 윤곽을 더 일관되게 보존하며, 오류는 전역적으로 축적되지 않고 진행 관련 영역에 집중됩니다.

의의 및 주장

본 논문은 계층적 복셀 공간 모델링이 종단적 볼륨 예측에 효과적임을 입증한다고 주장합니다. 안정적인 전역 구조 모델링과 미세한 국소 변화 모델링을 분리함으로써, 이 방법은 해부학적으로 충실하고 진행에 민감한 예측을 달성합니다. 저자들은 이를 통해 잠재 공간 방식(재구성 손실 문제)과 통합된 직접 방식(안정적인 배경 대비 미세한 변화 해결 문제)으로부터 한 단계 진보한 성과라고 위치시킵니다. 본 연구는 직접적인 트랜스포머 기반 모델링이 시공간적 그라운딩을 통해 계층적으로 구조화될 때, 미래의 뇌 MRI를 생성하기 위한 실행 가능하고 우월한 접근 방식임을 확립합니다.