One-sentence Summary

Researchers from the University of Science and Technology of China and ByteDance Seed propose GenFirst, a generation-before-reconstruction strategy that achieves stable end-to-end latent generative modeling by leveraging the entropy term in the Kullback-Leibler divergence objective to prevent latent collapse and by addressing asymmetric learning dynamics, yielding strong image generation results including a gFID of 0.97 on ImageNet $256 \times 256$256×256 and a GenEval score of 0.90 on text-to-image generation.

Key Contributions

• Identify a prior–entropy imbalance as the key cause of latent collapse in direct end-to-end training and demonstrate that preserving posterior entropy prevents the latent space from degenerating.
• Reveal asymmetric learning dynamics where reconstruction is fast and strongly supervised while generation is slower, and propose a GenFirst strategy that first shapes a generation-friendly latent space under weak reconstruction pressure, then progressively strengthens reconstruction to recover visual details.
• Validate GenFirst with exact-likelihood continuous autoregressive priors and SiT-based flow-matching priors, achieving a gFID of 0.97 with CFG on ImageNet 256×256 and a GenEval score of 0.90 on text-to-image generation; further experiments confirm its applicability across objectives and modalities, including shared visual latents and unified text–image modeling.

Introduction

Latent generative models compress images into a low-dimensional latent space using a variational autoencoder (VAE) and then train a generative model on that latent distribution, greatly reducing modeling complexity. However, a VAE optimized solely for reconstruction often produces a latent space that is overly dispersed and hard for generative models to learn, and naive end-to-end joint training of the VAE and generative model suffers from latent collapse. The authors diagnose this collapse as a prior–entropy imbalance: the generative objective introduces a strong prior-fitting force while the small KL weight inherited from standard VAE training suppresses posterior entropy, driving latent variance to zero. They further reveal asymmetric learning dynamics between reconstruction (fast, detail-oriented) and generation (slow, distribution-shaping), which create a conflict. To address both issues, the authors introduce an entropy-preserving end-to-end objective and propose GenFirst, a generation-before-reconstruction strategy that lets the generative objective shape the latent distribution early under weak reconstruction pressure, then progressively strengthens reconstruction to recover visual details, yielding a more generation-friendly latent space and a practical recipe for stable end-to-end latent generative modeling.

Method

A standard latent generative model is typically trained in two stages: first, a Variational Autoencoder (VAE) is trained to encode images into a continuous latent space, and second, a generative prior is trained on the frozen latent representations. A natural alternative is to jointly optimize the VAE and the generative prior in an end-to-end manner. However, naive end-to-end training often leads to latent collapse. To understand this, the authors decompose the KL divergence term in the VAE objective into a prior-fitting term and a posterior entropy term. The prior-fitting term pulls latent samples toward high-probability regions of the reference prior, while the entropy term counteracts this by preventing the posterior from degenerating into point masses. In naive end-to-end training, the added generative objective introduces strong prior-fitting pressure but overlooks the entropy-preserving role of the KL term, breaking the balance and driving the posterior variance toward zero.

To resolve this, the authors explicitly separate prior fitting and posterior entropy in the unified end-to-end objective:

Here, the entropy term acts as a crucial counteracting force to the learned prior fitting, preventing latent collapse.

As illustrated above, different objectives push the latent posterior in different directions. When reconstruction dominates, the posterior means become widely distributed but the variance decreases, creating gaps in the latent space. When prior fitting dominates without sufficient entropy, both the standard deviation of the means and the average variance decrease, causing the posterior to collapse toward a nearly constant code. Adding the explicit entropy term increases the posterior variance and overlap, preventing collapse and improving sampleability.

Even with the entropy term preventing collapse, end-to-end training faces a generation-reconstruction conflict. Increasing the prior loss weight improves generation quality but degrades reconstruction fidelity, indicating that the two objectives favor different operating points. The authors investigate whether scaling the latent dimension, the VAE, or the generative model can resolve this conflict, but find that these strategies merely shift the trade-off without eliminating it. Similarly, standard loss balancing techniques like constant weighting, cosine decay, or adaptive control fail to achieve a satisfactory balance.

The failure of simple loss balancing stems from the asymmetric learning dynamics between reconstruction and generation.

As shown in the figure above, reconstruction is a fast, strongly supervised mapping problem where the decoder quickly recovers visual details from latent samples. In contrast, generation is a slower distribution-modeling problem that requires substantially longer training to learn a sampleable latent distribution. Because reconstruction reaches good image fidelity in far fewer steps than generation, treating the conflict as a pure loss-balancing problem is insufficient.

The figure above demonstrates this trade-off, showing that no fixed prior weight achieves a favorable balance between generation and reconstruction. To address this, the authors propose a generation-before-reconstruction training strategy, which treats the conflict as an optimization-order problem.

Building on this analysis, the authors jointly optimize the VAE and the latent generative model. Given an image $x$x, the encoder predicts posterior parameters from which a latent vector $z$z is sampled, and the decoder reconstructs the image. The form of the prior loss depends on the generative model used.

For an exact-likelihood continuous autoregressive prior, the authors minimize the latent negative log-likelihood:

The spatial dimensions of the latent tensor are flattened into a sequence of continuous tokens. A causal Transformer with a Gaussian mixture model head is used to predict the conditional density of each continuous token. During end-to-end training, this loss updates both the autoregressive prior and the VAE encoder, reshaping the latent space to be easier for the prior to model.

For a SiT prior, a velocity field is trained instead of an explicit likelihood. Given noise $\epsilon_t$ϵt​ and time $t$t, the interpolated latent $z_t$zt​ and velocity $v$v are constructed, and the objective is:

Since SiT perturbs $z$z with standard Gaussian noise, a small standard VAE KL regularization term is retained to keep the latent value scale bounded and prevent shortcut solutions.

The generation-before-reconstruction strategy employs a two-stage schedule based on the unified end-to-end objective. Throughout training, the reconstruction, LPIPS, and GAN weights remain unchanged, and the entropy weight is tied to the prior weight.

In the first stage, a larger prior weight is used. The strong generative objective shapes a generation-friendly latent space, while reconstruction preserves image information and posterior entropy prevents collapse. In the second stage, the prior and entropy weights are reduced. This allows reconstruction to receive greater relative influence, recovering visual details without substantially disrupting the generation-friendly latent space formed in the first stage. By explicitly separating latent-space formation from reconstruction refinement according to the asymmetric learning dynamics, this strategy successfully resolves the generation-reconstruction conflict.

Experiment

The experiments evaluate end-to-end training of VAEs with autoregressive and diffusion priors on class-conditional ImageNet and text-to-image generation. They reveal a fundamental generation-reconstruction trade-off, where stronger prior fitting improves generation but degrades reconstruction, and show that simple loss balancing strategies cannot resolve this conflict. The proposed GenFirst strategy, which first optimizes generation then refines reconstruction, learns a more generation-friendly latent space that substantially accelerates convergence and improves generation quality. Further extensions demonstrate that the end-to-end framework can jointly support representation learning and unified text-image modeling, improving both generative performance and semantic representation.

The reconstruction loss increases the spread of latent means and reduces variance, improving reconstruction but potentially harming generation. The prior loss reduces both mean spread and variance, enhancing generation until collapse occurs. The negative entropy term raises variance without affecting mean spread, preventing collapse. Reconstruction loss raises Std(μ) and lowers E[σ²], aiding reconstruction but possibly hurting generation. Negative entropy increases E[σ²] without changing Std(μ), preventing latent collapse.

Stable end-to-end training with entropy preservation and the GenFirst strategy avoids the numerical collapse seen with naive end-to-end training and learns a latent space that is substantially more generation-friendly. Freezing this end-to-end learned VAE and training a new prior from scratch largely preserves the generative benefit, confirming that the latent space itself drives the improvement. The better generation comes with a modest decline in reconstruction fidelity. Naive end-to-end training collapses with NaN loss, while stable end-to-end training remains stable. A new prior trained on the frozen end-to-end VAE improves gFID from 36.33 to 5.67 over the reconstruction-trained VAE, and from 7.90 to 3.57 for SiT.

Stable end-to-end training prevents the collapse observed with a naive end-to-end objective. Freezing the end-to-end trained VAE for SiT training leads to the strongest generation quality, improving gFID and IS relative to both the frozen original VAE and the stable end-to-end setting. Reconstruction quality remains similar across the non-collapsed end-to-end configurations and only slightly below the frozen VAE baseline. Naive end-to-end training collapses, while the stable end-to-end variant avoids collapse. The frozen end-to-end trained VAE achieves the best generative performance, with lower gFID and higher IS than the frozen VAE baseline. Reconstruction metrics for the stable end-to-end and frozen end-to-end trained VAE settings are nearly identical and close to the frozen VAE baseline.

Simple loss balancing strategies such as constant weighting and cosine decay each favor one objective at the expense of the other, failing to achieve a satisfactory generation–reconstruction balance. GenFirst, which trains the generative prior before the reconstruction decoder, yields the best generation quality among all stable methods while keeping reconstruction competitive. These results indicate that the conflict is better resolved by optimizing the order of objectives rather than by tuning loss weights. Constant weighting achieves reasonable generation but substantially degrades reconstruction. Cosine decay improves reconstruction at the cost of noticeably worse generation. PI adaptive control becomes numerically unstable and collapses during training. GenFirst attains the lowest gFID and highest IS among all stable strategies, with much stronger reconstruction than constant weighting. The frozen VAE baseline shows the best reconstruction but the poorest generation, highlighting the inherent trade-off. The findings suggest that scheduling prior weights alone cannot reconcile the two objectives; the order of optimization is critical.

EiT consistently outperforms REPA and REPA-E across all VAE initialization settings, achieving lower gFID and higher IS. The best generation quality is obtained with VA-VAE initialization, where EiT reaches a gFID of 2.79 and an IS of 175.40, surpassing the corresponding REPA-E configuration. These results demonstrate that stable end-to-end training learns a more generation-friendly latent space. EiT with VA-VAE initialization reduces gFID from 7.90 (REPA) to 2.79 and increases IS from 122.60 to 175.40. Under the same VA-VAE initialization, EiT improves over REPA-E by 0.67 in gFID and 15.6 in IS.

The experiments evaluate VAE training strategies for balancing reconstruction and generation, revealing that reconstruction loss increases latent mean spread and reduces variance, while prior loss reduces both and negative entropy raises variance to prevent collapse. Stable end-to-end training with entropy preservation and the GenFirst strategy (optimizing the prior before the decoder) avoids collapse, learns a generation-friendly latent space, and outperforms simple loss weighting, with the resulting EiT model surpassing REPA. These findings highlight that the order of optimization and entropy control are critical for reconciling the reconstruction-generation trade-off.