PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Max Van Puyvelde, Ibrahim Gulluk, Wim Van Criekinge, Olivier Gevaert

arXiv:2606.19651v2cs.AIcs.CVcs.LG

TL;DR

Existing 3D brain-MRI tokenizers struggle to preserve clinically informative embeddings while reconstructing anatomically faithful volumes. BrainG3N decouples a frozen 3D MAE encoder from a CNN decoder and uses the embeddings for conditional diffusion, outperforming or matching three foundation baselines on 21 of 23 tasks while supporting controllable generation.

  • Problem

    3D brain-MRI latent-diffusion tokenizers must preserve clinical information in encoder embeddings while reconstructing anatomically faithful volumes.

  • Method

    BrainG3N decouples a frozen volumetric MAE encoder from a dedicated CNN decoder and uses the encoder embeddings for downstream tasks and conditional diffusion.

  • Results

    The encoder outperformed or matched BrainIAC, BrainSegFounder, and MedicalNet on 21 of 23 linear-probing tasks, while the DiT supported controllable generation.

  • Takeaways & Limitations

    A single embedding space supports clinical representation and controllable generation across four structural brain-MRI modalities.

  • Takeaways & Limitations

    Transfer beyond brain MRI is untested, and longitudinal forecasting recovers only approximately 27% of true aging-trajectory magnitude.

Abstract

from arXiv · show

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.

1 Introduction

3D brain MRI generative models could expand clinical research and privacy-preserving data sharing, but full-resolution generation is computationally infeasible and current approaches often rely on 2D slices. BrainG3N addresses this with a dual-purpose MAE-based tokenizer whose frozen encoder supports clinical representation and conditional diffusion, while a CNN decoder reconstructs volumes.

  • Motivation: Generative 3D brain MRI could augment under-represented cohorts, simulate patient-specific counterfactual disease trajectories, and enable privacy-preserving cohort sharing.These applications target neurology, neuro-oncology, and population-scale research on brain development and aging.
  • Problem: Full-resolution 3D generation remains challenging because direct voxel-space generation is computationally infeasible, while much of the field still operates on 2D slices.Latent diffusion has therefore become a broadly adopted approach using an encoder–decoder tokenizer.
  • Approach: BrainG3N uses a frozen MAE encoder to produce one embedding space for downstream clinical tasks and a conditional diffusion transformer.A CNN decoder is connected through a linear projection of the embeddings for integration into a latent-diffusion pipeline.
  • Contributions: 35,309 brain MRI volumes from 18 cohorts support pretraining, and the frozen encoder outperforms or matches BrainIAC, BrainSegFounder, and MedicalNet on 21 of 23 linear-probing tasks.The pretraining data span the work’s stated cohort collection, while the benchmark evaluates clinically informative embeddings.
  • Contributions: AUC 0.937 is achieved for isocitrate dehydrogenase 1 mutation-status prediction, a key genomic biomarker for glioma diagnosis and treatment.This example illustrates the encoder’s clinical utility on the linear-probing benchmark.

2 Method

The method builds a dual-purpose 3D brain-MRI tokenizer by freezing a clinically informative MAE encoder and training a separate projected-token CNN decoder for voxel reconstruction. The resulting embeddings support linear clinical probing and conditional latent diffusion across six patient and acquisition variables.

  • Pretraining corpus: The pretraining corpus contains 35,309 brain MRI volumes from 17,399 subjects across 18 public cohorts and 200+ acquisition sites.It spans four modalities, ten clinical categories, and subject ages from 5 to 98 years.
  • MAE encoder: The 3D MAE uses a 12-layer vision transformer with hidden dimension 1152 and 163 patches, producing 1200 tokens while masking 70% of patches.The encoder processes 360 visible patches and reconstructs 840 masked patches with a separate transformer decoder under per-patch mean-squared-error loss.
  • Two-phase tokenizer: The two-phase tokenizer freezes the MAE encoder, projects 1152-dimensional tokens to 32 channels, and reconstructs voxels with a dedicated 3D CNN decoder under an ℓ1 loss.The projected feature space z′=zP is shared by the clinical probe and diffusion transformer.
  • Two-phase tokenizer: Joint reconstruction training can shift encoders toward local intensity fidelity and degrade clinical content, so the d′=32 bottleneck balances diffusion input dimensionality with preserved clinical information.The passage identifies this bottleneck as a deliberate trade-off supported by a sweep in Table 5.
  • Linear probing: Linear probes assess frozen encoder embeddings by mean-pooling 1200 tokens into a d=1152 volume vector, using logistic regression for classification and ridge regression for regression.The encoder is not fine-tuned during probing.
  • Conditional latent diffusion: The conditional flow-matching DiT trains on z′∈R1200×32 and conditions generation through six embeddings: disease, sex, modality, acquisition site, age, and IDH1 mutation status.The six conditions are routed through adaLN-Zero modulation, with modality never dropped and classifier-free guidance using independently null-replaced conditions at p=0.1.

3 Experiments

Experiments validate the tokenizer’s architectural choices and evaluate its frozen embeddings for cross-sectional clinical probing, controllable generation, and longitudinal forecasting. The embeddings support strong performance across clinical tasks, conditional attribute control, and patient-specific temporal extrapolation.

  • 3.1 Architectural validation: On a 1,100-volume tumor cohort, experiments validate the projection bottleneck d′ and the MAE–CNN versus CNN-VAE tokenizer design.The cohort combines UCSF-PDGM and UPENN-GBM volumes.
  • 3.1 Architectural validation: IDH1 probing AUC is 0.861 at d′=32, only 0.022 below probing the raw 1152-dimensional embeddings, while reconstruction improves with larger bottlenecks but shows diminishing returns above d′=128.The projection sweep uses d′ ∈{32, 128, 512}.
  • 3.1 Architectural validation: At matched dimensionality, MAE exceeds AutoencoderKL by +0.064 AUC on IDH1 and +0.046 AUC on WHO tumor grade, with the same advantage at d′=32.AutoencoderKL is trained from scratch on the same cohort at matched compute.
  • 3.2 Cross-sectional probing: 21 of 23 clinical linear-probing tasks outperform or match BrainIAC, BrainSegFounder, and MedicalNet on the 35,309-volume corpus.The panel contains 15 classification and 8 regression tasks; GDS and UPDRS-III are the two non-wins.
  • 3.3 Conditional generation: Conditional generation recovers requested attributes, with 0.99 agreement for HC-versus-AD disease control and Pearson r=0.93 for age, while IDH1 reaches mean conditional agreement of 0.52.Age predictions are 35.4 / 54.4 / 76.0 for requested 30 / 50 / 70 years; pooled gen-vs-recons FID is 34.4 and gen-vs-real FID is 107.3.
  • 3.4 Longitudinal forecasting: Longitudinal forecasting recovers approximately 27% of true aging, with slope 0.268 and Pearson r=0.716 between predicted and requested ∆t.Evaluation uses 64 held-out validation baselines and requested intervals ∆t ∈{0, 1, 2, 5, 10} y.

4 Discussion

BrainG3N’s central contribution is a unified MAE-CNN tokenizer and DiT whose decoupled encoder preserves clinically informative embeddings while its decoder reconstructs voxels faithfully across four MRI modalities. The discussion also highlights clinical potential, limited longitudinal magnitude recovery, untested cross-modality transfer, and released model resources.

  • Contributions: The frozen MAE encoder and dedicated CNN decoder outperform or match three published 3D brain-MRI foundation baselines on 21 of 23 linear-probing tasks while reconstructing voxels faithfully.The two non-wins, GDS depression and UPDRS-III, are near-floor regression for every encoder.
  • Contributions: A single MAE encoder, MAE-CNN tokenizer, and DiT span T1, T2, FLAIR, and T1c across pretraining, probing, and conditional generation.This contrasts with approaches using 2D slicewise pretraining or separate per-modality models.
  • Broader impact: MRI-based prediction could support clinical decisions involving IDH1 mutation status, APOE genotype, and cognitive scores currently requiring biopsy, blood draw, or dedicated examiner time.The cited cognitive scores are CDR, MMSE, and MoCA.
  • Limitations: 27% of true aging magnitude is recovered in longitudinal forecasting, with slope 0.268, despite recovering trajectory direction and anatomical loci.The discussion attributes this limitation to the stochastic-bridge regularizer interpolating toward the conditional mean.
  • Reproducibility: Model weights and loading and feature-extraction code are available under a non-commercial research-use license, with approximate total training of 200 H100-hours.Probing runs take approximately 2 hours per task on a 16-core CPU.

A Dataset card · B Preprocessing pipeline · C Hyperparameters

BrainG3N is trained on a broad, longitudinal, multimodal brain-MRI corpus processed through a harmonized spatial pipeline. The paper also reports stage-specific training hyperparameters and an approximate total training cost of 200 H100-hours.

  • A Dataset card: 35,309 volumes from 17,399 subjects span 18 public cohorts, 200+ acquisition sites, four modalities, ages 5–98, and 6,576 subjects with longitudinal scans.The modalities are T1, T2, FLAIR, and T1c.
  • A Dataset card: Healthy controls comprise 15,274 scans, while disease coverage includes MCI, GBM, PD, paediatric-mixed, AD, ASD, glioma, SCZ, and ADHD.Healthy subjects represent 43% of the corpus to support clinical probing and generative modeling.
  • A Dataset card: Clinical metadata coverage is 90% for age-at-scan and sex, 96% for diagnosis, and 78% for site, with disease-specific fields more complete within relevant cohorts.CDR and MMSE reach 99% coverage within AD/MCI, while tumor-grade coverage reaches 100% within tumor cohorts.
  • B Preprocessing pipeline: Every volume undergoes N4 bias-field correction, HD-BET skull stripping, and affine registration to the SRI24 atlas at 240×240×155 voxels with 1 mm isotropic resolution.The registered output uses LPS orientation; the supplied passage truncates the subsequent multimodality step.
  • C Hyperparameters: Stage-specific training hyperparameters are reported for each component in Table 4.The supplied passage provides the table reference but not its individual parameter values.
  • C Hyperparameters: ≈200 H100-hours cover 100 MAE-pretraining epochs, 80 Phase-2 epochs, and ≈1100 DiT epochs, while probing takes ≈2 CPU-hours per task.Probing runs use a 16-core CPU.

D AKL baseline training details

The AKL baseline is a canonical CNN-VAE tokenizer trained from scratch on the same 1,100-volume tumor cohort as MAE+CNN, with matched latent size and reconstruction-plus-KL loss. MAE+CNN generally preserves stronger clinical probing performance, while its advantage reflects encoder information rather than voxel-detail reconstruction alone.

  • Training setup: AKL was trained from scratch on 1,100 UCSF-PDGM and UPENN-GBM tumor volumes using reconstruction-plus-KL loss and a 614K-element latent matching MAE+CNN d′=512.Held-out reconstruction used a 20% test split, while clinical probing used 5-fold subject-grouped cross-validation.
  • Reconstruction: PSNR 33.62 vs 32.54 shows AKL competitive with MAE+CNN at d′=32, whereas MAE+CNN exceeds AKL at d′≥128.These comparisons come from held-out voxel reconstruction on the 1,100-volume tumor cohort.
  • Clinical probing: MAE+CNN outperforms AKL on IDH1 at every bottleneck and on WHO tumor grade above d′=32, where the models are essentially tied.The clinical probing results use 5-fold logistic regression, with AUC reported for the architectural validation.
  • Interpretation: The MAE-as-tokenizer advantage is attributed to the clinical information captured by encoder embeddings rather than raw voxel-detail preservation.This conclusion follows the reconstruction-versus-probing trade-off observed across bottleneck sizes.

E Full 23-task probing breakdown (per modality)

Table 6 presents frozen linear-probe performance for each task across modality splits in the 35,309-volume corpus. Classification uses AUC, regression uses R2, and dashes indicate modality subsets with too few samples.

  • Evaluation setup: Table 6 reports per-task, per-modality frozen linear-probe results from the 35,309-volume corpus.Evaluation uses 5-fold stratified group k-fold cross-validation grouped by subject.
  • Metrics: AUC is used for classification tasks, while R2 is used for regression tasks.The table also marks the best modality per task in bold.
  • Reporting conventions: Dashes denote modality subsets with too few samples for evaluation.The table defines task abbreviations including APOE, IDH1, MGMT, CDR, MMSE, and MoCA.

F Head-to-head competitor probing — full 23-task breakdown

Across all 23 tasks, frozen-feature linear probing evaluates four encoders under identical splits and probe code, reporting each encoder’s strongest modality configuration. Ours outperforms or matches every competitor on 21 of 23 tasks, with non-wins on two near-floor regression tasks.

  • Evaluation setup: The benchmark probes four frozen encoders on identical splits and probe code using the 35,309-volume corpus and 5-fold stratified subject-grouped cross-validation.Each encoder is reported at its best-performing modality slice for each task.
  • Head-to-head results: 21 of 23 tasks show Ours outperforming or matching every competitor across 15 classification tasks measured by AUC and 8 regression tasks measured by R2.The figure reports the best modality per encoder for each task.
  • Head-to-head results: The two non-wins, GDS depression and UPDRS-III, are near-floor regression tasks for every encoder.These exceptions account for the remaining two tasks in the 23-task comparison.

G CFG scale sensitivity

CFG scales 1.5 and 2.0 produced comparable controllable-condition performance, whereas scale 3.0 degraded age control through over-extrapolation. IDH1 remained near chance across the sweep, and the study selected s=2.0.

  • CFG scale sensitivity: CFG s=1.5 and s=2.0 were within noise on every controllable condition.The sweep used s ∈{1.5, 2.0, 3.0} with n=32 samples per arm.
  • CFG scale sensitivity: 0.94 → 0.82: Pearson r for continuous age dropped at CFG s=3.0 as strong guidance over-extrapolated away from the conditional manifold.The age axis degraded specifically at s=3.0.
  • CFG scale sensitivity: 0.51–0.52: IDH1 mean agreement stayed near chance throughout the sweep, reflecting difficulty steering toward the rare mutant class on the small T1c tumor subset.The passage attributes this behavior to class rarity and subset size rather than a CFG-scale effect.
  • CFG scale sensitivity: s=2.0 was used for subsequent controllable generation.This choice followed the observed stability of s=1.5 and s=2.0 and the age degradation at s=3.0.

H Generation fidelity — full decomposition

The evaluation decomposes 3D-FID across controllability arms to separate generator quality from the tokenizer’s reconstruction floor, while a nearest-neighbor audit finds no evidence of memorization. FID is computed from 15 central 2D slices per volume under the stated conditional DiT sampling protocol.

  • Evaluation protocol: 3D-FID uses 15 central slices per volume—five per anatomical view across three views—processed with Inception-V3.Real references come from the filtered 35,309-volume corpus; generated samples use conditional DiT sampling with CFG s=2.0 and a 50-step Euler ODE.
  • FID decomposition: Three FIDs are reported for each controllability arm: gen-vs-real, gen-vs-recons, and recons-vs-real.gen-vs-recons measures alignment with the tokenizer’s conditional latent distribution, while recons-vs-real defines the tokenizer floor; lower is better.
  • Memorization audit: 57.7 is the mean synthetic-to-real nearest-neighbor distance, versus 37.2 for real-to-real, yielding a ratio of means of 1.55.Distances are computed in the 1200×32 token latent, flattened to 38,400 dimensions, using a real pool of 33,639 and 1,088 synthetic samples.
  • Memorization audit: Zero of 1088 generated samples falls below the real-to-real 5th-percentile threshold of 14.4, supporting the verdict of no memorization.The synthetic distribution lies entirely to the right of the lower tail of the real-to-real distribution.

I Longitudinal sweep visualization

Longitudinal bridge sampling at fixed baseline produces progressively larger forecast differences as the requested interval increases from 1 to 5 years. These changes concentrate along cortical and ventricular regions in two held-out ADNI cases, consistent with expected age-related loci.

  • Longitudinal sweep visualization: Difference maps grow monotonically with the requested ∆t after subtracting the round-trip sampler noise floor.This isolates structural change driven by the requested forecasting interval.
  • Longitudinal sweep visualization: Two held-out ADNI cases were bridge-sampled at requested ∆t values of 1, 2, and 5 years from a fixed baseline.The cases were a healthy-control 75-year-old male and an Alzheimer’s-disease 77-year-old male.
  • Longitudinal sweep visualization: The interval-driven changes concentrate along the cortex and ventricles, the expected loci of age-related change.The visualization uses absolute voxel differences between each forecasted sample and the ∆t=0 sample.
Loading 2606.19651v2…