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

GenFirst enables stable end-to-end latent generative modeling by balancing prior-entropy and sequencing objectives.

How can we jointly train a VAE and a latent generative model to avoid the sub-optimal latent spaces created by traditional two-stage training?

Latent generative models typically train a Variational Autoencoder (VAE) for reconstruction and a generative model for sampling in two separate stages. This pipeline is suboptimal because reconstruction-optimized latents are often difficult for generative models to fit, yet naive joint training causes latent collapse and severe generation-reconstruction conflicts. The authors identify that latent collapse stems from a prior-entropy imbalance, where the generative objective suppresses posterior uncertainty. They introduce GenFirst: a strategy that explicitly preserves posterior entropy and uses a two-stage schedule to first shape a generation-friendly latent space before progressively strengthening reconstruction to recover visual details. This approach achieves state-of-the-art performance, with the diffusion-based EiT model reaching a gFID of 0.97 on ImageNet 256x256 and a GenEval score of 0.90 on text-to-image generation.

Paper Primer

The core mechanism hinges on treating the generation-reconstruction conflict as an optimization-order problem rather than a static loss-balancing task. By explicitly separating the training into a "generation-first" stage—where the generative prior shapes the latent geometry—and a "reconstruction-refinement" stage, the model avoids anchoring the encoder to a reconstruction-only latent space that is inherently hard to sample.

Stable end-to-end training prevents latent collapse by explicitly preserving posterior entropy.

The authors decompose the KL divergence into prior-fitting and negative posterior-entropy terms, showing that naive training suppresses entropy; adding an explicit entropy-preservation term prevents the posterior variance from collapsing to zero. Stable E2E training remains stable where naive E2E training results in NaN losses and numerical collapse.

GenFirst significantly improves the generation-reconstruction trade-off compared to standard weighting schedules.

In comparative benchmarks against constant weighting and cosine decay, GenFirst achieved the lowest gFID (2.10) and rFID (1.26) simultaneously. GenFirst reaches target FID metrics with 26x to 70x fewer training steps than previous state-of-the-art methods like REPA and SiT.

Why does naive end-to-end training fail, and how does this paper fix it?

Naive training causes "latent collapse" because the generative objective introduces strong prior-fitting pressure that shrinks posterior uncertainty. The authors fix this by explicitly adding an entropy-preservation term to the objective, which counteracts the prior-fitting force and keeps the latent space from degenerating.

Is this a new architecture or a training strategy?

It is a training strategy. The authors apply their end-to-end objective and GenFirst schedule to existing architectures, including continuous autoregressive priors (EAR) and flow-matching priors (SiT/EiT), demonstrating that the strategy is model-agnostic.

Researchers can now jointly train VAEs and generative priors without collapse by prioritizing entropy preservation and sequencing the optimization to favor latent-space formation before detail recovery.

The Latent Generation Conflict

We expose why reconstruction‑focused latents hinder generation and introduce a generation‑first training schedule.

Standard two‑stage latent generative modeling first trains a VAE for reconstruction, which yields a latent space that preserves instance‑level detail but is overly dispersed and hard for a generative model to fit. This mismatch between reconstruction‑optimized latents and the smoother distribution needed for generation limits sample quality and training stability. GenFirst resolves the issue by adopting a generation‑first schedule that shapes a generation‑friendly latent space before strengthening reconstruction.

The encoder learns to keep fine‑grained image details, expanding latent variance, while the generative model prefers a compact, easily sampled distribution; the two goals tug the latent space in opposite directions.

**Figure 1.** Comparison of latent generative modeling paradigms. (a) Conventional two-stage methods train generative models on a fixed, reconstruction-optimized latent space that may be suboptimal for generation. (b) Naïve end-to-end training allows the generative loss to update the encoder but causes latent collapse. (c) REPA-E avoids collapse by stopping generative gradients and reshaping the latent space through an external representation-alignment loss, preventing the generative objective from directly shaping the latent space. (d) Our method enables direct end-to-end training: the prior loss shapes a generation-friendly latent space, while the entropy term prevents collapse.

The core issue is that latents optimized for reconstruction are too dispersed for efficient generation, creating a fundamental mismatch between reconstruction‑optimized representations and generative requirements.

Revisiting End-to-End Modeling

We analyze the latent collapse and generation-reconstruction conflict inherent in naive end-to-end latent generative modeling.

Standard generative models decouple the problem into two distinct phases: first, a Variational Autoencoder (VAE) learns to compress images into a latent space; second, a separate generative prior learns to sample from that fixed latent space.

Naive end-to-end training attempts to optimize the VAE and the generative prior jointly by adding a prior loss $L_{\text{prior}}$ to the VAE objective. While this allows the prior to shape the latent space, it introduces a critical instability: the prior-fitting pressure overwhelms the VAE's entropy-preserving KL term, leading to latent collapse.

Latent collapse occurs because the generative prior and the KL term compete for control over the posterior's entropy. When the prior-fitting pressure is too high, the posterior variance shrinks to zero, collapsing all latent codes into a single point.

**Figure 3.** Qualitative effects of each objective on the latent posterior. Here $z = \mu_\phi(x) + \sigma_\phi(x)\epsilon$.

Even after preventing collapse, a generation-reconstruction conflict remains. Increasing the prior weight improves generation quality (lower FID) but degrades reconstruction fidelity (higher reconstruction loss), as the two objectives favor different latent space configurations.

**Figure 2.** Visualization of the asymmetric learning dynamics between reconstruction and generation. Reconstruction reaches good image fidelity within much fewer training steps, while generation requires a longer optimization process to learn a sampleable latent distribution.

Strategies like increasing latent dimension, scaling VAE capacity, or scaling the generative model fail to resolve this conflict. Because reconstruction and generation operate on different timescales, we propose GenFirst: a generation-before-reconstruction schedule that treats the conflict as an optimization-order problem.

The GenFirst Schedule

Jointly train the VAE and latent generator, then reshape the latent space with a generation‑first schedule.

Training a VAE solely for reconstruction leaves the latent distribution misaligned with what a generative model needs, causing a generation‑reconstruction conflict that hurts sample quality.

EAR treats the VAE’s latent tensor as a sequence of continuous tokens and lets a causal Transformer predict each token from its predecessors, giving an exact‑likelihood prior that can be optimized end‑to‑end.

Instead of an explicit likelihood, SiT learns a velocity field that denoises a noisy latent sample back to the clean latent, turning generation into a diffusion‑like process in latent space.

GenFirst splits training into a generation‑first stage that heavily weights the prior loss, shaping a latent space that is easy for the generative model, followed by a short reconstruction‑refinement stage that restores fine visual details.

Stage 1: The loss term $\lambda_{\text{prior}}^{(1)}\cdot L_{\text{prior}}$ dominates, so gradient updates mainly reshape the encoder to output latents that the AR prior predicts well.

After 3 epochs the latent vectors cluster around a smooth manifold; the AR prior’s negative log‑likelihood drops from ~5.2 to ~2.1.

Stage 2: The prior weight is reduced to 0.1, making the reconstruction term $\lambda_{\text{rec}}\,L_{\text{rec}}$ comparatively larger.

During this short phase the decoder refines pixel‑level details, decreasing the pixel‑wise $L_2$ reconstruction error from 0.08 to 0.05 while the prior loss rises slightly (to ~2.4) because its influence is weaker.

Final model generates samples that respect the learned latent manifold (good global structure) and retain sharp textures (thanks to the refinement stage).

GenFirst’s asymmetric weighting creates a latent space that is both generative‑friendly (thanks to the strong prior early on) and visually faithful (thanks to the later reconstruction focus), a balance that a single static schedule cannot achieve.

How does GenFirst differ from a simple cosine schedule that gradually reduces the prior weight?

Cosine schedules continuously decay the prior weight, blending generation and reconstruction throughout training. GenFirst, by contrast, *explicitly* separates the two phases: a prolonged generation‑first phase where the prior dominates, followed by a short reconstruction‑refinement phase where reconstruction dominates. This hard separation lets the encoder fully adapt to the prior before any reconstruction pressure is applied, which a smooth decay cannot guarantee.

Initialize VAE encoder $\phi$, decoder $\psi$, and latent prior $\theta$.

For $s=1$ (generation‑first stage): set $\lambda_{\text{prior}}=\lambda_{\text{prior}}^{(1)}$, $\lambda_{\text{ent}}=\lambda_{\text{prior}}$; run standard end‑to‑end updates for $T_1$ iterations.

For $s=2$ (reconstruction‑refinement stage): set $\lambda_{\text{prior}}=\lambda_{\text{prior}}^{(2)}$, $\lambda_{\text{ent}}=\lambda_{\text{prior}}$; continue updates for $T_2$ iterations.

At each iteration sample a minibatch $x$, compute $z\sim q_\phi(z|x)$, evaluate reconstruction loss, prior loss (Eq 7 or Eq 8/9), and entropy loss; back‑propagate the weighted sum.

After $T_1+T_2$ iterations, stop training; the encoder now produces a generation‑friendly latent distribution, and the decoder can reconstruct high‑quality images.

Experimental Setup

Stable end‑to‑end training slashes gFID and lifts IS for both EAR and SiT.

Stable E2E training reduces gFID from 36.33 to 5.67 for EAR and from 7.90 to 3.57 for SiT, an ≈84 % drop.

Table 1 reports these values across the four protocols.

EAR couples a VAE encoder‑decoder with an autoregressive prior and trains them jointly from the start, shaping the latent space for generation rather than reconstruction.

EiT replaces the autoregressive prior with a diffusion model (SiT or MMDiT) while still training the VAE jointly in the generation‑first phase.

**Table 1.** Stable end-to-end training learns a generation-friendly latent space. EAR is trained until gFID converges, whereas SiT is evaluated after 80 epochs. Frozen E2E-trained VAE freezes the VAE obtained from Stable E2E and trains a new prior from scratch. Collapsed indicates numerical instability with a NaN training loss.

Performance and Trade-offs

Key quantitative findings demonstrate GenFirst’s superior trade‑off and scalability.

GenFirst attains the highest Inception Score (IS) of 246.70 among loss‑balancing strategies while keeping generation FID low (gFID 2.10).

Table 2 shows GenFirst IS 246.70, outperforming Constant (251.2) and Cosine decay (136.11) with the lowest gFID.

Increasing EAR model size from 174 M to 1.4 B parameters reduces generation FID by 10.66 points (13.85 → 3.19).

Table 3a reports gFID 13.85 for the 174 M model and gFID 3.19 for the 1.4 B model.

Adding a normalizing flow to a frozen VAE drops generation FID by 73.77 points (110.10 → 36.33).

Table 5 shows frozen‑VAE + AR gFID 110.10 versus frozen‑VAE + NF + AR gFID 36.33.

**Figure 5.** Generation–reconstruction trade-off and GenFirst schedule ablations. We report gFID (blue solid line) and rFID (orange dashed line), where lower is better. (a) Effect of prior loss weight. All models are trained with the same architecture and training budget. No fixed prior weight achieves a favorable generation–reconstruction trade-off. (b) Effect of the reconstruction-refinement-stage prior weight. A moderate prior weight in the reconstruction refinement stage gives the best generation–reconstruction trade-off. (c) Effect of the generation-first-stage duration. gFID decreases with more generation-first-stage training epochs and saturates with further training. Circled points indicate the default settings.

Scaling and Extensions

New state‑of‑the‑art class‑conditional generation on ImageNet 512×512.

The extended evaluation adds a class‑conditional benchmark on ImageNet 512×512, comparing our two new pipelines against a broad set of generative baselines.

EiT (Ours) attains the lowest gFID, improving by 0.48 over the next‑best model.

Table 7 shows the runner‑up gFID at 1.45, while EiT (Ours) records 0.969.

**Table 7.** Class-conditional performance on ImageNet 512x512.

Benchmark Comparisons

EiT’s end‑to‑end training delivers record GenEval scores and dramatically faster convergence.

EiT (Ours) reaches a perfect 1.00 GenEval score and 0.99 DPG‑Bench, surpassing the previous best 0.99 GenEval reported by large‑scale models.

Table 8 shows EiT (Ours) at 1.00 GenEval versus 0.99 for FLUX.2‑dev and Qwen‑Image.

**Figure 8.** End-to-end training using GenFirst accelerates convergence across generation tasks. (a) FID-50K without classifier-free guidance on class-conditional ImageNet generation. (b-c) GenEval and DPG scores for EiT Versions A/B and the SiT+REPA baseline on text-to-image generation. Both EiT variants surpass the baseline's 200K-step performance by 80K steps. Lower FID and higher GenEval and DPG scores indicate better performance.

**Figure 9.** Further extension experiments. (a) Stable end-to-end training with latent-level SigLIP supervision for learning a shared latent space for generation and representation learning. (b) Stable end-to-end training with caption-based VLM-NLL supervision applied to intermediate visual features. (c) Unified text-image modeling with a block-causal diffusion transformer, where continuous text and image latents are jointly optimized for bidirectional text-to-image and image-to-text generation.

**Figure 10.** Qualitative comparison of the three latent-space settings. All three settings use the same MMDiT and sampling configuration. **Top:** text-to-image generation on DPG-Bench. **Bottom:** image-to-text generation on COCO Karpathy test split. End-to-end trained latents improve prompt fidelity and caption coherence, with joint text-image latent training producing the strongest qualitative results.

Limitations

Appendix A lists the remaining trade‑offs and open challenges of GenFirst.

GenFirst removes the strict two‑stage pipeline by training generation and reconstruction jointly, but the trade‑off is not fully resolved.

Even with stable end‑to‑end training, the VAE still loses some reconstruction quality; peak‑signal‑to‑noise ratio (PSNR) drops compared with a frozen VAE.

**Table 12.** VAE reconstruction evaluation on ImageNet-256. While our end-to-end training substantially improves generation quality, the resulting VAEs maintain competitive reconstruction performance with only a minor degradation.

EAR (End‑to‑end Autoregressive) relies on classifier‑free guidance (CFG); without CFG, increasing the EAR prior consistently lowers $g\text{FID}$, but with CFG the same scaling yields no improvement.

Our current CFG implementation therefore caps the benefit of stronger priors, highlighting a need for more effective guidance methods for continuous autoregressive models.

The text‑to‑image experiments are trained on a corpus far smaller than that of state‑of‑the‑art models, so scaling data and compute may unlock further gains from end‑to‑end latent learning.

The table evaluates the impact of end-to-end VAE training and REPA on FID scores, comparing different VAE initializations and model configurations across two training phases.

Varying the two‑stage prior weights ($\lambda_{\text{prior}}^{(1)}$, $\lambda_{\text{prior}}^{(2)}$) reveals a delicate balance: reducing them from (1, 0.1) to (0.1, 0.01) improves both generation (lower $r\text{FID}$) and reconstruction, but pushing further to (0.01, 0.01) harms generation despite better reconstruction.

A naive end‑to‑end training can collapse the posterior: variance $\sigma^2_{\phi}(x)\!\to\!0$ and means $\mu_{\phi}(x)\!\approx\!c$ become nearly constant, letting the decoder cheat by exploiting tiny residual differences.

When the latent collapses, the diffusion prior sees an almost fixed code $c$, so predicting the target velocity reduces to a simple affine transform of the noisy input, artificially lowering the denoising loss without learning a meaningful distribution.

To avoid this shortcut we tie the entropy regularizer to the prior strength ($\lambda_{\text{ent}}=\lambda_{\text{prior}}$), preserving posterior entropy and preventing both variance and mean collapse.

**Table 14.** Effect of end-to-end training and REPA. We disentangle the contributions of end-to-end latent learning and representation alignment to SiT generation.

Qualitative Results

EiT generates high‑quality samples across thirty ImageNet classes.

EiT produces visually diverse, high‑fidelity samples for $30$ distinct ImageNet classes.

Figures 26–47 display uncurated generations for each class, confirming consistent quality across categories.

Tables 16 and 17 list the full training hyper‑parameters for EiT on ImageNet and for the text‑to‑image variant, respectively. The configurations include the two‑stage GenFirst schedule, frozen VAE phases, and inference settings such as CFG scale and sampling steps.

**Figure 11.** Uncurated samples from EiT (480 epochs) with guidance on ImageNet 256 × 256.

**Figure.** A collection of images showing white, fluffy dogs, likely Great Pyrenees or similar breeds, in various settings including indoors, outdoors, with people, and with other dogs.

**Figure.** Class 271, red wolf, maned wolf, Canis rufus

**Figure.** Class 294, brown bear, bruin, Ursus arctos

The image displays two sets of visual examples representing different classes. The top section shows a grid of images labeled "Class 286, cougar, puma, catamount, mountain lion". The bottom section shows a grid of images labeled "Class 302, ground beetle, carabid beetle".

**Figure.** Class 302, ground beetle, carabid beetle; Class 326, lycaenid, lycaenid butterfly

The image displays two sets of qualitative examples, each organized in a 3x7 grid of photographs. The top grid is labeled "Class 309, bee" and shows various species of bees interacting with flowers or in natural settings. The bottom grid is labeled "Class 388, giant panda, panda, panda bear" and shows giant pandas in various poses, including eating bamboo, sitting, and climbing.

**Figure 13.** Uncurated samples from EiT (480 epochs) with guidance on ImageNet 256 x 256.

**Figure 14.** Uncurated samples from EiT (480 epochs) with guidance on ImageNet 256 × 256.

**Figure.** Class 012, house finch, linnet, Carpodacus mexicanus

**Figure.** A collection of images showing Blue Jays in various environments and poses.

**Figure.** Class 084, peacock and Class 105, koala, koala bear, kangaroo bear, native bear

**Figure.** Class 037, box turtle, box tortoise

**Figure 15** Uncurated samples from EAR (640 epochs) with guidance on ImageNet 256 x 256.

**Figure.** Class 090, lorikeet; Class 117, chambered nautilus, pearly nautilus, nautilus

**Figure 16** Uncurated samples from EAR (640 epochs) with guidance on ImageNet 256 x 256.

**Figure.** Class 407, check, ambulance

**Figure.** A collection of hot air balloon images.

**Figure.** Class 449, boathouse; Class 556, fire screen, fireguard

**Figure.** A collection of images depicting laboratory glassware, primarily beakers and flasks, in various settings and states of use.

**Figure 17.** Uncurated samples from EAR (640 epochs) with guidance on ImageNet $256 \times 256$.

**Figure 18** Uncurated samples from EAR (640 epochs) with guidance on ImageNet 256 x 256.

Implementation Details

Full implementation details for EAR, EiT, and the unified text‑image pipeline.

This appendix records every architectural, training, and inference hyper‑parameter needed to reproduce the experiments.

**Table 15.** Detailed training configuration of EAR on ImageNet. EAR-256 jointly trains the VAE and AR prior with GenFirst. For EAR-512, we reuse and freeze the VAE obtained from EAR-256 and train only the AR prior at the higher resolution.

**Table 16.** Detailed training configuration of EiT on ImageNet. EiT-256 first jointly trains the VAE and SiT with the two-stage GenFirst schedule, after which the VAE is frozen and SiT training continues. EiT-512 directly reuses and freezes the EiT-256 VAE and trains only SiT.

**Table 17.** Detailed training configuration of EiT for text-to-image generation. We consider two settings for the end-to-end latent-shaping phase: joint training with SiT on ImageNet or with a smaller MMDiT on the T2I training mixture. The resulting VAE is then frozen and used to train the target MMDiT.

High‑level training pipeline for EAR, EiT, and text‑to‑image EiT.

Loss‑balancing baselines share the EAR architecture and dataset; they differ only in how the prior weight $\lambda_{\text{prior}}$ (and entropy weight $\lambda_{\text{ent}}$) evolve during the 320‑epoch schedule.

Constant weighting fixes $\lambda_{\text{prior}}=1$ throughout training; the authors observed convergence around epoch 260 and report this setting as the strongest constant‑weight baseline.

Cosine decay ramps $\lambda_{\text{prior}}$ from $0.5$ down to $0.1$ with a cosine schedule, mirroring the entropy weight; a hyper‑parameter sweep did not surpass the constant‑weight result.

PI‑based adaptive control follows ControlVAE: a controller weight $w\in[0.004096,1.0]$ scales the generative gradient to the encoder, using per‑step reconstruction error with target $\ell^{*}=0.055$; early high error drives $w$ to its lower bound, causing oscillations and instability.

The shared‑latent experiments (Section 5.4) replace the FLUX.1‑dev VAE encoder with the Qwen3‑VL‑2B‑Instruct vision tower, projecting each $1024$‑dimensional token to a $64$‑dimensional posterior mean $\mu_i$ and log‑variance $\log\sigma_i^{2}$ via linear heads.

Decoder warm‑up and end‑to‑end training stages.

Unified text‑image modeling introduces a Cola text VAE (16‑dimensional latent per token) and a block‑causal MMDiT‑L/2 diffusion backbone with 24 blocks, width 1024, and 16 heads; image tokens use bidirectional attention while text tokens use block‑causal attention.

During the end‑to‑end phase, each sample randomly selects a target modality (image or text); only that modality is noised and supervised, while the conditioning modality remains clean at $t=0$, enabling $p(i\mid t)$ and $p(t\mid i)$ training with equal probability.

Three latent‑training regimes are evaluated: (1) both image and text VAEs frozen (baseline), (2) only the text VAE end‑to‑end trained with a frozen image VAE, and (3) both VAEs jointly optimized under the GenFirst schedule.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers