Atlases Are Already Inside: Recovering Population Templates from Pretrained Diffusion Models

기술 요약: 아틀라스는 이미 내부에 존재한다

문제 정의
인구 집단 수준의 해부학적 아틀라스를 구축하는 전통적인 방식은 명시적인 이미지 등록(registration)과 변형 모델링(deformation modeling)에 의존한다. 이 과정은 많은 개인을 공통된 좌표계로 정렬하기 위해 메트릭, 정규화 항(regularizer), 유도 편향(inductive bias)에 대한 모달리티별 선택이 필요한 복잡한 최적화 문제를 해결해야 한다. 이 방식은 성공적이긴 하지만, 돌이킬 수 없는 비용을 수반한다. 즉, 아틀라스가 데이터 분포의 내재적 속성이 아니라 외부의 최적화 대상이라는 점이다. 저자들은 근본적인 질문을 던진다. 만약 확산 모델(diffusion model)이 전체 해부학적 인구 집단의 분포를 학습했다면, 그 집단의 정준 참조(canonical reference, 아틀라스)는 명시적인 등록이나 아틀라스 전용 목적 함수 없이도 모델의 학습된 역학(dynamics) 내에 이미 암시적으로 존재하는가?

방법론
본 논문은 사전 학습된 확산 모델로부터 결정론적 추론(deterministic inference)을 통해 인구 템플릿을 직접 복구하는 방법을 제안하며, 이는 등록이나 미세 조정(fine-tuning)의 필요성을 배제한다.

• 메커니즘: 이 방법은 확산 모델의 결정론적 역과정(probability-flow ODE)을 활용한다. 모델은 오직 이미지 합성(노이즈로부터 깨끗한 이미지를 예측)만을 위해 학습되었다. 저자들은 결정론적 역과정이 무작위 가우시안 노이즈로부터 인구 통계적으로 전형적인 이미지로 가는 지도 역할을 한다고 가정한다.
• 수렴 논리: 높은 노이즈 수준에서 모델의 조건부 평균은 인구 평균(흐릿한 평균값)에 근사한다. 노이즈가 감소함에 따라, 프로세스는 현재 상태를 기반으로 이 추정치를 반복적으로 정교화한다. 만약 학습 분포가 충분히 정렬되어 있고 인구 집단으로서 일관성(population-coherent)을 갖춘다면, 독립적인 노이즈 초기화로부터 시작된 궤적들은 동일하고 선명하며 해부학적으로 전형적인 이미지로 수렴한다.
• 추론 절차:
  1. 무작위 가우시안 노이즈($x_T \sim \mathcal{N}(0, I)$)로 초기화한다.
  2. 확률적 노이즈를 추가하지 않고 결정론적 역과정(반복적 평균 업데이트)을 실행한다.
  3. 특정 타임스텝 $t^*$ (예: 실험에서의 $t^*=99$)에서 멈춘다. 이때 궤적은 안정적이고 선명한 템플릿으로 수렴했지만, 고주파의 피험자 특이적 세부 사항이 다시 도입되기 전 단계여야 한다.
  4. 결과물인 이미지가 복구된 아틀라스이다.
• 조건부 생성: 이 방법은 조건부 모델(예: 연령 조건부 또는 대조 조건부 모델)로 확장된다. 역과정 동안 조건 변수(예: 연령 = 30)를 고정함으로써, 특정 하위 그룹에 대한 정준 참조를 복구할 수 있으며, 이는 요구에 따라 다양한 아틀라스 패밀리를 생성하는 효과를 갖는다.

주요 기여

1. 아틀라스 구축의 재정의: 저자들은 해부학적 아틀라스가 단순히 외부에서 최적화된 템플릿이 아니라, 학습된 확산 모델의 암시적 정준 구조임을 입증한다. 아틀라스 구축은 생성적 추론의 부산물로 재정의된다.
2. 등록이 필요 없는 복구: 인구 템플릿은 이미지 합성을 위해서만 학습된 모델로부터 순수하게 추론을 통해 복구될 수 있으며, 아틀라스 전용 목적 함수나 미세 조정, 등록 최적화를 필요로 하지 않는다.
3. 온디맨드 조건부 아틀라스: 하나의 조건부 모델은 조건 변수를 변화시킴으로써 연속적인 아틀라스 패밀리(예: 인간의 생애 주기 전반에 걸친 아틀라스)를 생성할 수 있으며, 이는 알려진 생물학적 경향(예: 노화에 따른 뇌척수액 확장)을 재현한다.
4. 일반화 가능성: 이 절차는 학습 분포가 충분히 일관적이라면 다양한 의료 영상 모달리티(T1/T2 MRI, CT, X-ray)와 차원(2D 및 3D)에 걸쳐 작동함이 입증되었다.

결과

• 질적 성능: T1 가중치 뇌 MRI에 대해, 복구된 참조 모델은 기존 아틀라스(MNI152, ANTs-IXI, ANTs-OASIS, VoxelMorph, MultiMorph)와 비교했을 때 일관된 전역적 조직과 인식 가능한 해부학적 패턴을 보여준다. "유령 해부학(ghost anatomy)"이나 거친 불연속성이 나타나지 않는다.
• 정량적 성능 (등록 대상으로서):
  • 기하학적 앵커: 복구된 참조 모델은 기존 아틀라스(IXI, OASIS, MNI152)보다 인구 집단을 더 긴밀하게 고정하며, 더 완만한 변형(낮은 굽힘 에너지)과 함께 더 낮은 사후 등록 복셀 분산을 달착한다. 이는 분산 대 변형 크기의 파레토 프런티어(Pareto frontier) 상에 위치한다.
  • 등록 메트릭: 세 가지 라벨 코호트(LPBA, IBSR, Mindboggle)에서 복구된 참조 모델은 기존 베이스라인과 경쟁하거나 이를 능가하는 교차 상관(Cross-Correlation, CC) 점수를 달성했다. 예를 들어, LP_BA 코호트에서 'null-age'로 복구된 참조는 테스트된 7개의 아틀라스 중 가장 높은 CC(0.9797)를 기록했다.
  • Dice 점수: 복구된 참조는 목적 지향적 학습 기반 아틀라스(예: LPBA에서 MultiMorph보다 약 4~5점 뒤처짐)에 비해 약간 낮은 Dice 중첩 점수를 보인다. 저자들은 이를 모델이 등록이나 라벨 보존을 위해 최적화되지 않았기 때문이라고 설명한다.
• 연령 조건부 검증: 연령 조건부 모델은 건강한 노화와 관련된 뇌척수액(CSF) 부피의 단조적이고 가속화된 확장을 성공적으로 재현했다. 또한, 연령이 일치하는 아틀라스를 등록 대상으로 사용했을 때, 연령이 맞지 않는(예: 80세) 아틀라스를 사용했을 때보다 훨씬 높은 교차 상관을 보였다.
• 실패 사례: 모델링 스케일에서 학습 분포의 인구 집단 일관성이 부족할 경우 선명한 아틀라스를 생성하는 데 실패한다. 이는 시드(seed)에 따른 발산(조건부 없이 다중 모드 분포가 나타나는 경우) 또는 시각적으로 흐릿한 이미지로의 수렴(해결할 수 없는 높은 클래스 내 변동성)으로 나타난다.

의의 및 주장
본 논문은 계산 해부학의 개념적 전환을 주장한다. 확산 모델은 단순한 생성기가 아니라 인구 구조의 학습된 표현체이다. 전통적인 방식이 최적화하는 정준 참조는 이미 확산 사전 확률(diffusion prior) 안에 암시적으로 존재한다.

저자들은 자신들이 새로운 등록 알고리즘을 제안하거나 기존 파이프라인을 대체하려는 것이 아님을 명시적으로 밝힌다. 대신, 추가적인 구축 비용 없이 인구 집단에 대해 학습된 생성기로부터 얻을 수 있는, 해부학적으로 일관되고 등록 대상으로서 경쟁력을 갖춘 "공짜 점심(free-lunch)" 형태의 인구 템플릿을 제공한다. 이 연구의 의의는 아틀라스 구축을 최적화 문제에서, 생성적 분포의 속성을 추론 시점에 접근 가능한 형태로 재정의한 데 있다. 이 방법은 특히 데이터 정렬이 불완전한 경우를 대비한 전문적인 아틀라스 구축 도구의 완전한 대체재라기보다는, 하나의 능력 증명(capability demonstration)으로 제시된다.