BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation

Three-dimensional (3D) brain MRI is central to clinical neurology and neuro-oncology, where generative models could augment under-represented cohorts, simulate disease trajectories, and support privacy-preserving data sharing. Latent diffusion has been the go-to solution for modeling imaging data, but it places two competing demands on the tokenizer: encoder embeddings must retain the clinical information that downstream tasks act on, and the decoder must reconstruct anatomically faithful volumes. Existing reconstruction-driven tokenizers achieve the second at the expense of the first. To address this, we introduce a fully volumetric masked-autoencoder (MAE) based tokenizer for 3D brain MRI latent diffusion, decoupling encoder and decoder: a frozen 3D MAE encoder produces clinically informative embeddings, while a dedicated CNN decoder reconstructs voxels from a linear projection of those embeddings. We pretrain the encoder on 35,309 volumes from 18 public cohorts spanning four modalities, ten disease categories, and 200+ acquisition sites, and demonstrate its dual utility in two settings. First, on a 23-task linear-probing benchmark, the encoder outperforms or matches SOTA models (i.e., BrainIAC, BrainSegFounder, and MedicalNet) on 21 of 23 tasks. Second, a conditional diffusion transformer (DiT) trained on these clinically informative embeddings supports both conditional generation across six variables and patient-specific longitudinal forecasting. Together these results establish a single 3D brain-MRI embedding space capable of both downstream clinical tasks and controllable generation.

한 줄 요약

BrainG3N은 3D 뇌 MRI 생성과 임상 태스크를 동시에 지원하는 이중 목적의 토크나이저를 제안한다.

핵심 기여도

35,309개의 3D MRI 볼륨으로 사전 학습한 3D MAE 인코더는 23개 태스크 중 21개에서 BrainIAC, BrainSegFounder, MedicalNet을 앞선다.
AUC 0.937 (IDH1 돌연변이 예측), 0.921 (종양 등급 분류), 0.967 (성별 예측) 등 임상적으로 중요한 지표에서 높은 성능을 보인다.
조건부 생성과 환자별 장기 예측을 지원하는 DiT 모델을 제시하며, 생성된 임베딩은 CNN 디코더를 통해 고해상도 볼륨으로 복원된다.
생성된 샘플은 실제 데이터 분류기에서 Pearson r = 0.93 (나이 조건), 0.72 (장기 나이 진행) 등 높은 상관성을 보인다.

핵심 아이디어

기존 3D 뇌 MRI 생성 모델은 인코더가 복원 정확도를 우선시함으로써 임상 정보를 잃는 문제가 있었다. BrainG3N은 이 문제를 해결하기 위해 인코더와 디코더를 분리한 이중 목적 토크나이저를 제안한다. 구체적으로, 3D MAE 인코더는 사전 학습 후 고정(frozen)되며, 임상적으로 유의미한 임베딩을 생성하고, 별도의 CNN 디코더는 볼륨 복원을 담당한다. 이는 인코더가 복원 정확도에 영향을 받지 않고 임상 정보를 유지할 수 있게 한다.

또한, 이 임베딩 공간은 조건부 생성 및 장기 예측을 위한 DiT 모델 학습에 활용된다. 이는 생성된 데이터가 실제 임상 분류기에서 조건을 정확히 반영함으로써, 생성 모델이 임상적으로 의미 있는 정보를 포함하고 있음을 입증한다.

기술적 접근법

**모델 구조**: 3D MAE 인코더 + CNN 디코더 + DiT 생성자
**인코더**: 35,309개의 3D MRI 볼륨으로 사전 학습한 3D MAE, 고정(frozen) 상태 유지
**디코더**: 인코더 임베딩의 선형 투영을 기반으로 볼륨 복원
**DiT**: 조건부 생성 및 장기 예측을 위한 조건부 확률 확산 변형자
**데이터**: 18개 공개 코호트, 4개 모달리티(T1, T2, FLAIR, T1c), 10개 질병 범주, 200개 이상의 측정 장소
**학습**: 인코더는 선형 프로빙, DiT는 조건부 생성 및 장기 예측을 위한 훈련

주요 결과

**선형 프로빙**: 23개 태스크 중 21개에서 SOTA 모델(BrainIAC, BrainSegFounder, MedicalNet)을 앞선다.
**IDH1 돌연변이 예측**: AUC 0.937
**종양 등급 분류**: AUC 0.921
**뇌 나이 회귀**: 평균 절대 오차 4.43년
**성별 예측**: AUC 0.967
**장기 예측**: Δt = 10년 조건에서 Pearson r = 0.72, 실제 나이 진행과 유사한 변화를 재현
**복원 정확도**: 볼륨 복원에서 정확도 높은 CNN 디코더 사용

의의 및 한계

BrainG3N은 3D 뇌 MRI 분야에서 임상 태스크와 생성 모델을 동시에 처리할 수 있는 단일 임베딩 공간을 제시하며, 이는 데이터 증강, 질병 시뮬레이션, 프라이버시 보호 데이터 공유 등 여러 임상 응용에 기여할 수 있다. 특히, 고정된 인코더를 통해 사전 학습된 임베딩이 다양한 조건에서 재사용 가능하다는 점이 학술적 가치를 높인다.

그러나, 일부 태스크(GDS 우울, UPDRS-III)에서는 모든 인코더가 낮은 성능을 보이며, 이는 데이터의 복잡성이나 특성에 따라 한계가 있을 수 있음을 시사한다. 또한, 생성된 볼륨의 변화량이 실제 나이 진행에 비해 약 27%만 재현되며, 이는 확률적 브리지 정규화의 영향으로 설명된다.

실용적 활용

BrainG3N은 뇌 종양 진단, 신경퇴행성 질환 모니터링, 환자별 디지털 트윈 생성, 병원 간 프라이버시 보호 데이터 공유 등 다양한 임상 및 연구 상황에 적용 가능하다. 특히, 조건부 생성과 장기 예측 기능은 임상 시뮬레이션 및 치료 계획 수립에 유용하게 활용될 수 있다.