← 返回 2026-09-01

GenFirst:生成先于重构的稳定端到端潜在生成建模 GenFirst: Generation Before Reconstruction for Stable End-to-End Latent Generative Modeling

Guangting Zheng, Yiyuan Zhang, Tao Yang, Yunpeng Chen, Rui Zhu, Jiajun Deng, Yanyong Zhang 📅 2026-08-29 👍 56 2026-09-01 18:30
VAE 扩散模型 文本到图像 潜在生成模型 端到端训练 自回归生成

熵正则防坍缩、先生成后重构的调度,首次实现VAE与生成模型的稳定端到端联训

前置知识

变分自编码器(VAE)

编码器-解码器结构:编码器 $q_\phi(z|x)$ 为每张图像预测高斯后验的均值 $\mu$ 和方差 $\sigma^2$,经重参数化 $z=\mu+\sigma\odot\epsilon$ 采样 latent,解码器从 $z$ 重建图像。目标是重构损失($\ell_1$+LPIPS+GAN)加 KL 正则 $\mathcal{D}_{KL}(q_\phi(z|x)\|p_0(z))$,实践中权重极小(LDM 取 $10^{-6}$)。

本文全部理论分析都建立在 VAE 目标函数上:端到端联训就是让生成损失梯度也回传到这个编码器,不理解后验采样与 KL 正则的构造,就无法理解坍缩从何而来、熵项为何是解药。

KL 散度的先验拟合–熵分解

对标准正态先验,KL 可精确拆成两项:$\mathcal{D}_{KL}(q_\phi(z|x)\,\|\,p_0(z))=\mathbb{E}_{q_\phi}[-\log p_0(z)]-\mathcal{H}(q_\phi(z|x))$。第一项(prior fitting)把 latent 样本拉向先验的高概率区;第二项(负熵)防止后验退化成点质量,保持不确定性。两项方向相反,构成一组对冲力。

本文的核心诊断——先验–熵失衡——正是基于这一分解:熵项长期被小 KL 权重当作工程附属品压制,而生成损失额外注入强大的先验拟合力,两者失衡导致 latent 坍缩。看懂这个分解等于看懂论文第一节。

潜在坍缩(latent collapse)

端到端联训时后验方差 $\sigma^2_\phi(x)\to 0$、后验均值 $\mu_\phi(x)$ 塌成跨图像近似常数的病态现象:latent 失去图像间区分度。更隐蔽的是附录 B 揭示的捷径:当 $z\approx c$ 时流匹配目标退化为 $v\approx c-z_t/(1-t)$,扩散模型只需对噪声输入做仿射变换即可压低损失,训练看似正常实则已坍缩。

这是本文要解决的第一个根本问题,REPA-E 正是因此放弃直接端到端训练;论文标题中的 Stable 一词就是针对它。

流匹配与 SiT

流匹配训练速度场 $v_\theta(z_t,t)$ 去预测目标 $v=z-\epsilon_t$,其中噪声与数据按 $z_t=(1-t)\epsilon_t+tz$ 线性插值,$\epsilon_t\sim\mathcal{N}(0,I)$、$t\sim\mathcal{U}(0,1)$。SiT 是基于该目标的类条件扩散骨干,只有隐式似然、无法直接算 latent 的精确概率。

EiT 的生成先验就是 SiT/MMDiT。理解插值与速度目标,才能看懂常数 latent 捷径的数学推导,以及为何 SiT 情形要额外保留 $10^{-6}$ 的 KL 来控制 latent 数值尺度。

连续自回归先验与 GMM 头

把 VAE latent 展平为 $N=hw$ 个连续 token,因果 Transformer 逐位置预测下一 token 的条件密度;因 token 连续,用 $K$ 分量高斯混合头 $p_\theta(z_i|z_{<i})=\sum_k\pi_{i,k}\mathcal{N}(z_i;m_{i,k},\mathrm{diag}(s^2_{i,k}))$ 替代离散 softmax,可算精确 NLL。

EAR(论文的连续 AR 实例化)是全文分析的主力平台:精确似然让先验拟合强度可以被直接、连续地控制,从而干净地暴露生成–重构 trade-off 的形状。

gFID 与 rFID

FID(Fréchet Inception Distance)衡量图像集与真实图像集在 Inception 特征空间中的分布距离,越低越好。gFID 指生成任务的 FID(常分有无 classifier-free guidance 两种口径),rFID 指重建任务的 FID,衡量 VAE 压缩-还原的保真度。此外论文还用 IS、Precision/Recall、PSNR、SSIM、LPIPS 做补充。

论文的核心叙事就是这对指标的 trade-off:提高生成损失权重时 gFID 降、rFID 升(120.60→8.06 对 0.53→4.20),而 GenFirst 的卖点是在两组指标上同时占优。

REPA 与 REPA-E

REPA 用 DINOv2 等预训练视觉特征对齐扩散模型中间层隐藏态,加速扩散训练;REPA-E 进一步联训 VAE,但发现生成梯度回传会引发坍缩,于是切断扩散梯度、改用表征对齐损失间接重塑 latent 空间。

它们是本文最重要的对照组与出发点:本文证明不需要切断梯度这种回避手段,显式熵项加 GenFirst 调度即可让生成损失直接塑造 latent,并取得更优结果。

研究动机

主流高分辨率图像生成几乎都采用两阶段范式:先用 VAE 把像素压到低维潜在空间,再在冻结 latent 上训练生成模型。但 VA-VAE、REPA、RAE 等近期工作一致发现,为重构优化的 latent 并不适合生成——重构倾向保留实例级细节,使 latent 分布过度分散、难以拟合采样。社区为此多方尝试:REPA-E 尝试把扩散模型与 VAE 端到端联训,却发现生成梯度回传到编码器会引发 latent collapse(论文实测 EAR 设定下 gFID 飙到 294.90、IS 跌到 1.48,SiT 设定下训练直接 NaN);JetFormer、FARMER 干脆改用可逆 normalizing flow 替代 VAE;SimFlow 则要固定 VAE 方差才敢联训。可见直接端到端联训是否可行悬而未决。而即便不坍缩,重构与生成对 latent 的要求也相互冲突:论文实测 $\lambda_{prior}$ 从 0 升到 1.0,gFID 从 120.60 降到 8.06,但 rFID 从 0.53 恶化到 4.20、PSNR 从 26.18 降到 19.66,如何权衡并无定论。

本文的目标是本文想回答两个根本问题。第一,直接端到端联训 VAE 与生成模型到底可不可行?具体地,要搞清 REPA-E 观察到的 latent collapse 究竟因何而起,并给出可诊断、可修复的训练目标,使生成损失能安全地直接塑造编码器,而不是像 REPA-E 那样切断梯度绕开问题。第二,在消除坍缩之后,如何缓解生成–重构冲突,让同一个 latent 空间既容易被生成模型拟合、又能高保真解码?作者的目标不是再造一个新的对齐损失,而是给出一个简单、通用、可跨先验类型(精确似然的连续自回归、隐式似然的流匹配)与跨模态(类条件生成、文生图、统一文本-图像建模)使用的端到端训练配方,并在 ImageNet 256/512 与 GenEval/DPG-Bench 等基准上拿到有竞争力的数字(如 ImageNet 256 引导 gFID 0.988、GenEval 0.90)。

与已有工作不同的是,独特之处在于分析视角。其一,作者回到 VAE 目标函数本身做机制诊断:把 KL 散度拆成先验拟合项与负后验熵项,$\mathcal{D}_{KL}(q_\phi(z|x)\,\|\,p_0(z))=\mathbb{E}_{q_\phi}[-\log p_0(z)]-\mathcal{H}(q_\phi(z|x))$,指出 LDM 以来沿用的极小 KL 权重($10^{-6}$)只是工程正则;一旦加入生成损失,等于额外注入强大的先验拟合力($\lambda_{prior}+\lambda_{KL}\gg\lambda_{KL}$),而熵项仍被小权重压制——先验–熵失衡才是坍缩根因。这把 collapse 从一个经验性失败模式变成目标函数层面可诊断、可校正的失衡,进而给出最小修复:显式保留后验熵。其二,作者把生成–重构冲突从损失加权问题重新定义为优化顺序问题:实验显示重构是快而强监督的映射学习、生成是慢的分布建模,因此关键不是给两个损失怎么配平(constant/cosine/PI 调权全部失败),而是谁先塑造 latent 空间——这正是 GenFirst 的由来。

核心方法

直觉上,端到端联训要同时解决两件事:别让 latent 坍缩成常数,也别让收敛更快的重构损失把编码器锚死在生成模型难学的几何上。技术路线分三层。第一层是稳定目标函数:显式写成 $\mathcal{L}_{E2E}=\lambda_{rec}\mathcal{L}_{rec}+\lambda_{LPIPS}\mathcal{L}_{LPIPS}+\lambda_{GAN}\mathcal{L}_{GAN}+\lambda_{prior}\mathcal{L}_{prior}-\lambda_{ent}\mathcal{H}(q_\phi(z|x))$,把后验熵独立出来对抗先验拟合;Figure 3 显示熵项抬高 $\mathbb{E}[\sigma^2]$、增加后验重叠,防止坍缩。第二层是 GenFirst 调度:先生成阶段用较大 $\lambda^{(1)}_{prior}$ 让生成目标主导塑造 latent 几何,重构强化阶段降为 $\lambda^{(2)}_{prior}$ 让收敛快的重构补视觉细节,全程绑定 $\lambda_{ent}=\lambda_{prior}$ 使防坍缩力与先验压力同步。第三层是两类先验实例化:EAR 用 FARMER 连续自回归先验(GMM 头、精确 NLL)配 MAR VAE;EiT 用 SiT/MMDiT 流匹配先验(隐式似然),验证配方普适性。

核心创新有二。其一,熵项的重新定位:以往 KL 里的熵被视为附属正则,本文证明它是防坍缩的主力——重构与先验拟合都在收缩后验,只有 $-\mathcal{H}(q_\phi(z|x))$ 提供反向力、保住非退化后验不确定性;附录 B 进一步揭示捷径机制:当 $z\approx c$ 时 $z_t\approx(1-t)\epsilon_t+tc$、$v\approx c-z_t/(1-t)$,流匹配损失可被平凡解压低,解释了朴素端到端为何表面在训练、实际已坍缩。其二,GenFirst 与已有方法的本质区别在于用「顺序」而非「权重」解决冲突:REPA-E 靠切断扩散梯度间接塑形,SimFlow 靠固定方差;本文证明 constant/cosine/PI 加权策略都无法同时优化两端(Table 2:constant gFID 3.04/rFID 3.27,cosine 5.45/0.66,PI 发散,GenFirst 2.10/1.26),唯有按不对称动力学安排先后——先让慢的生成塑造空间、再让快的重构补细节——才能把 trade-off 推到更优位置。

方法步骤详情

以 EAR 为例:图像 $x$ 经 MAR VAE 编码器得后验参数,重参数化采样 $z=\mu_\phi(x)+\sigma_\phi(x)\odot\epsilon$(f16d16 配置);解码器重建 $\hat{x}=D_\psi(z)$,施加 $\ell_1$/LPIPS/GAN 重构损失(权重沿用标准 VAE 配方);latent 展平为 $N=hw$ 个连续 token,因果 Transformer 逐位置预测 $K$ 分量高斯混合密度并最小化精确 NLL——该损失同时更新 AR 先验与编码器;GenFirst 调度:阶段一 500 epochs 取 $\lambda_{prior}=\lambda_{ent}=1.0$,阶段二 140 epochs 降为 0.25(batch 512、lr warmup 至 $10^{-4}$)。EiT 用流匹配目标 $\mathcal{L}=\|v_\theta(z_t,t)-(z-\epsilon_t)\|_2^2$($z_t=(1-t)\epsilon_t+tz$),额外保留 $10^{-6}$ KL 控制 latent 尺度;端到端阶段 80 epochs $\lambda=0.1$ 再 40 epochs 降为 0.01,随后冻结 VAE 只训 SiT 800 epochs。文生图有版本 A(SiT-XL/2 与 FLUX.1-dev VAE 在 ImageNet 联训)与版本 B(MMDiT-L/2 在文生图数据联训)两条路,之后统一训 MMDiT-XL/2:256 分辨率 200K 步、512 分辨率 80K 步、GPT-4o 数据微调 40K 步。

技术新颖性

新颖性体现在三点。第一,机制诊断而非补丁式修复:此前 REPA-E 只观察到梯度回传会坍缩并绕开它,本文第一次把坍缩归因到 KL 的先验–熵分解,并给出最小修复(显式熵项),使生成梯度得以真正回传到编码器——这是与 REPA-E 的根本分野,也解释了 Figure 1 四种范式的本质差异。第二,把冲突归因于学习动力学的不对称并用时序解决:相较于调权(constant/cosine/PI 全部失败或次优)或堆容量(latent 维度 16→128 只能容忍 $\lambda_{prior}\approx0.1$;加深 VAE 反而伤重构;放大先验改善生成但恶化重构,见 Figure 4),GenFirst 用同一目标函数、只改两阶段权重就同时改善了 gFID 与 rFID。第三,通用性验证扎实:同一配方同时适用于精确似然 AR(EAR)与隐式似然流匹配(EiT)、类条件与文生图,还外推到表征学习(SigLIP/VLM-NLL 监督共享 latent)与统一文本-图像扩散(连文本 VAE 也被生成目标重塑),说明这不是某个架构的特调技巧,而是潜在生成建模层面的一般训练原理。

Comparison of latent generative modeling paradigms
Figure 1: Comparison of latent generative modeling paradigms
Qualitative effects of each objective on the latent posterior
Figure 3: Qualitative effects of each objective on the latent posterior
Further extension experiments
Figure 9: Further extension experiments

实验结果

实验层层递进。机制验证(Table 1):朴素端到端在 EAR 上 gFID 294.90、IS 1.48,SiT 上直接 NaN;加熵项与 GenFirst 后 EAR 达 gFID 5.10;冻结端到端 VAE 重训新先验,EAR 从 36.33 提到 5.67、SiT 从 7.90 提到 3.57,证明 latent 本身变得生成友好。调度消融(Table 2):GenFirst 引导 gFID 2.10、IS 246.70、rFID 1.26,全面优于 constant(3.04)与 cosine(5.45),PI 发散。主结果(Table 6):312M 的 EAR 引导 gFID 2.10,超 1.67B 的 GIVT(2.59)、1.9B 的 FARMER(3.60)、2.8B 的 JetFormer(6.64);675M 的 EiT 480 epochs 引导 gFID 0.988、800 epochs 无引导 1.45,是首个不用 FD loss 就把 FID 压到 1.0 以下的扩散模型;ImageNet 512 上 EiT 达 1.23。文生图(Table 8):1.3B EiT GenEval 0.90、DPG 82.60(换 Qwen3-1.7B 编码器后 84.65),参数 0.1B+1.3B 超 FLUX.2-dev(0.87,24B+32B)。效率(Figure 8):100K 步 FID 6.91,比 SiT/REPA 终点少约 70×/26× 步,2M 步追平 REPA-E。扩展(Tables 9-10):VLM-NLL 监督使 linear probe 81.69 反超预训练 Qwen3-VL ViT 的 79.64;统一文本-图像联训把 GenEval 79.88→86.57、CIDEr 20.03→33.28。

Stable end-to-end training learns a generation-friendly latent space
Table 1: Stable end-to-end training learns a generation-friendly latent space
Comparison of loss balancing strategies
Table 2: Comparison of loss balancing strategies
Scaling behavior of EAR
Table 3: Scaling behavior of EAR
Comparison with REPA-E
Table 4: Comparison with REPA-E
Effect of an optional normalizing flow
Table 5: Effect of an optional normalizing flow
Class-conditional performance on ImageNet 256×256
Table 6: Class-conditional performance on ImageNet 256×256
Class-conditional performance on ImageNet 512×512
Table 7: Class-conditional performance on ImageNet 512×512
Quantitative results on GenEval and DPG-Bench benchmarks
Table 8: Quantitative results on GenEval and DPG-Bench benchmarks
Reconstruction and generation performance, and representation quality of different visual representations
Table 9: Reconstruction and generation performance, and representation quality of different visual representations
Text-to-image and image-to-text performance of frozen text/image VAEs versus end-to-end trained counterparts
Table 10: Text-to-image and image-to-text performance of frozen text/image VAEs versus end-to-end trained counterparts
Fine-grained evaluation of text-to-image generation
Table 11: Fine-grained evaluation of text-to-image generation
VAE reconstruction evaluation on ImageNet-256
Table 12: VAE reconstruction evaluation on ImageNet-256
Effect of end-to-end training and REPA
Table 14: Effect of end-to-end training and REPA
Generation–reconstruction trade-off and GenFirst schedule ablations
Figure 5: Generation–reconstruction trade-off and GenFirst schedule ablations
Class-to-image samples from EiT (row 1,2) and EAR (row 3,4) with guidance (CFG scale=4)
Figure 6: Class-to-image samples from EiT (row 1,2) and EAR (row 3,4) with guidance (CFG scale=4)
Qualitative comparison of text-to-image generation (FLUX.1-dev VAE vs. end-to-end trained VAE)
Figure 7: Qualitative comparison of text-to-image generation (FLUX.1-dev VAE vs. end-to-end trained VAE)
End-to-end training using GenFirst accelerates convergence across generation tasks
Figure 8: End-to-end training using GenFirst accelerates convergence across generation tasks
Qualitative comparison of the three latent-space settings
Figure 10: Qualitative comparison of the three latent-space settings
查看结构化数据
任务指标本文基线提升
ImageNet 256×256 类条件生成(扩散) gFID(w/ CFG) 0.988(EiT,480 epochs) REPA-E 1.12(800 epochs) 降低约 12%,且是首个 FID<1.0 的无 FD-loss 扩散模型
ImageNet 256×256 类条件生成(无引导) gFID 1.45(EiT,800 epochs) REPA 5.78 / SiT 8.61 相对 REPA 降低约 75%
ImageNet 256×256 连续自回归生成 gFID(w/ CFG) 2.10(EAR,312M) FARMER 3.60(1.9B)/ GIVT 2.59(1.67B) 降低约 42%/19%,参数量仅约 1/6~1/5
ImageNet 512×512 类条件生成 gFID 1.23(EiT,400 epochs) SimFlow+REPA-E 2.74(1.4B) 降低约 55%,且 VAE 从 256 直接迁移
文本到图像生成 GenEval 0.90(EiT,0.1B+1.3B) FLUX.2-dev 0.87(24B+32B)/ Qwen-Image 0.87(7B+20B) +0.03,参数量小一个数量级
文本到图像生成 DPG-Bench 84.65(替换为 Qwen3-1.7B 文本编码器) BAGEL-7B 84.19(带 prompt 改写) +0.46
统一文本-图像建模(图生文) CIDEr(COCO Karpathy) 33.28(文本+图像 VAE 联合端到端) 20.03(双冻结 VAE) 提升约 66%;GenEval 同步从 79.88 提至 86.57

局限与改进

作者承认三点局限:其一,GenFirst 只缓解而非消除生成–重构 trade-off,端到端 VAE 仍在重构保真度上让步(EAR 默认设定 PSNR 21.93/SSIM 0.58,明显低于冻结 VAE 的 26.18/0.72;Table 12 中 FLUX.1-dev VAE 端到端化后 PSNR 31.59→29.43/29.37);其二,EAR 受连续 AR 现有 resampling-based CFG 策略拖累,无引导 gFID 随规模下降(174M→1.4B:13.85→3.19)但引导 gFID 不随之改善;其三,文生图训练语料远小于 FLUX/Qwen-Image 等大规模模型。我的补充观察:Table 1 中 SiT 与 VAE 直接联合训练(80 epochs gFID 92.38)远不如冻结端到端 VAE 后重训(3.57),说明对隐式似然先验,「端到端」实际退化为先塑形再冻结的两段式,流程并未真正简化;GenFirst 引入阶段时长与两组权重的新超参,虽在 200-560 epochs 内不敏感,本质是用超参换权衡;$\lambda_{ent}=\lambda_{prior}$ 绑定是经验选择、缺乏最优性论证;扩展实验因算力限制训练量缩水(统一文本-图像仅用一半端到端步数、30.8% 先验数据),结论强度弱于主实验。

独立分析的弱点

独立分析出五点弱点。第一,重构保真度系统性下降(PSNR 26.18→21.93、SSIM 0.72→0.58),对图像编辑、修复、视频帧间压缩等需像素级一致性的任务不友好;可在重构强化阶段引入更细粒度感知约束或轻量残差解码器专补高频细节。第二,两阶段调度带来新超参面:阶段切换点、$\lambda^{(1)}_{prior}/\lambda^{(2)}_{prior}$ 均靠网格搜索(EAR 用 1.0/0.25、SiT 用 0.1/0.01),跨架构跨数据可能都要重调;可基于论文自提的 $\mathrm{Std}(\mu)$、$\mathbb{E}[\sigma^2]$ 统计量设计在线自动切换准则。第三,对隐式似然先验(SiT)联合优化效率低,实际采用端到端塑形加冻结重训,尚未真正一体化;可探索为流匹配先验配辅助 normalizing flow 密度头提供显式似然信号。第四,EAR 的引导机制是短板,模型越大引导收益越小(Table 3 中引导 gFID 不随规模改善),需为连续 AR 设计 CFG 之外的引导手段。第五,熵项只防方差与均值坍缩,未约束 latent 的各向同性或重尾性,更高维 latent 与其他模态上是否够用未知。

未来方向

作者明确提出的方向包括:在保持生成友好的同时进一步找回重构保真度;为连续 AR 模型开发更强的引导方法以替代现有 resampling-based CFG;扩大文生图训练语料与预算以充分释放端到端 latent 的潜力。基于本文成果还可自然延伸:把 GenFirst 推广到视频生成,处理时间维上更长的 latent 序列及其独特的不对称动力学;将共享 latent 框架(Table 9 的 SigLIP/VLM-NLL 监督路线)推向生成+表征+理解三合一的视觉基座,让同一 tokenizer 同时服务扩散模型与 VLM;在统一文本-图像模型上放大规模(当前联训仅用 MMDiT-L/2、文本 latent 每token 仅 16 维),挑战离散 tokenizer 主导的多模态统一路线;从理论上刻画先验–熵失衡引发坍缩的相变条件,给出 $\lambda_{ent}$ 的自适应闭式解或控制器,替代手工的 $\lambda_{ent}=\lambda_{prior}$ 绑定;把熵保持思想迁移到离散 latent(VQ)的端到端训练,检验是否同样存在 codebook 坍缩的对偶问题。

复现评估

复现条件总体友好,但算力门槛分化明显。信息透明度高:附录 Tables 15-17 给出 EAR/EiT/文生图的完整配置(深度宽度、GMM 分量数、两阶段 epochs 与权重、batch 512、lr $10^{-4}$ 与 10-epoch warmup、EMA 0.9999、CFG 采样温度等),附录 C 给出 PI 基线、SigLIP/VLM-NLL 监督、文本 VAE(Cola 初始化)的全部实现细节;数据全公开(ImageNet、公开文生图数据集、COCO Karpathy、GenEval/DPG-Bench),项目页为 https://tom-zgt.github.io/GenFirst/,但正文未见明确代码开源承诺,动手前需确认。算力未披露:从 EAR-256(312M、640 epochs、batch 512)与 MMDiT-XL/2 的 320K 步三阶段训练推断,完整复现文生图需大规模集群;相对地,类条件 ImageNet 的 EAR 路线参数量小、配置完全公开,是学术实验室最可行的切入点。方法改动极小(一个熵项加两段权重),在现成 REPA-E 或 FARMER 代码库上改造成本低,核心结论用 Table 1/2 的消融规模即可低成本验证。