BrainG3N: Ein Dual-Purpose-Tokenizer für die kontrollierbare 3D-Gehirn-MRT-Generierung

Zusammenfassung

Dreidimensionale (3D) Gehirn-MRT ist zentral für die klinische Neurologie und Neuroonkologie, wo generative Modelle unterrepräsentierte Kohorten ergänzen, Krankheitsverläufe simulieren und datenschutzfreundlichen Datenaustausch unterstützen könnten. Die latente Diffusion war die bevorzugte Lösung zur Modellierung von Bilddaten, stellt jedoch zwei konkurrierende Anforderungen an den Tokenizer: Die Encoder-Embeddings müssen die klinischen Informationen bewahren, auf die nachgelagerte Aufgaben wirken, und der Decoder muss anatomisch getreue Volumen rekonstruieren. Bestehende rekonstruktionsgetriebene Tokenizer erreichen Letzteres auf Kosten des Ersteren. Um dies zu adressieren, führen wir einen vollständig volumetrischen Tokenizer auf Basis eines maskierten Autoencoders (MAE) für die latente Diffusion von 3D-Gehirn-MRT ein, der Encoder und Decoder entkoppelt: Ein eingefrorener 3D-MAE-Encoder erzeugt klinisch informative Embeddings, während ein dedizierter CNN-Decoder Voxel aus einer linearen Projektion dieser Embeddings rekonstruiert. Wir trainieren den Encoder auf 35.309 Volumen aus 18 öffentlichen Kohorten, die vier Modalitäten, zehn Krankheitskategorien und über 200 Akquisitionsstandorte umfassen, und demonstrieren seinen doppelten Nutzen in zwei Szenarien. Erstens übertrifft oder erreicht der Encoder in einem 23-Aufgaben-Benchmark mit linearer Auswertung den Stand der Technik (d. h. BrainIAC, BrainSegFounder und MedicalNet) bei 21 von 23 Aufgaben. Zweitens unterstützt ein bedingter Diffusionstransformator (DiT), der auf diesen klinisch informativen Embeddings trainiert wurde, sowohl die bedingte Generierung über sechs Variablen hinweg als auch patientenspezifische Längsschnittvorhersagen. Zusammen etablieren diese Ergebnisse einen einheitlichen 3D-Gehirn-MRT-Embedding-Raum, der sowohl für nachgelagerte klinische Aufgaben als auch für kontrollierbare Generierung geeignet ist.

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.