BrainG3N: 제어 가능한 3D Brain MRI 생성을 위한 이중 목적 토크나이저

3D brain MRI의 생성-재구성 파이프라인에서 인코더-디코더 토크나이저는 임상 정보를 잃지 않으면서도 고해상도 재구성을 달성해야 한다. 본 논문은 frozen MAE encoder를 통해 임상 정보를 보존하는 이중 목적 토크나이저를 제시하고, 이 임베딩 공간을 통해 6가지 조건으로의 제어 가능한 생성과 환자별 종단 예측을 가능하게 한다.

출발점: 3D brain MRI의 latent diffusion 파이프라인에서 encoder embeddings가 임상 정보를 충분히 담아야 하므로, voxel 재구성에 초점을 맞춘 재구성 기반 토크나이저는 임상 내용을 손실시키는 경향이 있다. 이러한 한계를 극복하기 위해, 3D MAE 기반 토크나이저를 도입하고 encoder는 고정된 상태로 남겨 두되, 1200 토큰을 32채널로 축소하는 선형 프로젝션(P)과 CNN 디코더를 연결한다. 이 구성을 통해 임상 정보가 유지된 임베딩 공간을 형성한다. 두 번째 축은 이 임베딩 공간을 입력으로 삼아 DiT( diffusion transformer )를 학습시키는 것이다. DiT는 12-블록 구조에 adaLN-Zero conditioning으로 여섯 가지 조건을 주입하고, CFG를 통해 조건부 샘플링을 제어한다. 이 임베딩 공간은 질병 상태, 생물학적 성별, 촬영 모달리티, 촬영 사이트, 연령, IDH1 돌연변이 여부 등 다중 정보가 섞인 표현으로 작동한다. 결과적으로, 이 단일 임베딩 공간은 Downstream 임상 예측 태스크와 controllable generation을 모두 지원하며, cross-sectional 프로빙, 조건부 생성, longitudinal forecasting에서 서로 보강되는 성능을 달성한다.

Phase 1: 전체 접근의 핵심은 3D MAE encoder의 사전학습이다. 입력 x는 70%가 마스킹된 상태에서, encoder는 남은 360개의 패치를 처리하고, separate decoder가 840개의 마스킹 패치를 복원한다. [입력] x → [연산] Enc(x) [frozen], z′ = zP [선형 투영], x̂ = φ(z′) [복원], L(voxel) = ℓ1(x̂, x) → [결과] 임상적으로 풍부한 임베딩(z) 생성. 의미: encoder가 전역적 해부학적 맥락을 포착하도록 유도한다. Phase 2: Phase-2 tokenizer에서 encoder를 고정하고, 선형 projection P로 z를 32채널로 축소한다. CNN decoder가 z′에서 voxel을 재구성한다. 이때의 손실은 voxel ℓ1이다. [입력] x → Enc(x) → z′ = zP → x̂ = φ(z′) [연산] voxel loss ℓ1 → [결과] 임상 정보를 보존하는 임베딩 공간 생성 → [의미] Diffusion 모델의 입력 차원을 감소시키면서도 중요한 임상 정보를 유지한다. Phase-2의 임베딩은 DiT의 입력으로 사용되며, DiT는 flow-matching objective를 따른다. x1은 real latent, ε은 표준 정규분포, c는 여섯 조건의 임베딩 벡터이며, vθ를 이용해 x_t를 예측한다. [입력] x_t, t, c → [연산] vθ(·) ≈ (ε − x1) → [결과] x1에 조건을 부여한 샘플 생성 → [의미] 조건부 생성의 제어 가능성 확보. Phase-3: CFG(dropout)로 conditioning 신호를 조절하고, sampling에서 CFG-스케일 s를 통해 조건과 무조건 간의 샘플링 차이를 제어한다. Longitudinal Variant: baseline-followup latent를 이용한 브리지 샘플링으로 Δt의 예측 차이를 증가시키는 샘플링 방식을 도입한다. [입력] baseline x_b, follow-up x_f, Δt → [연산] interpolant + envelope로 x_t 생성 → [결과] longitudinal 샘플링 → [의미] 환자별 시간 경과 예측 가능성 확보.

3D 프로빙 및 벤치마크: 35,309개 코호트 데이터에서 frozen encoder 임베딩의 임상 정보 보존 능력은 23개 태스크 중 21개 태스크에서 SOTA에 필적하거나 우수하였다(예: IDH1 mutation AUC 0.937 ± 0.020, Tumor grade AUC 0.921 ± 0.019, Brain age MAE 4.43 ± 0.07년, Sex AUC 0.967 ± 0.004). DiT 기반 컨디셔널 생성: CFG를 이용한 조건부 생성에서 Real-data probe가 요청된 속성을 회복하며, per-arm 3D-FID 비교에서 gen-vs-real: 34.4, gen-vs-recons: 107.3, recons-vs-real: 77.4를 보였다. Modality는 샘플링 중-dropping 되지 않으며, age의 선형 스윕은 r=0.93으로 측정되었다. Longitudinal forecasting: Δt = {0,1,2,5,10}년에서 baseline을 기준으로 샘플링하고, ADNI 데이터에 대해 brain-age probe로 평가했다. slope는 0.268로 추정되며 Pearson r은 0.716으로 측정되었다. 예측 차이는ventricles 및 sulci의 변화에 주로 집중되었다. 27%의 실제 aging magnitude를 회복하는 수준이며, 이는 확률적 브리지 규제에 의해 조건 평균으로 수렴하는 특성 때문으로 해석된다.

[아키텍처] Phase 1 MAE encoder( ViT-12 layers, hidden 1152, patch 16^3, 1200 tokens) + Phase 2 projection(1152×32) + CNN decoder. [학습] Phase 1: 70% masking, patch MSE loss, 360 visible tokens. Phase 2: voxel ℓ1 loss, 32 채널 고정 임베딩. [DiT] 12-block DiT-L, 1152 hidden, 18 heads, 64 samples, CFG dropout p=0.1, conditioning 6 dims. [Longitudinal] baseline-followup latent 간의 Brownian-bridge interpolant, σ=0.5, 100-step Euler.

본 연구는 뇌 MRI에 한정되며 다른 3D 모달리티(CT, 전신 MRI 등)로의 일반화는 검증되지 않았다. longitudinal forecasting의 샘플링은 실제 aging의 27%만 반영하며, 브리지 샘플링의 한계상 향상 여지가 있다. 설계상의 차이로 인해 배포 코드는 논문 발표 시점 기준으로 공개되지 않았으며, 데이터 수집 분포 간 차이가 성능에 영향을 미칠 수 있다. 장치 및 취득 조건에 따른 임베딩 편향 가능성도 존재한다.