core

Reconstruction vs. Generation: Taming Optimization Dilemma in Latent Diffusion Models

Jingfeng Yao, Bin Yang, Xinggang Wang · Huazhong University of Science and Technology · 2025-01 · arXiv:2501.01423 · code

Why it matters

Names and measures the reconstruction-generation dilemma directly: a VAE tuned purely for reconstruction fidelity becomes a *harder* space for a diffusion model to learn in. Fixes it by adding a vision-foundation-model alignment loss (VF Loss) during VAE training itself, then pairs the resulting VA-VAE with an improved DiT baseline (LightningDiT) for a 21x convergence speedup and SOTA ImageNet FID.

What this paper does

read: full text

Before this

Stable Diffusion 3, Emu and Sana had pushed toward higher per-token latent dimensions to improve tokenizer reconstruction quality, following the intuition that a wider bottleneck simply carries more information.

The problem

Raising the per-token feature dimension of a VAE tokenizer reliably improves reconstruction quality but just as reliably degrades generation quality once a diffusion transformer is trained on that latent, an optimization dilemma the field had been working around rather than solving. The authors trace this to unconstrained high-dimensional latent spaces simply being harder to learn a generative model over.

The idea

Constrain the VAE's latent space during its own training by aligning it to a frozen pretrained vision foundation model, so that a higher-dimensional tokenizer can still be reconstructed well without becoming harder for a diffusion model to generate in.

How it works

A VF Loss is added to ordinary VAE training, on top of the usual reconstruction and KL terms, with no change to the VAE's architecture. It has two marginal terms compared against a frozen DINOv2's features. A marginal cosine loss penalizes patch pairs whose cosine similarity to DINOv2 falls below a 0.5 margin, and a marginal distance-matrix loss regularizes the VAE latent's pairwise relative-similarity structure to match DINOv2's, with a 0.25 margin. Both are combined with an adaptive weight set from the ratio of reconstruction-loss to VF-loss gradient magnitudes, times a manually tuned hyperparameter of 0.1. DINOv2 was found to outperform MAE, CLIP and SAM as the alignment target. VA-VAE is then paired with LightningDiT, a bundle of training and architecture upgrades to a standard DiT-XL/2, including rectified flow, logit-normal timestep sampling, larger batch size and learning rate, SwiGLU FFN, RMSNorm, rotary position embeddings, and patch size 1 instead of 2.

Evidence

Without VF loss, a 32-channel tokenizer reaches gFID 22.62 and a 64-channel one 36.83 at 160 epochs; with DINOv2 alignment these fall to 15.82 and 24.00. VA-VAE's 32-channel variant reaches reconstruction FID 0.28 against 0.49 for the 16-channel baseline. The full system, LightningDiT trained on VA-VAE, reaches gFID 2.11 at 64 epochs and the paper's headline SOTA of gFID 1.35 at 800 epochs on ImageNet 256x256, describing the 64-epoch result as an over 21x convergence speedup against the original DiT's 2.27 at 1400 epochs. Ablations attribute the gain specifically to VF loss producing a more uniform latent distribution, measured by lower density coefficient of variation and Gini coefficient, which correlates with lower gFID.

Limitations

The authors concede the margins and the 0.1 hyperparameter were manually tuned and that different foundation models likely need different margins, with no theoretical account of why DINOv2 works best. All experiments are class-conditional ImageNet at 256x256; the paper does not validate the approach at higher resolution or in text-to-image settings. VF loss gives little benefit at the standard 16-channel dimensionality, since the problem it targets only becomes severe at higher dimensions.

Why it matters

It relocates the fix for the reconstruction-generation dilemma from the diffusion model back into the tokenizer's own training, showing the dilemma is not intrinsic to high-dimensional latents but to leaving them unconstrained. This is the move REPA-E and MAETok build on directly, turning representation alignment from an auxiliary loss on the diffusion side into a property the latent space is built with.

Abstract, in the authors' own words

Latent diffusion models with Transformer architectures excel at generating high-fidelity images. However, recent studies reveal an optimization dilemma in this two-stage design: while increasing the per-token feature dimension in visual tokenizers improves reconstruction quality, it requires substantially larger diffusion models and more training iterations to achieve comparable generation performance. Consequently, existing systems often settle for sub-optimal solutions, either producing visual artifacts due to information loss within tokenizers or failing to converge fully due to expensive computation costs. We argue that this dilemma stems from the inherent difficulty in learning unconstrained high-dimensional latent spaces. To address this, we propose aligning the latent space with pre-trained vision foundation models when training the visual tokenizers. Our proposed VA-VAE (Vision foundation model Aligned Variational AutoEncoder) significantly expands the reconstruction-generation frontier of latent diffusion models, enabling faster convergence of Diffusion Transformers (DiT) in high-dimensional latent spaces. To exploit the full potential of VA-VAE, we build an enhanced DiT baseline with improved training strategies and architecture designs, termed LightningDiT. The integrated system achieves state-of-the-art (SOTA) performance on ImageNet 256x256 generation with an FID score of 1.35 while demonstrating remarkable training efficiency by reaching an FID score of 2.11 in just 64 epochs--representing an over 21 times convergence speedup compared to the original DiT. Models and codes are available at: https://github.com/hustvl/LightningDiT.

Research line

Representation-aligned latentsascendant

Reported results

Benchmark	Value	Guidance	Budget	Source
ImageNet 256x256 gFID	2.17	none	800 epochs	Table 3
ImageNet 256x256 gFID	1.35	cfg-interval + timestep shift (scale not stated)	800 epochs	Table 3
ImageNet 256x256 rFID	0.28	—	—	Table 3; independently confirmed by MAETok Table 4 and REPA-E

The widely quoted '21x faster than DiT' refers to the 64-epoch checkpoint (gFID 2.11), not to this 800-epoch headline number. The speed claim and the quality claim come from different rows.

The strongest cross-paper agreement in this table: three independent papers report the same value.

Design-axis choices

Representation

VAE latent

Training signal

Representation alignment (REPA-style)

Problem

A single latent space optimized for pixel-reconstruction fidelity is often a poor space for a generative prior to learn in, and pushing on one side tends to hurt the other.

Method note — the shared flow-matching interpolation

Every flow-matching / rectified-flow paper in this atlas trains toward a straight-line path between a noise sample x₀ and a data sample x₁:

x_t = (1-t)\,x_0 + t\,x_1, \qquad v_\theta(x_t, t) \approx x_1 - x_0

Builds on

builds_on Representation Alignment for Generation: Training Diffusion Transformers Is Easier Than You Think — VA-VAE's VF Loss extends REPA's alignment idea from the diffusion transformer into VAE training itself.

Built on by

REPA-E: Unlocking VAE for End-to-End Tuning with Latent Diffusion Transformers improves this — REPA-E goes end-to-end where VA-VAE trains the aligned VAE and the DiT as separate stages.
Diffusion Transformers with Representation Autoencoders builds_on this — Both attack the reconstruction-generation dilemma by aligning the encoder to a foundation model; RAE goes further and freezes the encoder entirely instead of just aligning it.