BrainG3N : un tokenizer à double usage pour la génération contrôlable d'IRM cérébrale 3D

Résumé

L'imagerie par résonance magnétique (IRM) cérébrale tridimensionnelle (3D) est essentielle en neurologie clinique et en neuro-oncologie, où les modèles génératifs pourraient enrichir les cohortes sous-représentées, simuler des trajectoires de maladies et favoriser le partage de données respectueux de la vie privée. La diffusion latente constitue la solution de référence pour la modélisation des données d'imagerie, mais elle impose deux exigences concurrentes au tokeniseur : les plongements de l'encodeur doivent préserver les informations cliniques sur lesquelles reposent les tâches en aval, et le décodeur doit reconstruire des volumes anatomiquement fidèles. Les tokeniseurs actuels basés sur la reconstruction privilégient la seconde exigence au détriment de la première. Pour y remédier, nous introduisons un tokeniseur entièrement volumétrique fondé sur un auto-encodeur masqué (MAE) pour la diffusion latente d'IRM cérébrale 3D, en dissociant l'encodeur du décodeur : un encodeur MAE 3D figé produit des plongements cliniquement informatifs, tandis qu'un décodeur CNN dédié reconstruit les voxels à partir d'une projection linéaire de ces plongements. Nous pré-entraînons l'encodeur sur 35 309 volumes provenant de 18 cohortes publiques couvrant quatre modalités, dix catégories de maladies et plus de 200 sites d'acquisition, et démontrons sa double utilité dans deux contextes. Premièrement, sur un banc d'essai de sondage linéaire comprenant 23 tâches, l'encodeur surpasse ou égale les modèles de pointe (BrainIAC, BrainSegFounder et MedicalNet) pour 21 des 23 tâches. Deuxièmement, un transformateur de diffusion conditionnel (DiT) entraîné sur ces plongements cliniquement informatifs permet à la fois la génération conditionnelle pour six variables et la prévision longitudinale spécifique au patient. Ensemble, ces résultats établissent un espace de plongement unique pour l'IRM cérébrale 3D, capable à la fois de tâches cliniques en aval et de génération contrôlable.

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.