BrainG3N: Двуцелевой токенизатор для управляемой генерации 3D МРТ головного мозга

Аннотация

Трёхмерная (3D) МРТ головного мозга является центральным инструментом в клинической неврологии и нейроонкологии, где генеративные модели могут дополнять недостаточно представленные когорты, моделировать траектории заболеваний и обеспечивать конфиденциальный обмен данными. Латентная диффузия стала стандартным решением для моделирования данных визуализации, но она предъявляет два конкурирующих требования к токенизатору: вложения энкодера должны сохранять клиническую информацию, на которую опираются последующие задачи, а декодер должен реконструировать анатомически достоверные объёмы. Существующие токенизаторы, основанные на реконструкции, достигают второго за счёт первого. Для решения этой проблемы мы предлагаем полностью объёмный токенизатор на основе маскированного автоэнкодера (MAE) для латентной диффузии 3D МРТ головного мозга, разделяющий энкодер и декодер: замороженный 3D MAE-энкодер формирует клинически информативные вложения, а специализированный CNN-декодер реконструирует воксели на основе линейной проекции этих вложений. Мы предварительно обучаем энкодер на 35 309 объёмах из 18 общедоступных когорт, охватывающих четыре модальности, десять категорий заболеваний и более 200 центров сбора данных, и демонстрируем его двойное применение в двух сценариях. Во-первых, в тесте линейного зондирования по 23 задачам энкодер превосходит или достигает уровня SOTA-моделей (BrainIAC, BrainSegFounder и MedicalNet) по 21 из 23 задач. Во-вторых, условный диффузионный трансформер (DiT), обученный на этих клинически информативных вложениях, поддерживает как условную генерацию по шести переменным, так и индивидуальное продольное прогнозирование. В совокупности эти результаты устанавливают единое пространство вложений 3D МРТ головного мозга, пригодное как для последующих клинических задач, так и для управляемой генерации.

English

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.