End-to-End Training for Unified Tokenization and Latent Denoising

Abstract: 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.

• The paper introduces UNITE, a unified framework that jointly optimizes tokenization and latent denoising to achieve near-SOTA reconstruction and generative performance.
• It employs a weight-shared Generative Encoder with a Transformer backbone, aligning tasks via dual objectives without relying on external supervision.
• Empirical results show competitive FID scores on ImageNet and high reconstruction accuracy on molecules, demonstrating efficiency and domain versatility.

End-to-End Training for Unified Tokenization and Latent Denoising: An Expert Overview

Introduction

The proliferation of latent diffusion models (LDMs) and foundation architectures across domains has foregrounded tokenization and generation as central operations. Conventionally, these tasks are decoupled: tokenization (typically via autoencoder) is trained as a distinct stage, producing a frozen latent space, upon which generative models are then fit. This compartmentalization, while convenient, precludes the possibility that generation and tokenization objectives could jointly induce more optimal latent representations. The paper "End-to-End Training for Unified Tokenization and Latent Denoising" (2603.22283) introduces UNITE—a framework that unifies these two objectives in a single-stage model, empirically substantiating the feasibility of true end-to-end training for high-fidelity reconstruction and synthesis across images and molecules. The approach obviates external supervision (pretrained encoders, e.g., DINO), relies solely on pixel-space and latent denoising losses, and yields performance matching or exceeding multi-stage pipelines.

Architectural Formulation and Training Dynamics

UNITE employs a shared-parameter Generative Encoder (GE), leveraging a Transformer backbone, which functions in two modes. In tokenization mode, it processes image patches concatenated with register tokens into compact latents; in generation/denoising mode, it transforms noisy latent inputs along a flow (via rectified flow matching) to their clean counterparts. The decoder reconstructs from these latents. Critically, both objectives (reconstruction, latent denoising) operate via weight-sharing, aligning representational geometry and allowing gradients from both tasks to sculpt the latent space in complement. Training consists of two forward passes per data point (for tokenization and denoising), optimizing for reconstruction $\ell_1$+LPIPS and flow-matching losses on the latents.

This approach departs from prior work that either (a) freezes the tokenizer, or (b) relies on guidance from pretrained, externally-supervised encoders to stabilize joint optimization—approaches that increase infrastructure complexity and entrench a sequential pipeline. Notably, UNITE is trained from scratch, without adversarial objectives. The empirical analysis demonstrates that weight-sharing, combined with stop-gradient on the clean latents for the denoising objective, yields more stable, interpretable, and efficient latent spaces compared to alternatives with separate parameters or unimpeded gradient flow through the latent variable.

Empirical Results

UNITE achieves strong quantitative performance and robust generalization:

• ImageNet 256×256 Generation: Base (UNITE-B) and large (UNITE-L) models reach FID of 2.12 and 1.73 respectively, competitive with two-stage latent diffusion (e.g., DiT-XL/2, FID 2.27), while using substantially fewer parameters, and outperforming pixel-space single-stage baselines (e.g., JiT-B/16, FID 3.66) (2603.22283).
• Reconstruction: UNITE tokenizers, trained without adversarial or external SSL objectives, achieve rFID as low as 1.01 (UNITE-B), improved to 0.51 with a lightweight GAN-finetuned decoder, surpassing pure autoencoder and SSL-dependent pipelines.
• Molecule and Crystal Generation: The method achieves 99.37% reconstruction match and 99.71% uniqueness on QM9—outperforming ADIT, which uses a separate VAE tokenizer. Comparable results are observed for crystal structure generation on MP20, demonstrating domain-agnostic applicability.
• Compute and Efficiency: Compared to approaches built atop large pretrained encoders (e.g., DINOv2, incurring $\sim$1022 model FLOPs of pretraining), UNITE eliminates such upfront overhead, matching or improving compute efficiency per target FID by requiring a single unified training run.

Analysis

Weight Sharing and Representation Alignment:

Centered Kernel Alignment (CKA) analysis reveals that, under weight sharing and with stop-gradient on the clean latent, the generator/denoiser and tokenizer pathways produce highly aligned internal representations, especially in later layers. Removing weight sharing or allowing direct gradient flow through the clean latent reduces CKA and degrades intermediate sample quality, introducing noise and instability in the denoising trajectory.

Latent Space Structuring:

Empirical entropy analysis (MDL perspective) shows that end-to-end trained, weight-shared models require no increase in functional parameter complexity (e.g., attention/MLP parameters), with adaptations focused in normalization/scale parameters for mode-specific calibration. This supports that the architecture is not overburdened by unification, instead capitalizing on inherent task compatibility.

Training Dynamics:

The joint optimization of reconstruction and denoising losses induces an “adversarial” interaction, where improvement in one objective (e.g., generative fidelity) coincides with an increase in loss for the other (denoising), until a steady state is reached. This tension ensures that the latent space remains dense in information yet robust to perturbation—a necessary condition for successful iterative generation.

Implications and Future Directions

UNITE’s results challenge entrenched sequential pipelines for high-fidelity generative modeling, demonstrating that high-quality tokenizer/generator coupling is achievable in a single-stage model. Importantly, by removing the requirement for large pretrained encoders for diffusion alignment, the method democratizes generative modeling for domains lacking such resources (e.g., scientific and low-data modalities).

Practically, this enables streamlined deployment (single training run, unified model, reduced code and parameter footprint) and facilitates transfer across domains by amortizing representation learning and synthesis. Theoretically, it recasts tokenization and generation as manifestations of conditional inference with varying observability, suggesting future research into alternative shared-objective regimes (e.g., integrating DINO/JEPA-style SSL objectives, optimizing for downstream control in robotics, or exploring iterative tokenization maps).

Potential future expansions include extending this unified approach to multimodal (e.g., vision-language) tokenization/generation, employing non-reconstruction objectives for tokenization, and leveraging the latent space’s robustness for accelerated one-step or few-step sample generation, further mitigating inference latency.

Conclusion

The UNITE framework evidences that end-to-end training for unified tokenization and latent denoising is not only viable but yields near-SOTA results in both reconstruction and sample generation, with substantial gains in simplicity, efficiency, and domain-transferability. The weight-shared Generative Encoder, trained solely via reconstruction and denoising objectives, occupies an efficient Pareto optimum, maintaining a latent space that is information-rich and generatively robust. These findings motivate further research into unified generative modeling, reducing reliance on pipeline staging and external pretraining, and facilitating extensibility to domains previously inaccessible to latent diffusion models.