BrainG3N: Um Tokenizador de Duplo Propósito para a Geração Controlável de RM 3D do Cérebro

Resumo

A ressonância magnética cerebral tridimensional (3D) é fundamental na neurologia clínica e na neuro-oncologia, onde modelos generativos podem aumentar coortes sub-representadas, simular trajetórias de doenças e apoiar o compartilhamento de dados com preservação de privacidade. A difusão latente tem sido a solução preferida para modelar dados de imagem, mas impõe duas demandas concorrentes ao tokenizador: as incorporações do codificador devem reter as informações clínicas sobre as quais as tarefas subsequentes atuam, e o decodificador deve reconstruir volumes anatomicamente fiéis. Os tokenizadores existentes baseados em reconstrução alcançam o segundo objetivo às custas do primeiro. Para resolver isso, introduzimos um tokenizador totalmente volumétrico baseado em autoencoder mascarado (MAE) para difusão latente de RM cerebral 3D, desacoplando codificador e decodificador: um codificador MAE 3D congelado produz incorporações clinicamente informativas, enquanto um decodificador CNN dedicado reconstrói voxels a partir de uma projeção linear dessas incorporações. Pré-treinamos o codificador em 35.309 volumes de 18 coortes públicas abrangendo quatro modalidades, dez categorias de doenças e mais de 200 locais de aquisição, e demonstramos sua utilidade dupla em dois cenários. Primeiro, em um benchmark de sondagem linear com 23 tarefas, o codificador supera ou iguala modelos de última geração (ou seja, BrainIAC, BrainSegFounder e MedicalNet) em 21 das 23 tarefas. Segundo, um transformador de difusão condicional (DiT) treinado nessas incorporações clinicamente informativas suporta tanto a geração condicional em seis variáveis quanto a previsão longitudinal específica do paciente. Juntos, esses resultados estabelecem um espaço de incorporação único de RM cerebral 3D capaz de realizar tanto tarefas clínicas downstream quanto geração controlável.

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.