BrainG3N: 一种用于可控3D脑部MRI生成的双用途分词器

# BrainG3N：面向可控3D脑部MRI生成的双用途分词器
来源：https://arxiv.org/html/2606.19651
Max Van Puyvelde∗1,2 maxvpuyv@stanford\.edu &H\. Ibrahim Gulluk∗3 gulluk@stanford\.edu Wim Van Criekinge2 wim\.vancriekinge@ugent\.be &Olivier Gevaert1 ogevaert@stanford\.edu 1生物医学数据科学系，斯坦福大学医学院 2数学建模、统计与生物信息学系，根特大学 3电气工程系，斯坦福大学

###### 摘要

三维（3D）脑部MRI是临床神经学和神经肿瘤学的核心，而生成模型可以增强代表性不足的队列、模拟疾病轨迹并支持隐私保护的数据共享。潜在扩散已成为建模成像数据的首选解决方案，但它对分词器提出了两个竞争性要求：编码器嵌入必须保留下游任务所需的临床信息，解码器必须重建解剖学上逼真的体积。现有的基于重建的分词器实现了第二个要求，却牺牲了第一个。为解决这一问题，我们引入了一种基于全3D掩码自编码器（MAE）的分词器，用于3D脑部MRI潜在扩散，将编码器和解码器解耦：冻结的3D MAE编码器产生临床信息丰富的嵌入，而专用的CNN解码器则从这些嵌入的线性投影重建体素。我们在来自18个公共队列的35,309个体积上预训练编码器，涵盖四种模态、十种疾病类别和200多个采集站点，并在两种设置下展示其双重效用。首先，在一个23任务的线性探针基准测试中，编码器在21/23任务上优于或匹敌最先进的模型（即BrainIAC、BrainSegFounder和MedicalNet）。其次，在这些临床信息丰富的嵌入上训练的条件扩散Transformer（DiT）既支持跨六个变量的条件生成，也支持患者特定的纵向预测。这些结果共同建立了一个单一的3D脑部MRI嵌入空间，既能用于下游临床任务，也能进行可控生成。

11footnotetext:同等贡献。

## 1 引言

在神经学和神经肿瘤学中，脑部MRI为从肿瘤诊断、治疗计划到阿尔茨海默症和帕金森病等神经退行性疾病的分期和监测等临床决策提供信息，并支持关于脑发育和老龄化的人群规模研究。3D脑部MRI上的生成模型可以在多个方向扩展这一实践：增强代表性不足的患者队列，生成患者特定的数字孪生[35 (https://arxiv.org/html/2606.19651#bib.bib58)]来模拟反事实疾病轨迹，以及实现跨机构的隐私保护队列共享——目前法规和后勤障碍阻止了真实成像的获取。实现这些应用需要全3D分辨率的生成模型，然而该领域仍多基于2D切片操作。由于在该尺度下直接体素空间生成在计算上不可行，生成流程已普遍收敛于潜在扩散[34 (https://arxiv.org/html/2606.19651#bib.bib55)]：编码器-解码器分词器首先将体积压缩为低维潜在空间，然后在该嵌入上训练扩散模型。此设置下的条件生成对分词器提出了两个不同要求。首先，编码器嵌入必须携带条件生成和下游临床任务所需的临床信息。其次，解码器必须重建解剖学上逼真的体素体积。现有的3D放射学潜在扩散流程[31 (https://arxiv.org/html/2606.19651#bib.bib14)、18 (https://arxiv.org/html/2606.19651#bib.bib8)、44 (https://arxiv.org/html/2606.19651#bib.bib15)、41 (https://arxiv.org/html/2606.19651#bib.bib16)、17 (https://arxiv.org/html/2606.19651#bib.bib17)]针对重建目标训练单个编码器-解码器，这使编码器偏向于体素保真度而牺牲临床内容；产生的潜在空间通常仅根据体素级重建指标进行评估。

我们提出了一种针对3D脑部MRI的双用途自监督方法，其中冻结的掩码自编码器（MAE）编码器产生一个服务于两个角色的嵌入空间：用于下游任务的临床表示，以及条件扩散Transformer（DiT）的特征空间。为了在潜在扩散流程中使用此嵌入空间，编码器通过嵌入的线性投影与CNN解码器配对（§2 (https://arxiv.org/html/2606.19651#S2)）。使用MAE作为下游扩散的分词器已有2D先例[9 (https://arxiv.org/html/2606.19651#bib.bib4)]；同样的方法是否能迁移到3D放射学数据——其中语料库规模小几个数量级，输入维度更高，相关轴线是亚视觉表型而非可见物体类别——正是我们在此解决的问题。

我们的工作有两项主要贡献。首先，在18个队列的35,309个脑部MRI体积上预训练的冻结MAE编码器产生临床信息丰富的嵌入：在一个23任务的线性探针基准测试（§3.2 (https://arxiv.org/html/2606.19651#S3.SS2)）中，冻结编码器在21/23任务上优于或匹敌BrainIAC[36 (https://arxiv.org/html/2606.19651#bib.bib13)]、BrainSegFounder[11 (https://arxiv.org/html/2606.19651#bib.bib12)]和MedicalNet[10 (https://arxiv.org/html/2606.19651#bib.bib39)]；例如，编码器在异柠檬酸脱氢酶1（IDH1）突变状态预测（胶质瘤诊断和治疗分层的关键基因组生物标志物）上达到AUC 0.937，在肿瘤分级分类上达到AUC 0.921，在脑龄回归上平均绝对误差为4.43年，在性别预测上达到AUC 0.967。其次，在同一嵌入上训练的条件扩散Transformer（DiT）[29 (https://arxiv.org/html/2606.19651#bib.bib2)、24 (https://arxiv.org/html/2606.19651#bib.bib5)]支持跨六个变量的可控生成（§3.3 (https://arxiv.org/html/2606.19651#S3.SS3)）和患者特定的纵向预测（§3.4 (https://arxiv.org/html/2606.19651#S3.SS4)）；生成的嵌入由CNN解码器映射回高保真3D体素体积。在这两种情况下，生成的样本被使用真实数据训练的 classifiers 正确恢复：例如，在横截面年龄条件上，Pearson r=0.93，在纵向年龄进展上，Pearson r=0.72。这一迁移测试将嵌入对临床表型的编码与其生成可控性联系起来。

## 2 方法

#### 预训练语料库。

我们的预训练语料库包含来自17,399个独特受试者的35,309个脑部MRI体积，这些体积来自18个公共队列和200多个采集站点。体积涵盖四种模态（T1、T2、液体衰减反转恢复[FLAIR]和T1对比增强[T1c]）和包括健康对照、神经退行性疾病（阿尔茨海默病、帕金森病）、神经发育状况、精神疾病和脑肿瘤在内的十种临床类别；受试者年龄从5岁到98岁。一个统一的预处理流程通过ANTs仿射配准[4 (https://arxiv.org/html/2606.19651#bib.bib22)]将每个体积配准到SRI24图谱[33 (https://arxiv.org/html/2606.19651#bib.bib20)]，使用HD-BET[22 (https://arxiv.org/html/2606.19651#bib.bib21)]进行颅骨剥离，并使用N4偏置场校正[38 (https://arxiv.org/html/2606.19651#bib.bib23)]校正强度不均匀性，产生160×192×160个体素、1mm各向同性间距的体积。特定队列的入口点处理处于不同预处理阶段（已剥离、去面部、原生DICOM等）的原始输入，而无需双重处理。完整的数据集卡片和每个队列的流程见附录A (https://arxiv.org/html/2606.19651#A1)和B (https://arxiv.org/html/2606.19651#A2)。

#### MAE编码器。

编码器是一个3D掩码自编码器[19 (https://arxiv.org/html/2606.19651#bib.bib1)]，构建在12层视觉Transformer[15 (https://arxiv.org/html/2606.19651#bib.bib24)]上，隐藏维度为1152，使用16^3大小的补丁，每个体积产生1200个令牌。在预训练期间，70%的补丁被随机掩码：编码器仅处理360个可见补丁，一个单独的Transformer解码器从编码器输出中重建840个掩码补丁，使用逐补丁的均方误差损失。从360个可见补丁重建840个缺失补丁需要建模长距离解剖上下文，这迫使编码器捕获全局结构关系而非局部体素统计。这一特性正是我们在下游中利用的，既用于临床预测，也作为扩散分词器的输入空间。

(a) 两阶段分词器

参见标题

\phantomcaption

参见标题(b) 条件流匹配 DiT

\phantomcaption

图1：架构。(a) 阶段1在70%掩码补丁重建上预训练3D MAE编码器；阶段2冻结编码器并训练线性投影P∈R^{1152×32} + 3D CNN解码器，使用体素l1损失。相同的冻结特征空间z' = zP 被探针使用，并由DiT生成。(b) 噪声令牌x_t通过12块DiT堆栈，使用adaLN-Zero调制和条件向量c，产生32通道速度v̂。分类条件使用嵌入查找，带有K+1个空槽位用于CFG dropout；年龄使用正弦+MLP头部。

#### 两阶段分词器（图1 (https://arxiv.org/html/2606.19651#S2.F1)(a)）。

分词器通过线性投影将冻结的MAE编码器与3D CNN解码器耦合。预训练编码器被冻结；一个线性投影P∈R^{1152×d'}将1200个令牌中的每一个从1152通道压缩到d'=32通道；一个3D CNN解码器φ在l1损失下从投影令牌重建体素：

z = Enc(x) [冻结]， z' = zP， x̂ = φ(z'). (1)

与CNN-VAE分词器一样，联合训练编码器和解码器针对同一重建目标会使编码器偏向于局部强度保真度并降低其临床内容。瓶颈d'=32是一个刻意权衡：较小的瓶颈减少了扩散模型的输入维度，同时仍保留了编码器的大部分临床内容（扫描见表5 (https://arxiv.org/html/2606.19651#A4.T5)）。

#### 线性探针。

我们通过线性探针（自我监督表示学习中的标准协议）评估冻结编码器嵌入的临床内容：在每个任务上，在冻结嵌入之上拟合一个线性分类器，无需微调编码器。每个探针的输入是编码器的1200令牌输出，经平均池化后得到每个体积的单个d=1152维向量。我们使用逻辑回归进行分类任务，岭回归进行回归任务，在5折按受试者分组的带分层交叉验证下评估，确保任何折的训练和测试分割中不会出现同一受试者。完整的23任务面板和每模态细分见附录E (https://arxiv.org/html/2606.19651#A5)。

#### 条件潜在扩散（图1 (https://arxiv.org/html/2606.19651#S2.F1)(b)）。

一个流匹配扩散Transformer（DiT）[29 (https://arxiv.org/html/2606.19651#bib.bib2)、24 (https://arxiv.org/html/2606.19651#bib.bib5)]在投影令牌序列z'∈R^{1200×32}上训练。DiT具有12个块，隐藏维度1152，18个注意力头，并使用流匹配目标训练：

x_t = (1-t) x_1 + t ε， L = ||v_θ(x_t, t, c) - (ε - x_1)||_2^2， (2)

其中x_1是真实潜在，ε∼N(0,I)，条件向量c是六个条件嵌入的总和，通过adaLN-Zero调制[29 (https://arxiv.org/html/2606.19651#bib.bib2)]实现：modulate(h; c) = h (1 + scale(c)) + shift(c)。六个条件是：疾病（8类）、性别、模态（4类，永不丢弃）、采集站点（19类）、年龄（连续）和IDH1突变状态（二分类）。无分类器引导（CFG）[20 (https://arxiv.org/html/2606.19651#bib.bib26)、12 (https://arxiv.org/html/2606.19651#bib.bib25)]通过在训练期间以概率p=0.1将每个条件独立替换为null嵌入来实现；在采样时，速度按条件外推为v = v_uncond + s (v_cond - v_uncond)，s=2.0（CFG尺度敏感性见附录G (https://arxiv.org/html/2606.19651#A7)）。模态从不丢弃，因为在推理时总是指定，所以空模态分支只会浪费容量。DiT的纵向变体（§3.4 (https://arxiv.org/html/2606.19651#S3.SS4)）重用同一冻结分词器和adaLN-Zero条件，将噪声到数据的插值替换为基线到随访的潜在桥接，以在请求的时间范围内实现患者特定的预测。

## 3 实验

第3.1节 (https://arxiv.org/html/2606.19651#S3.SS1)在小规模基准上验证架构选择；第3.2–3.4节 (https://arxiv.org/html/2606.19651#S3.SS2)在三种设置中评估产生的嵌入：横截面探针、条件生成和纵向预测。

### 3.1 匹配尺度下的架构验证

在扩展到完整语料库之前，我们在一个1100体积的肿瘤队列（UCSF-PDGM++UPENN-GBM）上验证了两个架构选择：投影瓶颈d'，以及MAE–CNN与CNN-VAE分词器的选择。

#### 瓶颈维度d'。

我们扫描投影维度d'∈{32,128,512}，并在每个值下训练阶段2解码器。重建质量随d'单调增加，在d'=128以上收益递减。为评估投影是否保留了下游任务依赖的临床信息，我们还使用线性分类器在每个d'下探针投影特征z'。d'=32时的IDH1探针AUC为0.861，比探针原始1152维编码器嵌入（0.883）的AUC低0.022，维度不到总维度的3%（38K对比1.4M）；WHO肿瘤分级上的相应差距为0.055。因此我们在本文其余部分采用d'=32：较小的瓶颈减少了扩散模型的输入维度和训练成本，而线性投影在z'中保留了足够的编码器临床内容以支持下游条件（完整扫描见附录D (https://arxiv.org/html/2606.19651#A4)，表5 (https://arxiv.org/html/2606.19651#A4.T5)）。

#### 分词器架构。

我们将冻结的MAE–CNN分词器与自编码器进行比较。