Search papers, labs, and topics across Lattice.
This paper introduces BrainG3N, a dual-purpose tokenizer designed for 3D brain MRI generation that effectively balances the need for clinically informative embeddings and anatomically accurate reconstructions. By employing a frozen 3D masked-autoencoder (MAE) for encoding and a dedicated CNN for decoding, the model successfully retains critical clinical information while enabling high-quality voxel reconstruction. The results show that the encoder surpasses or matches state-of-the-art models on 21 out of 23 tasks in a linear-probing benchmark, while also supporting conditional generation and patient-specific forecasting through a conditional diffusion transformer.
Achieving a dual-purpose tokenizer that excels in both clinical task performance and controllable 3D brain MRI generation could revolutionize how we approach medical imaging.
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.
Stanford HAI