Latent diffusion models (LDMs) enable high-fidelity synthesis by operating in learned latent spaces. However, training state-of-the-art LDMs requires complex staging: a tokenizer must be trained first, before the diffusion model can be trained in the frozen latent space. We propose UNITE - an autoencoder architecture for unified tokenization and latent diffusion. UNITE consists of a Generative Encoder that serves as both image tokenizer and latent generator via weight sharing. Our key insight is that tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes: tokenization infers latents from fully observed images, whereas generation infers them from noise together with text or class conditioning. Motivated by this, we introduce a single-stage training procedure that jointly optimizes both tasks via two forward passes through the same Generative Encoder. The shared parameters enable gradients to jointly shape the latent space, encouraging a "common latent language". Across image and molecule modalities, UNITE achieves near state of the art performance without adversarial losses or pretrained encoders (e.g., DINO), reaching FID 2.12 and 1.73 for Base and Large models on ImageNet 256 x 256. We further analyze the Generative Encoder through the lenses of representation alignment and compression. These results show that single stage joint training of tokenization & generation from scratch is feasible.
Figures & tables
Figure 1: Unified Tokenization & Generation via Generative Encoder (UNITE): We propose a single-stage architecture that unifies tokenization and generation through shared parameters. The Generative Encoder operates in two modes: (top) as a tokenizer, it processes image patches and register tokens to produce a latent representation z0 ; (bottom) as a generator/denoiser, it evolves latents along a flow-matching trajectory to synthesize z0 from Gaussian noise. The latent space is jointly shaped by recon. & generative objectives from scratch, without external supervision.
Figure 2: Tokenization and generation can be viewed as the same latent inference problem under different conditioning regimes. In tokenization, the full observation x strongly constrains the clean latent, z0∼pθ(z∣x) ; in generation, a noisy latent zt provides weaker evidence, and the same target latent z0 is recovered by denoising. This view motivates using a single shared Generative Encoder ( GEθ ) for both tokenization and latent denoising.
Figure 3: UNITE Training Pipeline uses two forward passes through the Generative Encoder: first, mapping (distilling) image patches into latent registers, and second, denoising a noised version of those latents, with weights shared across both passes. Training combines reconstruction losses with a denoising loss ∣z^~0−sg(z~0)∣ .
Figure 4: UNITE’s Training dynamics: The conflicting nature of the reconstruction and denoising objectives leads to an adversarial training behavior when trained jointly. The dotted lines (zoomed in) represent different ablations (see Appendix) based on the scale of noise added in reconstruction pathway for decoder robustness.
Figure 5: Weight-shared vs. Separate Enc-Denoiser training. UNITE uses a single Generative Encoder, sharing weights between tokenization and generation. To isolate the effect of weight sharing, we keep the rest of the end-to-end training pipeline fixed, including the stop-gradient that prevents denoising gradients from flowing into the tokenized output. Both UNITE and the separate encoder-denoiser ablation attain competitive performance, with UNITE benefiting from more denoising-to-reconstruction steps ratio during training, achieving the best overall rFID–gFID trade-off.
Figure 6: Representation alignment between tokenization and generation pathways. We measure alignment between tokenization and denoising activations using CKA and cosine similarity. Given an input image, we first record intermediate activations along the tokenization pathway, then corrupt the encoded latent and record the corresponding denoising-pathway activations. Left: both the weight-shared UNITE model and the separate encoder–denoiser ablation exhibit strong alignment, especially in later layers, indicating that tokenization and denoising are intrinsically aligned tasks . Middle: removing the stop-gradient and backpropagating denoising gradients through the latent weakens late-layer alignment , even though the denoising objective still matches the final latent target. Right: cosine similarity on the final latents decreases at lower denoising timesteps in the no-stop-gradient setting, suggesting that direct gradient backpropagation from denoising into tokenization leads to a less cleanly shared representation (see Fig. 7 for visual interpretation).
Figure 7: Analyzing the denoising trajectory. Given an input image, we first encode it into latents, corrupt the latent with noise, and then decode the denoised prediction at different noise levels (first three columns). The final column shows direct decoding of the clean latent. Although all four models achieve competitive aggregate rFID/gFID, the stop-gradient variants (first two rows)–UNITE and the separate encoder-denoiser ablation–exhibit markedly cleaner intermediate denoising trajectories, with higher PSNR to the input image across all noise levels. This result is consistent with the representation-alignment in Fig. 6 , which shows drop in alignment at final layers.
Method
Aux. Token
Params
FID ↓
IS ↑
Single-stage Frameworks
JiT-B/16
-
131M
3.66
275.1
UNITE-B (Ours)
Joint
217M
2.12
294.1
RIN
-
410M
3.42
182.0
JiT-L/16
-
459M
2.36
298.5
ADM-G
-
554M
4.59
186.7
Table 1: ImageNet 256 × 256 generation. Our approach outperforms both recent single-stage pixel baselines and standard two-stage latent diffusion frameworks by a large margin.
Tokenizer
Adv.
Pretrained Encoder
rFID ↓
With adversarial / external supervision
SD-VAE
✓
-
0.62
DC-AE-f32
✓
-
0.69
RAE
✓
DINOv2
0.58
VA-VAE
✓
DINOv2
0.28
Without adversarial or external supervision
Table 2: ImageNet 256 × 256 reconstruction. Our tokenizer achieves competitive rFID without adversarial loss (Adv.) or pretrained encoders. All UNITE rows use base backbone at 120 eps.
Figure 8: Selected samples from UNITE-XL. Generated using 50 steps with CFG. This model achieves FID 1.75.
Reconstruction
Generation
Method
Match (%)
RMSD (Å)
Valid (%)
Unique (%)
EDM ( Hoogeboom et al., 2022 )
–
–
91.9
90.7
GeoLDM ( Xu et al., 2023 )
–
–
93.8
92.9
ADiT Tokenizer ( Joshi et al., 2025 )
97.20
0.075
–
–
ADiT-S QM9-only ( Joshi et al., 2025 )
–
–
96.02
97.76
UNITE-S (Ours)
99.37
0.039
94.90
99.71
Table 3: QM9 molecule generation. UNITE-S achieves the best reconstruction accuracy (99.37% match) and uniqueness (99.71%) under single-stage training. Crystal generation results on MP20 are provided in Appendix C.2 .
Appendix figures & tables6 assets
Supplementary material from the paper’s appendix.
Appendix
ODE Solver
Class Sampling
NFE
FID ↓
IS ↑
Heun (50 steps)
Balanced
100
2.789
287.6
dopri5 (adaptive)
Balanced
∼ 108
2.735
268.1
dopri5 (adaptive)
Random
∼ 108
2.885
274.7
Appendix
Table 4: Effect of evaluation protocol on reported FID. All rows use the same UNITE-B checkpoint (240 epochs). “Balanced” denotes 50 images per class; “Random” denotes uniformly sampled class labels. NFE = number of function evaluations.
Figure 9: Uncurated, class-conditional samples on ImageNet 256×256 using UNITE-XL. We show images using CFG ω=4.0 . Each grid contains 21 randomly sampled images, demonstrating consistent quality across diverse categories including animals, objects, and scenes.
Table 6: Ablation study on ImageNet-256 generation. We systematically evaluate key design choices across architecture, normalization, training dynamics, and augmentation strategies. All ablations are done using the base backbone and trained for 120 epochs.
Component
Encoder/Denoiser
Decoder
DiT-B
DiT-L
DiT-XL
ViT-B
ViT-L
Hidden Dimension
768
1024
1152
768
1024
Layers
12
24
28
12
24
Attention Heads
12
16
16
12
16
MLP Ratio
4
4
4
4
4
Patch Size
16
16
16
-
-
Appendix
Table 7: Detailed architecture configurations for UNITE models.
Hyperparameter
Value
Hyperparameter
Value
Base Learning Rate
1×10−4
Warmup Epochs
20
Global Batch Size
1024
Total Epochs
240
Optimizer
Muon
LR Schedule
Cosine
AdamW Betas
(0.9, 0.999)
Min LR
1×10−6
Weight Decay
0
Gradient Clip
3.0
Reconstruction Noise ( σ )
0.7
EMA Decay
0.9978
Appendix
Table 8: Training configuration for ImageNet-256 experiments. All models trained with mixed precision BF16 & gradient clipping at 3.0.
Latent Diffusion Models (LDMs) typically adopt a two-stage pipeline: an auto-encoder (AE) is first pre-trained to define a latent space, then a diffusion model is trained to perform denoising within it. Such a two-stage design introduces a representation mismatch, as the latent space is optimized for reconstruction rather than adapting the denoising dynamics. We reveal that the LDM itself is an AE, and consequently present LDM-is-AE, an end-to-end one-stage LDM training framework that eliminates the need for a separately trained tokenizer. Our key observation is that the LDM backbone actually performs a latent-to-feature-to-latent transformation at each denoising step, which can be interpreted as an internal decoding--encoding process. Leveraging this structure, we split the DiT backbone into two reciprocal components, DiT-E (i.e., DiT Encoding) and DiT-D (i.e., DiT Decoding), and impose image-space supervision on the intermediate features across all timesteps. Our model encourages the internal representation to align with the image domain throughout denoising, thereby establishing an explicit latent-to-image-to-latent path. At the zero-noise timestep, our model further performs an image-to-latent-to-image mapping, corresponding to an auto-encoding process. As a result, LDM-is-AE jointly learns latent representations and denoising dynamics in an end-to-end manner, yielding a diffusion-native latent space tailored to the generation process. Experiments demonstrate that LDM-is-AE exhibits highly competitive generation performance, achieving an FID of 1.80 and 1.90 on 256x256 and 512x512 class-conditional image generation, respectively.
Zhengqiang Zhang, Lingchen Sun, Rongyuan Wu +4
The Hong Kong Polytechnic University · OPPO Research Institute
Latent diffusion models offer an attractive alternative to discrete diffusion for non-autoregressive text generation by operating on continuous text representations and denoising entire sequences in parallel. The major challenge in latent diffusion modeling is constructing a suitable latent space. In this work, we present the Latent Diffusion Language Model (LDLM), in which the latent encoder, diffusion model, and decoder are trained jointly. LDLM builds its latent space by reshaping the representations of a pre-trained language model with a trainable encoder, yielding latents that are easy to both denoise and decode into tokens. We show that naive joint training produces a low-quality diffusion model, and propose a simple training recipe consisting of an MSE decoder loss, diffusion-to-encoder warmup, adaptive timestep sampling, and decoder-input noise. Ablations show that each component substantially impacts generation performance. On OpenWebText and LM1B, LDLM achieves better generation performance than existing discrete and continuous diffusion language models while being 2-13× faster, indicating that jointly learning the latent space is a key step toward making latent diffusion competitive for text generation.
Viacheslav Meshchaninov, Alexander Shabalin, Egor Chimbulatov +4
Constructor University, Bremen, Germany · HSE University, Moscow, Russia · Applied AI Institute, Moscow, Russia +2
We present Unified Latents (UL), a framework for learning latent representations that are jointly regularized by a diffusion prior and decoded by a diffusion model. By linking the encoder's output noise to the prior's minimum noise level, we obtain a simple training objective that provides a tight upper bound on the latent bitrate. On ImageNet-512, our approach achieves competitive FID of 1.4, with high reconstruction quality (PSNR) while requiring fewer training FLOPs than models trained on Stable Diffusion latents. On Kinetics-600, we set a new state-of-the-art FVD of 1.3.