GenFirst: Generation Before Reconstruction for Stable End-to-End Latent Generative Modeling

AI summary

Overview

Research area: Computer Vision — latent generative modeling, specifically end-to-end training of variational autoencoders together with generative priors (diffusion/flow-style models and autoregressive models), with extensions to representation learning and text-to-image generation.

Technical level: Advanced.

Scope: The paper diagnoses why jointly training a latent encoder/decoder and a generative model tends to collapse, and proposes a "generation-before-reconstruction" training order that makes stable direct end-to-end latent learning possible across several priors and modalities.

What This Paper Is About

Most latent generative models are trained in two stages: first a variational autoencoder is trained to reconstruct images, then a generative model is trained on that frozen latent space. Because latents optimized purely for reconstruction are not necessarily well suited to generation, the authors argue it would be better to train both parts jointly — but naive joint training collapses the latent space and creates a conflict between the generation and reconstruction objectives.

The goal of the paper is to make direct, stable end-to-end training of the whole pipeline work, without latent collapse, and to show that the resulting recipe generalizes beyond a single model family or modality.

Key Contributions

1. A diagnostic analysis of how different objectives shape the latent space. The authors identify two insights: the entropy term in the Kullback-Leibler objective is essential for preventing collapse (both reconstruction and prior fitting tend to shrink the posterior, while entropy preserves non-degenerate latent uncertainty), and reconstruction and generation have asymmetric learning dynamics — reconstruction is fast and strongly supervised, while generation is slower and harder to optimize.

2. The first direct end-to-end training without latent collapse. Based on those insights, the authors report achieving joint training in which the latent space does not collapse, which the abstract presents as the central technical hurdle being cleared.

3. GenFirst, a simple generation-before-reconstruction strategy. Instead of optimizing reconstruction first, the generative objective shapes the latent space under weak reconstruction pressure, and reconstruction is then progressively strengthened to recover visual details.

4. Validation across priors and modalities. The recipe is tested with continuous autoregressive priors that have exact likelihoods and with SiT priors that have implicit likelihoods, and is extended to shared visual latents for both generation and representation learning, and to continuous unified text-image generation.

Main Findings

• The KL entropy term is the anti-collapse ingredient. The abstract states that reconstruction and prior fitting both tend to shrink the posterior, whereas the entropy term preserves non-degenerate latent uncertainty — making it essential rather than incidental.

• Reconstruction and generation learn at different speeds. Reconstruction is described as fast and strongly supervised; generation is slower and harder to optimize. The paper treats this asymmetry as the reason a naive joint objective fails, and as the motivation for reordering the objectives.

• Order matters: generation first, reconstruction later. The GenFirst schedule lets the generative objective shape the latent space under weak reconstruction pressure, then strengthens reconstruction progressively to recover visual detail.

• Image generation results. With the end-to-end objective plus GenFirst, SiT is reported to reach a gFID of 0.97 with CFG and 1.45 without CFG on ImageNet-256.

• Text-to-image results. MMDiT is reported to reach a GenEval score of 0.90 on text-to-image generation.

• The framework generalizes. Beyond image generation, it is extended to shared visual latents used for both generation and representation learning, and to continuous unified text-image generation, which the authors present as evidence of generality across generative priors and modalities.

Note: the abstract reports these headline numbers but does not give dataset sizes, baseline-by-baseline comparisons, or ablations, so those details are not available here.

Methodology in Plain English

The starting point is the standard two-stage recipe — compress images into a latent space with an autoencoder trained for reconstruction, freeze it, then train a generator on it. The authors instead train both parts together, which normally breaks because the two objectives pull the latent space in incompatible directions.

Their approach has two parts. First, they examine what each term in the training objective does to the latent space, and conclude that the entropy component of the KL term is what keeps the posterior from shrinking into something degenerate; without it, the combined pressure of reconstruction and prior fitting collapses the representation. Second, they note that if you train both objectives simultaneously, the fast, strongly supervised reconstruction objective dominates before the slower generative objective has a chance to organize the latent space.

GenFirst resolves that imbalance by changing the order rather than the architecture: the generative objective is optimized first, with only weak reconstruction pressure, so the latent space is laid out in a way that suits generation; reconstruction is then gradually turned up so that fine visual detail is recovered. The authors test this schedule with two different prior families — continuous autoregressive priors with exact likelihoods and SiT priors with implicit likelihoods — and then apply the same idea to shared latents for generation plus representation learning and to unified text-image generation.

Why This Matters

Impact on research: The paper reframes end-to-end latent generative training from an architecture problem into an objective-balancing and scheduling problem, and offers a concrete mechanism (the entropy term) for why collapse happens. If the claims hold, it removes a long-standing reason researchers keep the encoder frozen, and it suggests that the same ordering principle may transfer to other prior families and modalities.

Real-world applications:

• Text-to-image content creation tools, where the paper reports a GenEval score of 0.90 with MMDiT.
• High-fidelity class-conditional image synthesis, where the reported SiT gFID figures are relevant to production image pipelines.
• Representation learning, since the framework supports shared latents that serve both generation and downstream visual understanding.
• Unified multimodal generation, via the extension to continuous text-image generation.

Industry relevance: Training latent generative systems in one stage instead of two can simplify pipelines and reduce the cost and coordination of staged training, while potentially improving the match between the latent space and the generator that depends on it. A single latent space that supports both generation and representation tasks is also attractive for teams that want one backbone rather than separate models per task.

Future Directions

• Scaling the schedule beyond the tested settings. The abstract validates GenFirst on image generation, shared-latent representation learning, and text-image generation; whether the generation-before-reconstruction ordering holds for other priors, larger models, or other modalities is not addressed.

• Theory of the entropy term. The paper argues the KL entropy term prevents collapse, but the abstract does not specify how much entropy is needed, how it should be scheduled, or whether the argument generalizes beyond the objectives tested.

• Balancing shared latents. When one latent space must serve both generation and representation learning, the trade-off between the two is an open question the abstract does not resolve with numbers.

• Comparison against two-stage training. The abstract reports absolute results for the end-to-end recipe; how consistently it beats a frozen-latent two-stage baseline across settings is a question the abstract does not answer.

Target Audience

Researchers and engineers working on latent generative models — diffusion/flow-based image generators, autoregressive generative models, and variational autoencoders — as well as those interested in multimodal generation and in representation learning with generative objectives. The paper assumes familiarity with VAEs, KL regularization, latent collapse, classifier-free guidance, and FID-style evaluation, so it is best suited to readers with prior background in generative modeling rather than newcomers.

Authors’ abstract

Latent generative models typically follow a two-stage pipeline, training a variational autoencoder for reconstruction and then a generative model on the frozen latent space. Since reconstruction-optimized latents are not necessarily generation-friendly, jointly training both models is an appealing alternative. However, direct end-to-end training remains challenging, as it is prone to latent collapse and faces a generation-reconstruction conflict. We revisit this problem by analyzing how different objectives shape the latent space and identify two key insights. First, the entropy term in the Kullback-Leibler divergence objective is essential for preventing collapse: reconstruction and prior fitting tend to shrink the posterior, while entropy preserves non-degenerate latent uncertainty. Second, reconstruction and generation exhibit asymmetric learning dynamics: reconstruction is fast and strongly supervised, whereas generation is slower and harder to optimize. Based on these insights, we achieve the first direct end-to-end training without latent collapse and propose GenFirst, a simple generation-before-reconstruction strategy. The generative objective first shapes the latent space under weak reconstruction pressure, after which reconstruction is progressively strengthened to recover visual details. We validate GenFirst with continuous autoregressive priors with exact likelihoods and SiT priors with implicit likelihoods. With our end-to-end objective and GenFirst, SiT achieves a gFID of 0.97 with CFG and 1.45 without CFG on ImageNet-256, while MMDiT reaches a GenEval score of 0.90 on text-to-image generation. Beyond image generation, we extend the framework to shared visual latents for generation and representation learning, and to continuous unified text-image generation. These results demonstrate the generality of stable end-to-end latent learning across generative priors and modalities.