Tokenizer-Generator Coupling in Medical Image Generation¶
会议: NeurIPS2026
arXiv: 2608.07713
代码: https://github.com/liamchalcroft/tokenizer-generator-coupling
领域: 医学图像 / 图像生成
关键词: 图像分词器、离散量化、生成器耦合、采样预算、可建模性
一句话总结¶
这是一项以“分词器—生成器—采样器”三元组为单位的受控因子研究:在 ChestMNIST-64 的 70 个生成配置中,重建质量不能可靠预测生成质量,量化器排序随生成器变化,而验证集选择的低步数采样可将 LFQ-1024 + D3PM 的 FID-192 从 0.44 降至 0.09。
研究背景与动机¶
潜空间医学图像生成通常先训练一个 VQ 或 KL 正则化自编码器,再冻结它,把生成器接到潜变量上。实践中,分词器往往按重建 PSNR、SSIM 或 LPIPS 挑选,之后才比较扩散模型。这个流程隐含了一个未经充分检验的假设:保存图像信息最好的表示,也是最容易被后续生成器学好的表示。自然图像研究已经显示两者可能脱钩,但低分辨率医学影像的解剖结构更集中,是否仍有这种关系,需要直接实验而不是照搬结论。
VQ 的学习码本、LFQ 的二值符号码、FSQ 的标量格点可以得到相似重建,却产生不同的离散 token 分布;自回归预测、置信度解掩码和吸收态扩散又以不同方式学习这些分布。若只换分词器而固定生成器,或只换生成器而固定分词器,就可能把某个组合的优势误认为某一组件的普遍优势。即使训练检查点不变,采样温度与步数也可能改变排序,因此“某模型优于另一模型”还必须交代推理预算与采样规则。
本文没有提出一个统一的新网络,而是把这些选择放入共同实验协议中:先匹配离散词表和空间网格,再比较生成器训练配方,最后在独立验证集选择采样设置。核心 idea:分词器的生成价值取决于它与生成器、采样器及推理预算的组合,不能用重建保真度或单一 token 统计量替代三元组的实际生成评估。
方法详解¶
整体框架¶
输入为 \(64\times64\) 单通道图像,分词器将其压缩成 \(8\times8\) 潜网格,生成器在冻结分词器产生的潜表示上训练,采样后经对应解码器还原图像。主要离散面板包含 VQ/LFQ/FSQ、1,024/2,048/4,096 三档词表与六类生成器的交叉组合,共 54 个配置;另有八种不同的 VAE/AE 设置各配 LDM/RF,形成 16 个连续参考配置,总计 70 个。
连续通道扫描与 KL 扫描共享 \(c=4,\lambda_{\mathrm{KL}}=10^{-6}\) 这一角点,只算一个设置,不能重复计数。全网格每个配置只有一个训练种子;三种量化器与 AR/MaskGIT/D3PM 的词表 1,024 子块,另加 LFQ + SEDD,才在 42/43/44 三个种子上重训生成器,分词器保持冻结。
下图表示研究中的训练、采样与评估流程,不是一套新提出的网络架构。虚线表示训练目标或验证选择对后续阶段的约束,实线表示潜表示与生成样本的数据流。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
X["真实训练图像"] --> A["匹配量化面板"]
L["重建与正则目标"] -.-> A
A -->|冻结编码器与解码器| B["生成器配方交叉"]
O["各家生成训练目标"] -.-> B
B -->|固定检查点| C["验证集采样选择"]
V["独立验证集"] -.-> C
C -->|采样并解码| D["分层证据评估"]
D --> Y["三元组比较结果"]
关键设计¶
1. 匹配量化面板:先区分量化机制,再讨论表示容量
所有分词器使用共同卷积编码器—解码器骨干,通道为 64/128/256,每个分辨率两个残差块,在 \(16\times16\) 使用单头自注意力,空间压缩倍数为 8。共同骨干约 8.5M 参数;控制的是骨干与空间网格,不意味着所有潜表示的信息量相等。每个位置只保留一个离散 token,故没有把残差量化加入主比较:多个码的层间依赖会改变后续建模问题。
VQ 在学习码本中按余弦相似度选择条目,用 EMA 更新码本、直通估计器传梯度,并用 commitment loss 约束编码输出;解码前码向量做 L2 归一化。LFQ 对各维取符号,将二值向量映射成 token,没有学习码本;它降低单位置分配熵、提高批量平均熵,使编码既确定又覆盖更多码。FSQ 将各维先有界化、再独立舍入到固定等级,联合等级决定 token;没有辅助量化损失,但码频率不受 LFQ 式熵正则约束。
在词表 1,024 时,LFQ 用 10 个二值维度,FSQ 用五个四级标量;FSQ 在 2,048/4,096 时分别使用等级向量 (8,4,4,4,4)/(8,8,4,4,4)。三档词表匹配的是离散原始 token 位数,而不是实际熵编码码率:
连续参考则扫描 \(c\in\{1,2,4,8,16\}\)、固定 \(\lambda_{\mathrm{KL}}=10^{-6}\),以及固定 \(c=4\) 时的 \(10^{-5}/10^{-6}/10^{-7}/0\) KL 权重;零权重对应 AE。连续位数以 \(c\times8\times8\times16\) 的 float16 存储量作代理,不与离散码率匹配,也不代表临床信息含量。
这个控制允许问“相同词表大小下哪种量化更适合某生成器”,却不能据此宣布离散表示比连续表示更高效。词表 1,024 的三种离散方法均达到 100% 码利用率,归一化边际熵为 VQ 0.997、LFQ 0.999、FSQ 0.959,因此结果不是简单的 VQ 码本坍缩对比;所有离散生成器仍处理展平的类别 ID,没有专门利用 LFQ 位结构或 FSQ 标量分解。
2. 生成器配方交叉:把同一潜分布交给不同生成机制
共同 DiT 风格骨干有 12 层、宽度 512、8 个注意力头,使用 RoPE 和 SiLU,但这不是等参数、等训练预算实验。AR/MaskGIT 不带时间调制,词表 1,024 时为 38.9M 参数;DFM/D3PM/SEDD/BFN 带 adaLN-Zero 时间条件,为 57–58M;连续 LDM/RF 为 57.1M。学习率和训练轮数也按家族调整,故跨生成器差异应解释为完整配方差异,而不是架构的纯因果效应。
AR 按光栅顺序进行 next-token prediction,生成 64 个位置需要 64 次前向调用,并用 KV 缓存复用前缀计算。MaskGIT 从全掩码网格开始,同时预测未确定位置,依据带 Gumbel 噪声的置信度逐步提交 token,默认 12 轮;它的条件边际预测不同于 AR 的严格顺序联合分解。二者比较的是顺序建模与双向、置信度驱动的并行提交。
D3PM 用余弦日程将 token 逐渐吸收到 MASK,训练网络预测原始类别,再用闭式后验执行逆转移。SEDD 在这里共用 D3PM 的前向过程、骨干与逆采样器,换成简化 score-entropy 目标并加入权重 0.001 的交叉熵稳定项;它不是使用 tau-leaping 的 canonical SEDD,不能把这里的相近结果推广成两种原始算法等价。
DFM 以 MASK 到数据的连续时间概率路径训练离散跳跃速度,使用三次多项式日程,默认 100 个 Euler 步。BFN 的状态则是每个位置的类别概率分布,网络预测 receiver 分布,再以 Bayesian update 更新状态;其准确度日程尺度设为 \(\beta=\sqrt{2\ln K}\),默认 1,000 步。这样主面板不只比较不同 loss,也包含不同状态表示与更新动力学。
连续 LDM 在 VAE/AE 潜表示上预测高斯噪声,默认用 1,000 个随机 DDPM 逆步;RF 预测噪声到数据直线路径的速度,默认积分 100 个 ODE 步。本文在提取训练潜变量时使用 VAE 的后验均值,而不是每个 epoch 重新采样后验。AE 和 VAE-c1 使用训练集估计的潜变量标准化,生成后反标准化再解码,其他连续行使用原始潜变量;这种不统一的归一化也限制了目标函数层面的归因。
3. 验证集采样选择:不能把默认步数当作模型上限
主表默认温度为 0.9,各生成器使用自己的默认步数。作者固定已经训练好的检查点,在词表 1,024 的三种量化器上分别扫描温度、步数、top-k 或 mask 日程;每个候选生成 2K 样本,与官方验证集的 11,219 张图像比较。选定设置之后才在测试集生成 10K 样本评估一次,因此主张不是在测试集直接挑选最优值。
每个分词器的候选数为 AR 18、MaskGIT 60、DFM 25、D3PM 20、SEDD 20、BFN 30。搜索规模并不相等,但 D3PM/SEDD 的大幅提升不是来自更大的候选网格。2K 验证筛选的 real-vs-real FID-192 噪声基线为 0.010,小幅改善仍可能不稳定;“验证选择”不消除有限样本与模型选择的不确定性。
D3PM/SEDD 减少步数时会重新构造同一余弦家族的 \(N\) 步吸收过程,不是对原 1,000 步链做任意跳步,也不重新训练。对已经解掩码的位置,逆过程保持类别不变;对仍为 MASK 的位置,预测的干净类别分布乘以下式的解掩码概率,剩余质量留给 MASK:
这里 \(\bar\alpha_n\) 是累计保留概率,SEDD 接收归一化时间 \(n/N\)。该形式也使 D3PM 不必存储各时间步的稠密转移矩阵:转移矩阵存储从 \(O(TK^2)\) 降至 \(O(T)\),但每步预测仍涉及 \(O(BLK)\) 的类别计算,logits、嵌入、激活与优化器状态仍依赖词表;“matrix-free”不是总计算或总显存与词表无关。
验证选择后的 LFQ D3PM 用温度 0.7、100 步,SEDD 用温度 0.7、500 步;FSQ 的两者都选温度 1.0、250 步;VQ 分别选温度 0.5、500 步与温度 0.7、250 步。最佳步数随分词器变化,正是三元组而不是单个组件应成为比较单位的直接证据。
4. 分层证据评估:区分指标噪声、训练方差与临床意义
测试使用 EMA 权重,每个生成配置产生 10K 图像,与 22,433 张 ChestMNIST 测试图像比较。FID-192 取 ImageNet 预训练 InceptionV3 第二个最大池化块的全局平均 192 维特征;它是内部排序指标,不是临床准确率,其绝对数值不能直接与常见 FID-2048 文献比较。
10K real-vs-real bootstrap 的 FID-192 均值为 0.002、95% 区间为 [0.001,0.004];这是指标估计噪声,不包含训练种子方差。对 54 个默认离散配置的 log(FID-192) 做描述性平方和分解时,生成器家族占 59.5%、量化器占 10.1%、量化器—生成器交互占 8.4%、词表—生成器交互占 15.0%。由于每格只有一个训练种子,这不是带实验误差项和显著性检验的因子分析,不能把残余交互当作随机噪声。
真正的训练重复只覆盖上述词表 1,024 子块。九个成对量化器比较中四个超过三倍合并种子标准差,支持量化器—生成器非可分性;其余五个不作可靠排序。用同一批图像计算 FID-2048 时 D3PM 的 LFQ 优势保持,但 AR 与 MaskGIT 的赢家可能变化,因此应保留“存在交互”,而不是固化一份全局量化器排行榜。
作者另以 domain-FID、无标签 classifier two-sample test 和最近邻筛查补充检查。domain-FID 仅有六个代表配置,相关性 0.71、精确置换 \(p=0.14\),不足以单独验证细粒度排序,更没有验证全部调优结果。最近邻比率定义为生成到训练的最近邻距离中位数除以训练集内部最近邻距离中位数;AuthPct 是生成样本比其最近训练邻居自身的最近邻还近的比例。这些量只能筛查特定距离空间中的复制迹象,不能证明隐私安全。
一个完整示例¶
以LFQ-1024与D3PM组合为例,真实训练图像先由冻结分词器转成64个离散位置,生成器学习该token分布;验证阶段不重新训练模型,而是比较不同温度与逆过程步数。验证选定温度0.7、100步后,才生成10K测试图像并解码评估。
表7中FID-192由默认0.44降至0.09,说明默认采样并不是这一检查点的性能上限。这里是已报道配置的数据流,不是一张特定胸片的诊断过程;FID下降也不能证明临床结构或患者隐私得到保证。
损失函数 / 训练策略¶
分词器训练 50 epochs,使用 L1 重建(权重 4.0)、VGG-16 感知损失(0.5)、多尺度 PatchGAN(0.05,第 10 epoch 启动)与 LeCam 正则(0.001);各量化器再加入适用的专属项,FSQ 没有辅助量化损失。VAE 的 KL 权重用 10 epochs 线性预热。
输入 min–max 归一化到 [0,1],不做数据增强;分词器使用 AdamW,学习率 \(10^{-4}\),batch size 128,按验证总 loss 选择检查点。生成器随后读取冻结分词器的 token 或 float16 连续潜文件,不联合更新分词器。
离散生成器训练 100 epochs,连续生成器训练 200 epochs。AR/MaskGIT/DFM 学习率为 \(3\times10^{-4}\),其余为 \(10^{-4}\);共同使用 AdamW、余弦学习率、5% warmup、EMA 0.9999、梯度裁剪 1.0 和 BFloat16。主网格在单张 48GB A6000 上训练,部分重复实验使用 L40S;这些共同条件不能消除参数与训练配方差异。
实验关键数据¶
主实验¶
下表摘自原文表 5:ChestMNIST-64、10K 测试生成样本、默认采样、单种子,指标为 FID-192,越低越好。连续参考仅可在完整配方层面比较。
| 分词器 | AR | MaskGIT | DFM | D3PM | SEDD | BFN | LDM | RF |
|---|---|---|---|---|---|---|---|---|
| VQ-1024 | 0.52 | 1.86 | 4.42 | 1.30 | 1.31 | 5.98 | — | — |
| LFQ-1024 | 0.33 | 1.91 | 1.77 | 0.44 | 0.41 | 2.27 | — | — |
| FSQ-1024 | 0.39 | 1.15 | 3.02 | 1.21 | 1.05 | 7.60 | — | — |
| VQ-4096 | 0.31 | 2.03 | 3.39 | 3.69 | 3.54 | 4.22 | — | — |
| VAE-c8 | — | — | — | — | — | — | 1.61 | 0.32 |
| AE | — | — | — | — | — | — | 0.09 | 0.07 |
重建方面,LFQ-1024 PSNR 为 28.0 dB、FSQ-1024 为 29.4 dB,却并不能推导 FSQ 在 D3PM 下更好;连续 VAE-c8 的 PSNR 达 34.4 dB,却是通道扫描中生成最差的配置。九个离散分词器的 PSNR 与各自最佳默认生成 FID 的 Spearman 相关性为 −0.03,连续通道扫描为 +0.10。
原文 §4.2 将 LFQ-1024 + AR 的 0.33 称作“最佳离散结果”,但表 5 以及 §4.1/§6.1 明列 VQ-4096 + AR 为 0.31。这里按表值保留两者,不改写为一致;作者同时指出其小差距不超出已观察的 AR 种子波动,不能据此可靠排名。
消融实验¶
这里的“消融”主要是固定检查点改变采样设置,而非删除一个新模块。原文表 7 的数字均为 10K 测试样本结果;候选设置先经 2K 验证样本筛选。
| 生成器 | LFQ-1024:默认 → 调优 | FSQ-1024:默认 → 调优 | VQ-1024:默认 → 调优 |
|---|---|---|---|
| AR | 0.33 → 0.33 | 0.39 → 0.29 | 0.52 → 0.46 |
| MaskGIT | 1.91 → 1.89 | 1.15 → 0.88 | 1.86 → 1.42 |
| DFM | 1.77 → 1.58 | 3.02 → 1.47 | 4.42 → 3.13 |
| D3PM | 0.44 → 0.09 | 1.21 → 0.13 | 1.30 → 0.15 |
| SEDD | 0.41 → 0.10 | 1.05 → 0.10 | 1.31 → 0.13 |
| BFN | 2.27 → 2.13 | 7.60 → 5.80 | 5.98 → 4.10 |
下表摘自原文表 9 的默认采样三种子结果;均值与标准差保持原文精度。它与前两表不是同一次评估,不能把重新计算的值拼进主表。
| 词表 1,024 配置 | 三种子均值 FID-192 | 种子标准差 |
|---|---|---|
| VQ + MaskGIT | 2.16 | 0.214 |
| LFQ + MaskGIT | 1.92 | 0.052 |
| FSQ + MaskGIT | 1.15 | 0.012 |
| VQ + D3PM | 1.32 | 0.043 |
| LFQ + D3PM | 0.39 | 0.022 |
| FSQ + D3PM | 1.20 | 0.042 |
关键发现¶
- FID-192 下,MaskGIT 更适合 FSQ,D3PM 更适合 LFQ;这两个方向的反转在三种子子块中有支持。但 FID-2048 改变了 MaskGIT 的赢家,稳健结论是存在耦合,D3PM 的 LFQ 优势跨指标保持。
- LFQ + D3PM 调优后三个种子的 FID 为 0.09/0.15/0.10,SEDD 为 0.10/0.11/0.13,两者平均均为 0.11。各种子独立调优才得到这一结果;直接把种子 42 的 D3PM 设置用于种子 43 会得到 0.34,不能声称具体超参数无需重新选择。
- LFQ 调优 D3PM 的吞吐为 33.2 samples/s,对比默认 3.67;调优 SEDD 为 6.60,对比默认 3.67。缓存 AR 仍为 146,说明较少 NFE 不等于相同墙钟成本;计时包含采样与解码,不含指标计算,默认与调优 batch size 分别为 128/64。
- 六个代表配置在另两数据集上复现时,AE + RF 的 FID 为 Chest/Pneumonia/OrganAMNIST 0.07/1.05/0.94。Pneumonia 上 VQ + AR 的 2.09 优于 LFQ + AR 的 2.73,但将 Chest 训练集缩至 4,700 后 LFQ 0.95 仍优于 VQ 1.11,否定了“仅由样本量解释反转”的假设。
亮点与洞察¶
- 把采样器纳入实验单位,避免把不合适的默认 1,000 步预算误判为离散扩散的固有劣势。固定检查点就能显著改变质量与成本,值得在增加训练算力之前检查。
- 共同骨干与匹配词表让量化比较更可解释,但作者没有把它包装成严格等计算实验。最有价值的是辨认组合效应,而不是宣布一个全局最优模型。
- 边际熵、位置熵离散度、左右邻居条件预测增益都不能可靠预测最佳生成 FID。模型无关的压缩统计并不是特定生成器和预算下的可建模性的充分代理。
局限与展望¶
- 主要证据来自无条件 \(64\times64\) ChestMNIST,另两套 MedMNIST 的复现仅覆盖六个配置。代码支持高分辨率和 3D 不等于已经在这些任务上验证,更不能据此推断真实临床影像的诊断保真度。
- 全 70 配置并非三种子实验,连续通道异常尤其可能受训练波动影响。下一步应扩展重复种子,统一连续潜变量归一化,再分离表示容量、目标函数和噪声日程的影响。
- 调优离散扩散的 0.09/0.10 接近默认连续参考,但连续模型没有得到等价采样搜索,且训练 epochs 不同。这不是公平预算下“离散追平连续”的最终结论。
- FID-192 与 FID-2048 在 12 配置上的相关性为 0.89,与无标签两样本判别排序为 0.86;domain-FID 的六配置检查更弱。仍需要放射科医师评估、病理条件生成、下游效用与更贴近医学的特征指标。
- 最近邻特征筛查的最小比率为 0.975、最大 AuthPct 为 0.107,没有明确检出复制迹象,但 AE + LDM 在原始像素/LPIPS 空间的比率为 0.873/0.821,值得更强审计。没有成员推断和患者级重复检查,不能把这些结果当作隐私认证。
- 已知区域的反事实修补仅在六个人造亮椭球示例上展示 IoU 0.31–0.96,区域事先提供。这是 oracle-mask inpainting 演示,不是盲异常定位或临床异常检测验证。
相关工作与启发¶
- vs VQ-VAE / MAGVIT-v2 / FSQ:这些方法提供不同量化机制;本文的贡献是匹配词表并交叉多个生成器,而非发明新的量化头。也没有验证按位 LFQ 或因子化 FSQ 输出头是否能进一步利用几何结构。
- vs canonical SEDD / D3PM:本文 SEDD 是 hybrid score-entropy loss 配 D3PM sampler 的特定实现。matrix-free 吸收态计算是可复用实现,但训练目标结论和低步数结论均需限定到这套接口与日程。
- vs rate-distortion-perception / modelability:本文支持“重建不足以选生成分词器”的经验解释,并未证明普适定律或医学图像的容量阈值。更直接的延伸是统一特征空间的重建 FID、生成到重建 FID gap,以及生成器条件预测误差。
- vs MedVAE / MedITok / CheXGenBench:前两者侧重医学表示,后者侧重端到端合成的保真、隐私与效用。本文提供组件级研究协议,不能替代这些面向临床或应用的评价轴。
- 实现分别见 medtokenizers 与 medlatents;适合复用的是可交换分词器、生成器及验证采样选择流程,而不是直接部署这些低分辨率无条件模型。
评分¶
- 新颖性: 4/5 — 将量化器、生成器和采样器作为耦合对象做受控医学图像研究,贡献主要是实验视角而非新架构。
- 实验充分度: 4/5 — 有 70 配置、局部三种子与多种稳健性检查,但完整重复、统一预算和临床验证仍不足。
- 写作质量: 4/5 — 配方、种子与指标边界较明确,仍有默认最佳离散结果表述不一致。
- 价值: 4/5 — 为潜空间生成提供实用选择协议和可复用实现,结论须限定于当前低分辨率研究场景。