GenFirst enables direct end-to-end training of a Variational Autoencoder and a generative model by fixing a prior-entropy imbalance that causes latent collapse, then trains generation before reconstruction to resolve their conflict, hitting gFID (generation FID) 0.97 on ImageNet 256.
If you’ve shipped a diffusion or autoregressive image model, you almost certainly used the two-stage recipe from Latent Diffusion Models: train a VAE to reconstruct pixels, freeze it, then train your generator on its latents. This wastes signal. The VAE’s latent space was optimized for pixel-perfect reconstruction, not for being easy to sample from. Prior attempts to jointly train them either collapse (the VAE stops encoding anything useful) or need indirect tricks like REPA-E, which stops the generative gradient from touching the VAE and reshapes the latents through an auxiliary alignment loss to a pretrained encoder. This paper is the first to make direct end-to-end joint training work stably, which matters if you care about how much of your compute goes into a latent geometry your generator can actually model.
The authors diagnose why naive end-to-end training collapses. The standard Variational Autoencoder loss includes a small KL divergence regularizer with weight around 10⁻⁶. That KL term decomposes into two forces: a prior-fitting term that pulls the encoder’s outputs toward a Gaussian, and an entropy term that keeps the encoder’s per-image noise from shrinking to zero. When you add a generative loss (autoregressive likelihood or flow-matching), you inject more prior-fitting pressure but leave entropy at that tiny weight. The posterior variance collapses. Every image encodes to nearly the same point. Fix: split the KL into its two pieces and give the entropy term its own weight, tied to the generative loss weight.
That prevents collapse but leaves a second problem. Reconstruction learns fast (direct pixel supervision). Generation learns slowly (it has to model a whole distribution). If reconstruction shapes the latent geometry first, it locks in a spread-out, detail-preserving space that the generator then has to chase. GenFirst flips the order:
# Stage 1: generation-first (long), λ_prior = λ_ent = high
for step in range(long):
loss = recon + lpips + gan + λ_prior*L_gen - λ_ent*H(posterior)
update(vae, generator, loss)
# Stage 2: reconstruction-refinement (short), weights reduced ~4x
for step in range(short):
loss = recon + lpips + gan + λ_prior_low*L_gen - λ_ent_low*H(posterior)
update(vae, generator, loss)
Stage 1 lets the generative objective sculpt a smooth, modelable latent geometry while entropy blocks collapse. Stage 2 lowers generative pressure so the decoder can recover fine visual detail without disturbing the geometry. They instantiate this with two generator families: a continuous autoregressive prior called EAR (following FARMER) and a flow-matching prior called EiT (using SiT or MMDiT).
The prevailing view treats VAE-plus-generator training as either a two-stage pipeline (freeze the VAE) or an unstable joint problem that needs indirect supervision like REPA-E. This paper shows the opposite. Direct end-to-end training works once you stop treating the KL term as an engineering afterthought and recognize its entropy piece as a load-bearing counter-force to the generative loss, and once you sequence generation before reconstruction to match their asymmetric learning speeds. The clearest evidence is not the headline FID but the ablation showing that constant weighting, cosine decay, and PI-adaptive weighting all fail to match GenFirst. Ordering matters more than balancing.
•
The load-bearing ablation is loss-balancing versus ordering. Under matched training budget, constant weighting gets gFID 3.04 / rFID 3.27, cosine decay gets 5.45 / 0.66, PI-adaptive collapses numerically. GenFirst gets gFID 2.10 / rFID 1.26. No fixed weighting schedule reaches the same trade-off. Optimization order is the mechanism.
•
Naive end-to-end collapses; entropy-preserving end-to-end doesn’t. With SiT, naive joint training produces NaN losses. With the entropy fix and GenFirst, training is stable, and the resulting VAE, when frozen and used to train a fresh generator, gives gFID 3.57 vs 7.90 for the same generator on the standard SD-VAE at 80 epochs. So the learned latent space really is more generation-friendly, not just a co-adapted trick.
•
Headline numbers. On ImageNet 256, EiT reaches gFID 0.97 with CFG and 1.45 without CFG at 800 epochs, which the authors flag as the first diffusion result below 1.0 that doesn’t rely on a Fréchet-distance training loss. EAR at 312M parameters reaches gFID 2.10, beating larger continuous-AR baselines like FARMER (3.60) and JetFormer (6.64).
•
Text-to-image transfer. A 1.3B MMDiT trained on top of the end-to-end VAE hits GenEval 0.90, above much larger systems like FLUX.2-dev and Qwen-Image in their reported numbers.
•
Convergence. The reshaped latent space cuts training cost roughly 2× versus REPA-E at matched final FID, and matches SiT+REPA’s final result with ~26× fewer steps.
•
Cost. Reconstruction fidelity drops modestly (PSNR ~29.4 vs ~31.6 for the original FLUX.1-dev VAE), the price for a more modelable latent.
Reach for this when you’re training a latent generative model from scratch and you control both the tokenizer and the generator. The recipe is concrete: split the KL loss into prior-fitting and entropy with separate weights, tie the entropy weight to your generative loss weight (they use λ_ent = λ_prior), and run two phases where the generative weight starts high (e.g., 0.1 for flow-matching, 1.0 for AR) and drops roughly 4-10× for a shorter refinement phase. For flow-matching priors, keep a tiny 10⁻⁶ standard KL for latent-scale control since SiT-style noise injection breaks if latent magnitudes drift. The generation-first stage should be much longer than the refinement stage. If you have a pretrained VAE (FLUX.1-dev, SD-VAE, VA-VAE), you can initialize from it rather than training from scratch.
Code and a project page are at the project site. The paper releases the training recipe across both the EAR autoregressive setting and the EiT flow-matching setting, plus extensions to shared representation-learning latents (using SigLIP or VLM-caption supervision on the same latent) and unified text-image diffusion. No dataset release; text-to-image training uses a mix of public datasets they cite.
The KL term in your VAE isn’t a regularizer knob, it’s two opposing forces; when you add a generative loss you tilt the balance and have to compensate, and even after compensating, the order in which reconstruction and generation shape the latent matters more than how you weight them.
•
The generation-reconstruction trade-off is mitigated, not eliminated. End-to-end trained VAEs still lose ~2 PSNR versus their reconstruction-only counterparts. If your product needs pixel-perfect reconstruction (medical imaging, watermark preservation), this recipe costs you.
•
The specific weight schedule (λ_prior values, stage durations) was tuned per generator family on ImageNet and text-to-image at moderate scale. The paper doesn’t show that the same numbers transfer to video, 3D, or much larger training corpora, and the authors flag that their text-to-image data is small relative to FLUX-scale runs.
•
For flow-matching priors specifically, they observe that jointly optimizing VAE+SiT throughout is worse than end-to-end training followed by freezing the VAE and retraining a fresh SiT. So “end-to-end” here often means “end-to-end to learn the latent, then two-stage anyway.” The pure joint-training story is cleaner for exact-likelihood AR than for diffusion.