BrainG3N: Un Tokenizador de Doble Propósito para la Generación Controlable de Resonancias Magnéticas Cerebrales en 3D

Resumen

La resonancia magnética cerebral tridimensional (3D) es fundamental en neurología clínica y neurooncología, donde los modelos generativos podrían aumentar cohortes subrepresentadas, simular trayectorias de enfermedades y apoyar el intercambio de datos que preserva la privacidad. La difusión latente ha sido la solución preferida para modelar datos de imágenes, pero impone dos demandas contrapuestas sobre el tokenizador: los embeddings del codificador deben retener la información clínica sobre la que actúan las tareas posteriores, y el decodificador debe reconstruir volúmenes anatómicamente fieles. Los tokenizadores existentes impulsados por reconstrucción logran el segundo objetivo a expensas del primero. Para abordar esto, presentamos un tokenizador basado en autoencoder enmascarado (MAE) completamente volumétrico para difusión latente de resonancia magnética cerebral 3D, que desacopla codificador y decodificador: un codificador MAE 3D congelado produce embeddings clínicamente informativos, mientras que un decodificador CNN dedicado reconstruye vóxeles a partir de una proyección lineal de esos embeddings. Preentrenamos el codificador en 35,309 volúmenes de 18 cohortes públicas que abarcan cuatro modalidades, diez categorías de enfermedades y más de 200 sitios de adquisición, y demostramos su doble utilidad en dos entornos. Primero, en un punto de referencia de sondeo lineal de 23 tareas, el codificador supera o iguala a los modelos de última generación (BrainIAC, BrainSegFounder y MedicalNet) en 21 de las 23 tareas. Segundo, un transformador de difusión condicional (DiT) entrenado con estos embeddings clínicamente informativos admite tanto generación condicional en seis variables como pronóstico longitudinal específico del paciente. En conjunto, estos resultados establecen un espacio de embeddings único para resonancia magnética cerebral 3D capaz de realizar tanto tareas clínicas posteriores como generación controlable.

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.