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

2026-06-17

AI Summary

Existing tokenizers for 3D brain MRI generation struggle to simultaneously preserve the semantic information required for clinical tasks and the anatomical fidelity essential for high-quality synthesis. This work proposes a dual-purpose tokenizer based on a 3D Masked Autoencoder (MAE) with a decoupled architecture: a frozen 3D MAE encoder extracts embeddings rich in clinical semantics, while a dedicated CNN decoder enables high-fidelity voxel reconstruction from linearly projected embeddings. Integrated with a conditional diffusion Transformer (DiT), the framework supports controllable generation. For the first time, this approach unifies multi-task clinical analysis and conditional synthesis within a single embedding space, achieving state-of-the-art or comparable performance on 21 out of 23 linear probing tasks, enabling six-variable conditional generation and patient-specific longitudinal prediction, thereby overcoming the longstanding trade-off between semantic retention and structural reconstruction in conventional tokenizers.

Abstract

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.

Problem

• 3D brain MRI
• tokenizer
• clinical information
• anatomical fidelity
• latent diffusion

Innovation

• masked autoencoder
• latent diffusion
• 3D brain MRI
• dual-purpose tokenizer
• controllable generation