PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2608.29335v1cs.CV

TL;DR

Reconstruction-oriented latents can hinder generation, while naïve joint training risks latent collapse and a generation–reconstruction conflict. The paper separates prior fitting from entropy preservation and introduces GenFirst, which shapes the latent space through generation before strengthening reconstruction. GenFirst enables stable end-to-end latent learning across priors and modalities, including strong ImageNet and text-to-image results.

  • Problem

    Reconstruction-optimized latent spaces may not suit generation, while direct joint training faces latent collapse and a generation–reconstruction conflict.

  • Method

    GenFirst uses entropy-preserving end-to-end training, first allowing generation to shape the latent space under weak reconstruction pressure, then progressively strengthening reconstruction.

  • Results

    The method improves the generation–reconstruction trade-off across exact-likelihood and flow-matching priors, achieving a FID of 0.97 on ImageNet and a GenEval score of 0.90 on text-to-image generation.

  • Takeaways & Limitations

    Stable direct end-to-end latent learning generalizes across generative priors, objectives, modalities, and shared visual representations.

  • Takeaways & Limitations

    GenFirst does not completely eliminate the trade-off, and continuous autoregressive generation remains limited by the current CFG strategy.

Abstract

from arXiv · show

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.

1 Introduction

Latent generative models use compressed latent spaces, but reconstruction-optimized representations may be poorly suited to generation. GenFirst addresses collapse and the generation–reconstruction conflict by shaping the latent space through generation before strengthening reconstruction.

  • Reconstruction-optimized latent spaces may be suboptimal for generation because they preserve instance-level details and become harder to model and sample.
  • Naïve end-to-end training can cause latent collapse when generative gradients update the VAE.
  • Reconstruction improves rapidly under direct supervision, whereas generation requires longer optimization to learn a globally sampleable latent distribution.
  • GenFirst first lets the generative objective shape a generation-friendly latent space under weak reconstruction pressure, then progressively strengthens reconstruction to recover visual details.
  • GenFirst is validated with exact-likelihood continuous autoregressive and SiT flow-matching priors across image and text-to-image generation.

2 Related Work

Related work studies latent representations, representation alignment, and joint training to improve generative modeling beyond fixed reconstruction-oriented latents.

  • Latent Generative Models: Latent generative models compress images with VAEs or visual tokenizers before training generative models on the lower-dimensional latent space.
  • Representation Learning for Generation: Representation-alignment methods seek more generation-friendly representations, including REPA, VA-VAE, and REPA-E.
  • End-to-End Generative Modeling: End-to-end generative modeling jointly trains autoencoders and generative models, using approaches such as LSGM, REPA-E, JetFormer, and FARMER.

3 Revisiting End-to-End Latent Generative Modeling

The analysis attributes latent collapse to prior–entropy imbalance and finds that generation and reconstruction favor different operating points. GenFirst therefore treats their conflict as an optimization-order problem.

  • Latent Collapse: Naïve end-to-end training can degenerate the posterior and cause latent collapse without proper balancing.
  • Prior–Entropy Imbalance: Prior fitting pulls samples toward high-probability prior regions, while entropy prevents the posterior from degenerating into point masses.
  • Prior–Entropy Imbalance: A tiny KL weight leaves learned prior fitting dominant, driving posterior variance toward zero and posterior means toward nearly identical codes.
  • Prior–Entropy Imbalance: The entropy term counteracts learned prior fitting and prevents latent collapse.
  • Generation–Reconstruction Trade-off: Increasing prior weight improves gFID from 120.60 to 8.06 but worsens rFID from 0.53 to 4.20, exposing a generation–reconstruction trade-off.
  • Generation–Reconstruction Trade-off: Increasing latent dimension, VAE capacity, or prior capacity shifts or amplifies the trade-off rather than removing it.
  • GenFirst: Because reconstruction is faster and more strongly supervised than generation, GenFirst lets generation shape the latent space before strengthening reconstruction.

4 Method

The method jointly optimizes a VAE and latent generative prior, using exact-likelihood autoregressive or SiT priors and GenFirst’s staged schedule to shape latents before refining reconstruction.

  • End-to-end formulation: The unified end-to-end objective jointly optimizes the VAE and latent generative model, with the prior loss determined by the chosen generative model.The encoder samples latents from posterior parameters, and the decoder reconstructs images from those latents.
  • SiT prior: SiT trains a velocity field with Gaussian perturbations rather than an explicit likelihood, while a small KL term keeps latent scale bounded.The KL regularizer is distinct from the explicit posterior-entropy term that provides primary anti-collapse regularization.
  • Continuous autoregressive prior: EAR uses a continuous autoregressive prior with exact likelihood to model the VAE latent distribution and provide a tractable prior loss.The latent tensor is flattened into continuous tokens, whose density is factorized autoregressively and modeled with a causal Transformer.
  • Continuous autoregressive prior: During end-to-end training, the autoregressive negative log-likelihood updates both the prior and VAE encoder, reshaping latents into a distribution easier for the prior to model.A Gaussian mixture head predicts each continuous token’s conditional density from the Transformer hidden state.
  • Generation-before-reconstruction training: GenFirst separates latent-space formation from reconstruction refinement by using a larger prior weight first and a smaller one later.The generation-first stage emphasizes a generation-friendly latent space while reconstruction preserves information and entropy prevents collapse; refinement then recovers visual details.

5 Experiments

Experiments show that entropy-preserving end-to-end training with GenFirst avoids latent collapse, improves the generation–reconstruction trade-off, scales across priors and resolutions, and supports strong image generation, representation learning, and multimodal results.

  • End-to-end stability: Stable E2E avoids NaN collapse, while its learned VAE substantially improves newly trained-prior gFID from 36.33 to 5.67 for EAR and from 7.90 to 3.57 for SiT.For EAR, joint Stable E2E slightly outperforms retraining; for SiT, freezing the learned VAE and retraining the prior is more effective.
  • Trade-off analysis: Increasing prior weight improves gFID from 120.60 to 8.06 but worsens rFID from 0.53 to 4.20, confirming a generation–reconstruction trade-off.The same change raises IS from 13.07 to 113.51 while reducing PSNR from 26.18 to 19.66 and SSIM from 0.72 to 0.45.
  • Trade-off analysis: GenFirst reaches the best evaluated balance with gFID 2.10, IS 246.70, and rFID 1.26, whereas simple weighting schedules fail to consistently resolve the conflict.Constant weighting favors generation, cosine decay favors reconstruction, and PI control becomes numerically unstable.
  • Scaling: Scaling the EAR prior from 174M to 1.4B parameters reduces gFID from 13.85 to 3.19 and increases IS from 79.00 to 181.99.Increasing GMM components from 32 to 1024 improves gFID from 6.08 to 5.01 and IS from 132.70 to 144.63, with marginal gains beyond 1024 components.
  • Image generation: EiT achieves guided gFID 0.988 and unguided gFID 1.45 on ImageNet-256, while EAR reaches gFID 1.87 and EiT 1.23 on ImageNet-512.EiT also outperforms REPA-E after 80 epochs, reaching unguided gFID 2.79 versus 3.46 for REPA-E.
  • Text-to-image generation: EiT reaches GenEval 0.90 and DPG-Bench 82.60 for text-to-image generation, with DPG-Bench increasing to 84.65 using the Qwen3-1.7B text encoder.The reported strengths include compositional generation across counting, color, spatial relations, and attributes.
  • Training efficiency: GenFirst accelerates convergence: class-conditional FID reaches 6.91 at 100K steps and 4.97 at 150K steps, using at least 70× and 26× fewer steps than SiT and REPA final results.It matches REPA-E at approximately 2M steps, corresponding to a 2× reduction in training cost, and reaches FID 1.45 at 4M steps.
  • Representation and multimodal extensions: Latent-level supervision raises linear-probe accuracy from 77.71 to 79.25 with SigLIP and to 81.69 with VLM-NLL, while joint text-image VAE training reaches GenEval 86.57 and DPG-Bench 72.68.The shared-latent experiments also report CIDEr 33.28 for image-to-text generation.

6 Conclusion

The paper identifies prior–entropy imbalance and asymmetric learning dynamics as central issues in end-to-end latent training. It introduces entropy-preserving joint training and GenFirst to enable stable training and improve the generation–reconstruction trade-off.

  • Latent collapse in naive end-to-end training is mainly caused by a prior–entropy imbalance.The paper identifies this imbalance as the first key observation underlying collapse.
  • Asymmetric learning dynamics underlie the generation–reconstruction conflict.Reconstruction and generation do not optimize at the same pace, motivating a staged strategy.
  • The entropy-preserving end-to-end objective enables stable direct joint training.
  • GenFirst uses a generation-before-reconstruction strategy to mitigate the generation–reconstruction conflict.
  • The method learns more generation-friendly latent spaces and improves the generation–reconstruction trade-off.

7 Limitations

The paper reports remaining trade-offs and scope boundaries despite stable end-to-end training. Reconstruction fidelity declines moderately, continuous autoregressive guidance remains limited, and text-to-image experiments use less data than large-scale models.

  • GenFirst does not completely eliminate the generation–reconstruction trade-off.The end-to-end VAE still sacrifices some reconstruction fidelity, reflected by metrics such as PSNR.
  • EAR is limited by the current CFG strategy for continuous autoregressive models.Without CFG, scaling the EAR prior consistently improves gFID, whereas gFID with CFG does not improve accordingly.
  • The text-to-image experiments use a substantially smaller training corpus than large-scale text-to-image models.Scaling the data and training budget may further reveal the potential of end-to-end learned latent spaces.

A More Analysis

Additional analyses examine text-to-image variants, reconstruction quality, prior-weight choices, and the separate contributions of end-to-end training and REPA. They show close text-to-image variants, moderate reconstruction degradation, favorable weight effects, and complementary generation benefits.

  • Fine-grained evaluation of text-to-image generation: Version A performs slightly better overall than Version B, while both text-to-image variants remain close.Version A performs end-to-end training on ImageNet; Version B uses text-to-image training data.
  • Effect of end-to-end training on VAE reconstruction: End-to-end trained VAEs maintain strong reconstruction quality despite moderate degradation in reconstruction metrics.For FLUX.1-dev, rFID changes from 0.18 to 0.23/0.21 and PSNR from 31.59 to 29.43/29.37 for Variants A/B.
  • Effect of end-to-end training and REPA: End-to-end VAE training improves vanilla SiT-XL/2 FID from 17.20 to 8.09 even without REPA.REPA further reduces FID to 6.83 during end-to-end training and 3.57 after the prior-only phase.
  • Effect of end-to-end training and REPA: End-to-end training and REPA provide complementary rather than conflicting benefits for SiT generation.End-to-end training reshapes the latent space, while REPA further improves generation through representation alignment.

B A Shortcut Solution in Naive End-to-End Training

Naive end-to-end training can make the latent distribution highly concentrated, allowing reconstruction and denoising objectives to appear successful without learning a meaningful generative representation. Preserving posterior entropy counters this shortcut while reconstruction maintains image-dependent means.

  • Strong prior pressure with weak posterior entropy can shrink posterior variance and reduce the diversity of posterior means.In the extreme case, the latent distribution becomes highly concentrated while residual mean differences can still support reconstruction.
  • A concentrated latent distribution makes the diffusion or flow-matching denoising objective artificially easy.The model can approximately recover target velocity through an affine transformation around a nearly constant code without learning a meaningful latent distribution.
  • Preserving posterior entropy prevents posterior variance from vanishing.With non-trivial posterior noise, reconstruction also encourages posterior means to remain image-dependent and distributed across the latent space.
  • Entropy preservation and reconstruction together discourage both variance collapse and mean collapse.

C.1 Training and Inference Details

The appendix details GenFirst training configurations for EAR and EiT across ImageNet and text-to-image settings, including staged end-to-end and prior-only training. Higher-resolution experiments reuse frozen VAE latents, while selected ImageNet-256 EiT results include decoder fine-tuning.

  • EAR on ImageNet: EAR-256 jointly trains the MAR VAE and continuous AR prior from scratch with GenFirst, while EAR-512 reuses the frozen VAE and trains only a larger prior.The 512 × 512 experiment tests transfer of the GenFirst-learned 256 × 256 latent space without tokenizer adaptation.
  • EiT on ImageNet: EiT-256 jointly trains the VAE and SiT with GenFirst, then freezes the VAE for continued prior training; EiT-512 directly reuses the frozen EiT-256 VAE.The ImageNet EiT pipeline therefore separates latent shaping from subsequent prior optimization.
  • EiT on ImageNet: ImageNet-256 EiT results additionally fine-tune the VAE decoder for 10K steps, improving unguided gFID from 3.00 to 2.79.This decoder fine-tuning is used only for the ImageNet-256 EiT results in table 6.
  • Text-to-image EiT: Text-to-image EiT uses either joint VAE training with MMDiT-L/2 on text-to-image data or with SiT-XL/2 on ImageNet before freezing the VAE for target MMDiT training.The training schedule proceeds at 256 × 256, then 512 × 512, followed by OpenAI-4o fine-tuning at 512 × 512.

C.2 Details of Loss-Balancing Baselines

The appendix specifies loss-balancing baselines and implementation settings for shared visual latent experiments. It also describes stable end-to-end and discrimination-supervised variants, including latent-level SigLIP and trunk-level VLM-NLL supervision.

  • Loss-balancing baselines: All loss-balancing baselines use the same EAR architecture, dataset, and reconstruction losses, differing only in generative-objective control.This isolates schedule or controller effects from other training changes.
  • Loss-balancing baselines: The constant-weight baseline uses λprior = λent = 1 throughout 320 epochs and is reported as the strongest constant-weight comparison.Training largely converges around 260 epochs, while smaller constant prior weights produce worse trade-offs.
  • Loss-balancing baselines: Cosine decay from 0.5 to 0.1 and the tested PI controller did not improve over constant weighting, with the controller producing instability.The PI controller rapidly reduced its update factor toward the lower bound, then oscillated and caused abrupt encoder-gradient changes.
  • Shared visual latents: The end-to-end shared-latent training combines reconstruction, flow matching, entropy regularization, REPA alignment, and discrimination objectives.Trunk-level VLM-NLL supervision branches from native ViT tokens before the bottleneck latent and feeds them to a frozen Qwen3-VL language model.
  • Shared visual latents: The shortened schedule trains 200K steps with prior weight 0.1, then 100K steps with prior weight 0.01, using global batch size 256.For the discrimination variants, the discrimination-loss weight matches the prior weight.
  • Shared visual latents: The shared-latent study compares a stable end-to-end baseline, latent-level SigLIP supervision, and additional visual-language supervision settings.The stable baseline uses reconstruction, generative, entropy, and REPA objectives without explicit discrimination supervision.
  • Shared visual latents: Latent-level SigLIP applies discrimination supervision to the normalized 16 × 16 sequence of 64-dimensional tokens consumed by MMDiT.A linear stem, attention pooling, and projection map the latent tokens into the 768-dimensional CLIP image–text embedding space.

D More Qualitative Results

The appendix provides additional uncurated ImageNet-256 samples for EiT and EAR. It also summarizes the corresponding higher-resolution and text-to-image EiT configurations.

  • EiT qualitative results: Figures 11–14 show additional uncurated EiT samples on ImageNet 256 × 256 with guidance.The EiT examples correspond to 480-epoch training.
  • EAR qualitative results: Figures 15–18 show additional uncurated EAR samples on ImageNet 256 × 256 with guidance.The EAR examples correspond to 640-epoch training.
  • Configuration context: The EiT configurations distinguish ImageNet-256 latent shaping from ImageNet-512 prior training with the VAE frozen.The text-to-image configuration likewise uses an end-to-end latent-shaping phase before target MMDiT training.
Loading 2608.29335v1…