BrainG3N: een tweeledige tokenizer voor beheersbare 3D-hersen-MRI-generatie

Samenvatting

Driedimensionale (3D) hersens-MRI staat centraal in de klinische neurologie en neuro-oncologie, waar generative modellen ondervertegenwoordigde cohorten kunnen aanvullen, ziektetrajecten kunnen simuleren en privacybeschermende gegevensuitwisseling kunnen ondersteunen. Latente diffusie is de standaardoplossing geweest voor het modelleren van beeldgegevens, maar het stelt twee concurrerende eisen aan de tokenizer: encoder-embeddings moeten de klinische informatie behouden waar downstream-taken op werken, en de decoder moet anatomisch getrouwe volumes reconstructen. Bestaande reconstructiegedreven tokenizers bereiken het tweede ten koste van het eerste. Om dit aan te pakken, introduceren wij een volledig volumetrische gemaskeerde autoencoder (MAE) gebaseerde tokenizer voor 3D-hersens-MRI latente diffusie, waarbij encoder en decoder worden ontkoppeld: een bevroren 3D MAE-encoder produceert klinisch informatieve embeddings, terwijl een speciale CNN-decoder voxels reconstructeert uit een lineaire projectie van die embeddings. We trainen de encoder vooraf op 35.309 volumes uit 18 openbare cohorten die vier modaliteiten, tien ziektes categorieën en meer dan 200 acquisitielocaties beslaan, en tonen het dubbele nut ervan aan in twee settings. Ten eerste, op een 23-taak linear-probing benchmark, presteert de encoder beter dan of gelijk aan SOTA-modellen (d.w.z. BrainIAC, BrainSegFounder en MedicalNet) op 21 van de 23 taken. Ten tweede, een conditionele diffusietransformer (DiT) getraind op deze klinisch informatieve embeddings ondersteunt zowel conditionele generatie over zes variabelen als patiëntspecifieke longitudinale voorspelling. Samen vestigen deze resultaten een enkele 3D-hersen-MRI-embeddingruimte die zowel downstream klinische taken als beheersbare generatie mogelijk maakt.

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.