RoAR-VAE: Improving Latent Diffusability through Robust Prediction
Yang Yue, Fangyun Wei, Yitong Wang, Jinjing Zhao, Henry Zheng, Dong Chen, Gao Huang
cs.CV
2026-10-19
把 REPA-E 简化为从噪声潜变量预测干净目标,RoAR-VAE 生成均分 63.86 超最强基线 2.19 分,小字重建比 FLUX.1-VAE 高 16.93 个百分点
潜在扩散模型不直接在像素上学习。视觉 tokenizer(通常是个 VAE)先把图压成潜变量,扩散模型再在潜变量上生成。设计上有一对矛盾:潜变量要保住足够细节才能重建得像,又要排布得让扩散模型学得快,后者论文称为 diffusability(可扩散性)。高容量、弱正则的潜变量往往细节丰富却难建模。
改进路数主要有两条。一条在 VAE 训练时加语义监督,如 VA-VAE 把潜变量对齐到预训练视觉特征;另一条联合训练 tokenizer 和扩散模型,代表是 REPA-E。REPA-E 有个别扭的细节:flow-matching 损失对 VAE 编码器做了梯度截断,原论文发现放进这些梯度会导致潜变量方差坍缩,真正训到编码器的只有语义对齐那一项。联合训练到底贡献了什么,就是这篇的起点。
先做减法。把 REPA-E 的扩散目标和专属生成模块删掉,只留「从加噪潜变量预测 DINOv2 特征」一条通路,得到 RobustAlign。对照实验用同一 VAE 结构、同一下游训练:32k 步 GenEval 上 RobustAlign-VAE 拿 60.89,REPA-E tokenizer 拿 60.60,打平。为一个编码器联合训练完整扩散模型,增益约等于零。
由此提炼出原则:生成相关的信息,应从被严重污染的潜变量里仍然可预测。RoAR-VAE 用两个互补目标实现它:
两个目标都经加噪潜变量把梯度传回编码器,标准 VAE 损失照常作用在干净潜变量上,总损失为 LVAE + 1.5×RobustAlign + 0.5×RobustRecon。
机制上,各向同性高斯腐蚀下潜变量高方差方向的信噪比更高((1-t)²νi/t²),更抗噪;鲁棒预测把有用信息压进这些主导方向,而 Wang & Pehlevan(2025)证明扩散模型会先学高方差模态。PCA 可视化与之吻合:RobustAlign 的潜空间物体组织更清晰,RobustRecon 局部变化更丰富,合并则兼有。
tokenizer 为 84M 参数卷积自编码器,8× 下采样、32 通道,只在 ImageNet 256² 上训练:200k 步 VAE 预热加 250k 步 RoAR 精修。评估时冻结 tokenizer,同一个约 0.8B 的 MMDiT 在 LAION 上从零训 32k 步,测 GenEval、GenEval2、DPG-Bench 三个文生图基准。
| 方法 | 重建 LPIPS↓ | 生成均分 |
| FLUX.1-VAE | 0.0510 | 52.90 |
| VA-VAE(最强基线) | 0.1052 | 61.67 |
| REPA-E-FLUX | 0.0546 | 61.45 |
| FLUX.2-VAE(仅参考) | 0.0398 | 62.73 |
| RoAR-VAE | 0.0349 | 63.86 |
同预算下生成均分超最强基线 2.19 分,重建同步占优:PSNR 32.88 dB,比 FLUX.1-VAE 高 1.80 dB,rFID 0.1309。细粒度差距更大。TokBench 小字准确率 69.00%,FLUX.1-VAE 是 52.07%,高出 16.93 个百分点;小脸相似度 0.65 对 0.53,连 FLUX.2-VAE(67.33%、0.63)也被小幅超过。
消融定位了关键因素:Vanilla-KL 均分 50.50,单加 RobustRecon 到 59.38,单加 RobustAlign 到 62.49,合并 63.86;把两个目标的输入换成干净潜变量,掉到 58.47。起作用的是腐蚀本身,预测目标只是配合。架构上,语义预测器换成简单 MLP 会从 62.49 掉到 52.70,预测通路需要足够表达力。
tokenizer 质量同时决定扩散模型的训练效率和细节上限,是近期 latent 改造工作的焦点。RoAR 的贡献有两条。它证明 REPA-E 里最贵的部分,为 VAE 联合养一个扩散模型,可以整个砍掉:tokenizer 训练不再需要生成器,潜变量恢复可缓存,大规模训练便宜很多。它也打破了常见取舍:学得快的方法重建明显变差(VA-VAE 的 LPIPS 是 RoAR 的三倍),重建导向的 tokenizer 下游又学得慢。代码和模型已开源,做文生图系统的人可以直接换上。
论文自己把机制解释标为「可能的解释」:信噪比加谱偏置的论证只有 PCA 可视化这类定性证据,没有因果验证。训练和评估停在中等规模,ImageNet 256² 训 tokenizer,0.8B 下游模型、32k 步做比较,更高分辨率、更大模型、更长训练下 2.19 分的领先是否保持,没有验证;图 4 给了 120k 步的定性对比,但没给完整基准数字。FLUX.2-VAE 用未公开的数据训到 62.73,工业配方未必在同一比较平面。语义教师只测了 DINOv2 一种。2.19 分的绝对优势不大,这是一次机制讲得通的渐进改进。