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

TL;DR

UNITE collapses the standard two-stage latent-diffusion pipeline (train tokenizer, freeze, train denoiser) into a single network — a Generative Encoder (GE) whose weights are shared between an image-tokenization forward pass and a latent-denoising forward pass, paired with a ViT decoder. The conceptual unification is that tokenization and generation are the same latent-inference problem under different conditioning regimes (strong vs weak observability). Training is single-stage with just reconstruction + flow-matching losses — no adversarial loss, no DINO/pretrained-SSL anchor, no second-stage fine-tune. Reaches FID 2.12 (Base) and 1.73 (Large) on ImageNet-256 from scratch, and explicitly frames itself as a parameter-tied counterpart to concurrent Unified Latents Unified Latents (UL): How to train your latents (separate encoder/denoiser networks) and a counterpoint to REPA/RAE recipes that require an external pretrained teacher.

Key claims

Tokenization and generation are the same latent-inference problem under different conditioning regimes: tokenization is “generation under strong observability” (the data point is fully observed), generation is the weak-observability case (only noise + class). This view motivates parameter sharing across the two modes [§1, Fig. 2].
A single Generative Encoder with shared parameters across tokenization and denoising modes is sufficient to train a SOTA-class latent diffusion model end-to-end in a single stage, without adversarial losses or any pretrained encoder [§1, Abstract].
Latents are represented as fixed register tokens; tokenization concatenates patch tokens with Gaussian-initialized register tokens and keeps the updated registers, while generation initializes the registers with noisy latents and does a second forward pass through the same network [§3.2].
A stop-gradient on the clean latent before forming the noised input is necessary in the weight-shared setting; allowing denoising gradients to flow through the clean latent yields worse results under parameter tying (the opposite of the separate-network case, where removing stop-grad helps) [§4, Fig. 5].
Even without explicit weight sharing (the separate-networks ablation), CKA shows strong per-layer feature alignment between an independently-trained encoder and denoiser — tokenization and denoising are intrinsically compatible tasks, not just made compatible by tying weights [§4, Fig. 6].
The model differentiates tokenization vs denoising modes primarily through normalization scales (final LayerNorm is load-bearing), while attention and MLP sublayers are highly reusable across modes [§4].
Parameter tying yields the best overall rFID/gFID trade-off vs the separate-networks ablation, while also halving stored parameters and using a single training job [§4, Fig. 5].
Increasing the number of flow-matching steps per reconstruction step during training monotonically improves gFID (3.33 → 2.12 with a single configuration sweep) while maintaining rFID, indicating the latent space becomes more sampleable as denoising pressure increases [§4].
Joint training has “adversarial” dynamics: denoising loss can rise while gFID/IS improve, because the encoder is being pushed toward a richer information-dense latent that is harder to denoise but produces better samples — convergence is about reaching a stable equilibrium, not minimizing both losses to zero [§3.3, Fig. 4].
ImageNet-256 numbers: FID 2.12 (Base), 1.73 (Large), achieved from scratch with no GAN loss and no pretrained SSL teacher [Abstract].
Method also applies to non-image modalities (molecules), used as a feasibility demonstration of the single-stage joint-training principle [Abstract].

Method

UNITE has exactly two networks: a ViT Generative Encoder f_θ and a ViT decoder g_φ. The GE operates in two modes that differ only in input format. Tokenization pass: image patch tokens are concatenated with K register tokens initialized from N(0, I) (the same distribution as the maximum-noise endpoint of the flow); one forward pass through f_θ updates everything via self-attention; the patch tokens are discarded and the updated registers form the clean latent z. Generation pass: z is corrupted to z_t via rectified-flow interpolation z_t = (1-t)z + tε; only z_t (no image patches) is passed through f_θ to predict z (x₀-parameterization). A final LayerNorm normalizes both the tokenizer output and the denoiser output to the same space.

Training does both passes per step. The loss is L = L_recon(decoded(z + small noise), image) + L_flow(f_θ(z_t, t), sg[z]), where L_recon is a pixel + perceptual term (no GAN loss) and sg[·] is stop-gradient on the clean latent so denoising gradients shape f_θ only through the second pass, not directly through z. Increasing the number of flow steps per reconstruction step (so denoising loss is sampled more densely) drives gFID down monotonically in the parameter-tied setting.

At inference, the GE serves as the tokenizer in one mode (image → latent) and as the iterative denoiser in the other (sample registers from noise, run the GE multiple times conditioned on noise level and class, decode the final z).

Results

ImageNet-256, from scratch, no adversarial loss, no pretrained encoder: FID 2.12 (Base model), 1.73 (Large model) [Abstract].
Flow-step scaling: at fixed model size, increasing flow-matching steps per reconstruction step drives gFID 3.33 → 2.12 (≈36% relative improvement) without degrading rFID, indicating the latent space becomes more denoiseable under increased generative pressure [§4, Fig. 5].
Parameter tying vs separate networks: weight sharing yields the best rFID/gFID Pareto trade-off; the separate-networks ablation is competitive but strictly dominated [§4, Fig. 5].
Stop-gradient ablation: under weight sharing, keeping stop-grad on the clean latent gives the best gFID/rFID; under separate networks, removing stop-grad improves gFID — i.e. the optimal gradient-flow structure depends on whether parameters are tied [§4].
CKA alignment: even with parameter tying removed, independently trained encoder and denoiser show strong per-layer feature correspondence [§4, Fig. 6].
Loss dynamics: denoising loss can rise during training while generation metrics (FID/IS) keep improving — non-monotone losses, like GAN training [§3.3, Fig. 4].

Why it’s interesting

UNITE is the third filed paper this month that attacks the joint-training-of-tokenizer-and-denoiser question, and it lands as a natural pair to Unified Latents (UL): How to train your latents (concurrent Google work; UL = separate encoder + denoiser networks without parameter tying, UNITE = same architecture but with parameters tied — UNITE explicitly identifies UL as a special case of its separate-weights ablation). Where UL frames the question as “what should the latent’s bitrate be” (KL-as-noise-floor), UNITE reframes it as “what’s the right parameter-sharing structure” — and reports that tying weights wins. It also contrasts cleanly with Latent Forcing: Reordering the Diffusion Trajectory for Pixel-Space Image Generation: Latent Forcing keeps two modalities (pixels + DINOv2 latents) in one network but uses a frozen DINO encoder; UNITE keeps tokenization + denoising in one network and learns everything from scratch with no external prior. Versus the RAE/REPA family (Improved Baselines with Representation Autoencoders, No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA), Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis), the paper takes a position closer to Self-Flow’s: external pretrained encoders are dispensable; the right move is to let the generative objective shape the encoder weights directly, just without the EMA-teacher second forward pass Self-Flow uses. Worth a head-to-head with UL at matched compute and at the FVD-1.3 video scale where UL holds SOTA.