BrainG3N：一种用于可控3D脑MRI生成的双用途分词器
BrainG3N: A Dual-Purpose Tokenizer for Controllable 3D Brain MRI Generation
摘要
三维脑MRI是临床神经病学和神经肿瘤学中的核心技术。生成式模型可以补充那些被忽视的病例数据，模拟疾病发展过程，同时支持隐私保护的数据共享。潜在扩散模型是建模影像数据的常用方法，但这种方法对分词器提出了两个相互冲突的要求：编码器生成的嵌入必须保留有用的临床信息，而解码器则必须能够重建具有准确解剖结构的体积数据。现有的基于重建的分词器虽然能实现后端的重建功能，但却以牺牲编码器的临床信息为代价。为了解决这一问题，我们提出了一种基于全体积掩模自编码器的分词器，该分词器将编码器和解码器分离开来：冻结的3D MAE编码器能够生成具有临床意义的嵌入，而专门的CNN解码器则可以通过对这些嵌入进行线性投影来重建体积数据。我们在来自18个公共数据集的35,309个体积数据上对编码器进行了预训练，这些数据集涵盖了四种模态、十种疾病类别以及200多个采集地点。我们的分词器在两种应用场景中都表现出良好的性能：首先，在23项任务的线性探测基准测试中，编码器在23项任务中的21项上优于或等同于最先进的模型（如BrainIAC、BrainSegFounder和MedicalNet）；其次，基于这些临床信息生成的嵌入所训练的条件扩散变换器能够支持六种变量条件下的条件生成以及针对特定患者的长期预测。这些结果表明，我们可以拥有一个能够同时满足下游临床任务和可控生成需求的单一三维脑MRI嵌入空间。
English Abstract
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.