跳到主要内容
返回时间线
arXiv来源发表:

扩散模型训练不再依赖外部视觉编码器:复用VAE特征,SiT-XL/2的FID从17.2降到7.4

核心概要

CoVA把生成模型自带VAE的重建特征在线重组为语义与空间结构更强的监督信号,在ImageNet 256×256上把SiT-XL/2的FID从17.2降到7.4,且不引入外部视觉编码器。

Source-provided article image: CoVA: Composing VAE Representations for Generative Learning
Figure 1 · arXiv

深度剖析

扩散Transformer的训练可以不再依赖DINO或CLIP这类外部视觉编码器,而是复用生成模型自身VAE的特征作为中间监督。 此前REPA一类方法把中间生成特征对齐到预训练视觉编码器的表示,需要额外视觉骨干贯穿训练,带来计算和显存开销。 在ImageNet 256×256、相同400K步预算下,CoVA把SiT-B/2、SiT-L/2、SiT-XL/2的FID分别从33.0降到21.4、18.8降到9.3、17.2降到7.4。

一个轻量在线组合器把面向重建的VAE潜变量重新组织成兼具类别语义和空间结构的稠密特征,再分阶段使用。 原始VAE特征主要为重建学习,语义区分度和空间结构不足以直接作为生成学习的监督目标。 组合器用类别监督、掩码潜变量重建和水平翻转一致性三个目标联合训练,只增加8.5M训练期辅助参数,占SiT-XL/2的1.3%,推理时不保留。

表示对齐被当作分阶段信号:早期用组合特征对齐中间SiT层,后期撤掉对齐,改用组合器的类别后验做样本相关的软标签条件。 持续对齐在训练后期会拖累生成性能,而简单平均类别嵌入做软条件会同时改变语义方向和条件幅度。 在ImageNet 256×256上,CoVA在400个epoch、带无分类器引导时达到FID 2.03和IS 316.4,而原始SiT-XL/2在1400个epoch后为FID 2.06和IS 270.3。

启示与展望

报道覆盖的是ImageNet 256×256上的类别条件图像生成,模型为SiT-B/2、SiT-L/2和SiT-XL/2,图像经预训练Stable Diffusion VAE编码为潜变量。组合器与对齐投影器只在训练期使用,推理时丢弃,采样仍走标准SiT流程。适合关注扩散Transformer训练效率与表示对齐的研究人员和工程师阅读。

组合器容量从最小到最大变化时FID仅小幅波动,说明增益不依赖辅助模型容量。对齐在训练后期会变得有害,软标签条件若不做幅度修正会同时改变条件强度,带来保真度与覆盖率的权衡。这些结论来自ImageNet 256×256上的实验设置。