VFM-VAE: Vision Foundation Models Can Be Good Tokenizers for Latent Diffusion Models

AI summary

Overview

• Research area: Computer vision and generative modeling — specifically visual tokenizers for Latent Diffusion Models (LDMs) and the use of Vision Foundation Models (VFMs) as encoders.
• Technical level: Advanced. The paper assumes familiarity with VAEs (variational autoencoders, KL divergence, reparameterization), latent diffusion, adversarial and perceptual losses, transformer/ConvNeXt decoder blocks, and representation-similarity metrics such as CKA/CKNNA.
• Scope: The paper proposes VFM-VAE, a tokenizer that uses a frozen Vision Foundation Model directly as the encoder front-end of a VAE with a newly designed multi-scale decoder, and evaluates how this tokenizer affects diffusion training on ImageNet 256×256.

What This Paper Is About

Latent diffusion models first compress images into a latent space with a visual tokenizer (usually a VAE), then train a diffusion model in that space. Recent work tries to make those latents more semantically meaningful by distilling features from Vision Foundation Models into the VAE. The authors find that distillation-based alignment actually weakens the robustness of the inherited VFM representation, so they instead keep the VFM frozen and use it directly as the tokenizer encoder, building a specialized decoder to reconstruct pixels from its semantic-rich but spatially coarse features. The goal is faster, better diffusion training with no distillation-induced representation loss.

Key Contributions

1. A frozen-VFM tokenizer. The authors propose VFM-VAE, which uses a frozen VFM as the encoder front of a VAE for latent diffusion, removing distillation-induced representation degradation while maintaining high reconstruction quality through a purpose-built decoder.
2. A recovery-oriented decoder design. The decoder introduces Multi-Scale Latent Fusion (separating a global/pooled style component from spatial components) and Progressive Resolution Reconstruction Blocks (modulated ConvNeXt blocks with ToRGB heads and per-resolution supervision) to reconstruct realistic images from semantically rich, detail-poor VFM features.
3. A new diagnostic metric, SE-CKNNA. Semantic-Equivariant CKNNA extends CKNNA by averaging alignment scores over semantic-preserving perturbations, and is used to analyze how tokenizer latents affect representation learning inside diffusion models.
4. A dual-side alignment strategy and strong ImageNet results. Aligning both the tokenizer and the diffusion model yields a gFID (without classifier-free guidance, CFG) of 1.62 at 640 epochs, with an earlier checkpoint reaching 2.22 gFID in 80 epochs, described as a 10× speedup over prior tokenizers.

Main Findings

• Distilled alignment is brittle. Measured with CKNNA, VA-VAE (aligned with DINOv2-Large) shows a relative change of −33.2% between CKNNA (0.202) and SE-CKNNA (0.135), while VFM-VAE (measured with SigLIP2-Large) shows only +1.6% (0.188 vs 0.191). SD-VAE scores 0.004 CKNNA and 0.005 SE-CKNNA.
• VFM-VAE reconstructs well with far less training data. In Table 1, VFM-VAE used 44M images versus 160M for VA-VAE and 108M for SD-VAE; it reached rFID 0.52 and rIS 214.1, versus VA-VAE's 0.30 and 213.6, and SD-VAE's 0.62 and 212.1. The paper states VFM-VAE used only 25% of the training images used by VA-VAE.
• Linear probing favors VFM-VAE. Top-1 accuracy improves from 31.9 (VA-VAE) to 43.2 (VFM-VAE); SD-VAE is 8.0.
• Better tokenizer latents produce better diffusion representations. With no explicit alignment, the VFM-VAE–based diffusion model achieves higher average and peak layer-wise CKNNA, with a peak of 0.52, above the 0.50 reference computed between SigLIP2-Large and DINOv2-Giant.
• Dual-side alignment compounds the benefit. Combining tokenizer-side alignment (VFM-VAE) with shallow-layer alignment (REG) gives high and uniform layer-wise CKNNA across depths. VFM-VAE + REG reaches gFID 2.22 (w/o CFG) at 80 epochs, versus 2.20 for SD-VAE + REG at 480 epochs; at 640 epochs it reaches 1.62 w/o CFG and 1.31 with CFG (gIS 241.6, sFID 4.63, gIS 300.2 with CFG, precision 0.78, recall 0.66, 685M parameters).
• Gains hold without diffusion-side alignment. VFM-VAE + LightningDiT outperforms the VA-VAE variant by 1.34 gFID, reaching 3.80 at 64 epochs versus VA-VAE's 5.14.
• Decoder ablations show large steps. Starting from an SD-VAE-style baseline at rFID 19.69 (the text states 19.79), adding Multi-scale Latent Fusion gives 14.35, then Modern Blocks 1.08, then Encoder Modifications 0.71 (rIS 206.8).
• Compatibility across VFMs. Table 6, at 100k steps with LightningDiT-L/1 without CFG: EVA-CLIP-Large (rFID 1.35, rIS 188.4, PSNR 19.33, gFID 4.40, gIS 146.4), DINOv2-Large (1.55, 199.8, 17.60, 4.00, 147.1), SigLIP2-Large (1.61, 178.0, 18.73, 5.59, 127.8).
• Other-domain ablations. At ImageNet-512 after 100k steps with LightningDiT-B/1 without CFG: VFM-VAE + LightningDiT-B/1 gives gFID 18.05, gIS 69.6, precision 0.78, recall 0.60 versus VA-VAE's 21.42, 55.3, 0.75, 0.60. For text-to-image with BLIP3-o (256px, 1 pretraining epoch), VFM-VAE scores 59.1 on DPG-Bench and 17.0 gFID on MJHQ-30K, versus VA-VAE's 55.4 and 23.0.

Methodology in Plain English

The authors take an off-the-shelf Vision Foundation Model (SigLIP2-Large by default, with VA-VAE aligned to DINOv2-Large for comparison) and freeze it — never updating its weights. Instead of using only its last layer, they pull features from shallow, middle, and final layers, concatenate them, and pass them through a small learnable projection that outputs the mean and log-variance of a diagonal Gaussian, from which a compact latent is sampled (the standard VAE reparameterization trick, using an f16d32 configuration to match VA-VAE).

Because VFM features are built for semantics, not pixels, the decoder is redesigned. The latent is split into a global component, obtained by global pooling and fed to every decoder block as style control, and spatial components created by reshape operations such as pixel shuffling. For a 256×256 output the decoder progresses through the stages 8→16→32→64→128→256, with spatial components injected only into the early low-resolution blocks (i ≤ 4) and high-resolution blocks (5 ≤ i ≤ 6) left to refine detail. The core block is a Modulated ConvNeXt block, where the global component passes through a learned affine transform to produce per-channel scaling that modulates the 1×1 convolution weights. Each block has a lightweight ToRGB head that emits an image at its own resolution, supervised directly against a downsampled ground-truth image.

Training combines a representation loss (KL divergence plus a VF cosine/matrix-distance loss), multi-resolution L1 reconstruction losses at each block, an adversarial loss using a DINOv2-based discriminator, and an LPIPS perceptual loss. Separately, the authors define SE-CKNNA, which averages CKNNA over semantic-preserving transformations — additive noise at strengths {0.05, 0.10, 0.15, 0.20} in the [0,1] range, scale interpolation ratios {0.25, 0.50, 0.75, 1.0}, and rotations {0°, 90°, 180°, 270°}. All tokenizer and generative model training is on ImageNet at 256×256, with reconstruction FID/IS reported on the full 50K validation set and generation metrics following the ADM setup (gFID, sFID, IS, precision, recall).

Why This Matters

• Impact on research: The paper reframes how VFMs should be used in generative pipelines — not as a distillation teacher whose representation degrades, but as a frozen encoder that retains native semantics. SE-CKNNA also offers a diagnostic that exposes brittleness that plain CKNNA misses, which is relevant to any work measuring representation alignment.
• Real-world applications:
  • Faster, cheaper training of text-to-image and class-conditional image generators, since fewer epochs and less data reach competitive fidelity.
  • High-resolution image synthesis pipelines (the paper shows ImageNet-512 results with LightningDiT-B/1).
  • Unified or multimodal generation systems, illustrated by the BLIP3-o text-to-image experiments on DPG-Bench and MJHQ-30K.
  • Latent caching and storage in production diffusion stacks, since VFM-VAE keeps standard low-dimensional latents (unlike concurrent approaches that use high-dimensional VFM features).
• Industry relevance: Because VFM-VAE keeps standard latent dimensions and projection layers, it is described as compatible with existing VAE-LDM infrastructure and avoids the prohibitive storage overhead of high-dimensional latents. Training-cost reductions of the kind measured here translate directly into compute savings for teams training diffusion models.

Future Directions

• Recovering high-frequency detail. The authors acknowledge that using a frozen VFM sacrifices some high-frequency fidelity; improving the decoder to close this gap is a natural next step.
• Simpler training objectives. Inheriting complex objectives from prior work complicates tuning, so a more unified or better-principled loss formulation is an open problem.
• Metric grounding. SE-CKNNA scores remain relative to the specific aligned VFMs used for comparison, which limits how absolute or cross-paper those scores can be.
• Broader VFM and task coverage. The paper demonstrates compatibility across EVA-CLIP-Large, DINOv2-Large, and SigLIP2-Large; extending this to more VFMs, higher resolutions, and unified text-to-image settings (only preliminary results are reported for the latter two) remains open.

Target Audience

Researchers and engineers working on diffusion-based image generation, visual tokenizer design, or representation alignment between generative models and self-supervised/weakly-supervised vision backbones. It is also useful for practitioners who train or fine-tune latent diffusion systems at scale and care about convergence speed, data efficiency, and compatibility with existing VAE-LDM pipelines. Readers without a background in VAEs, diffusion models, and representation-similarity metrics will find the method sections difficult.

Note on reported figures: where the paper's prose and tables differ, the values above follow the tables unless the text is explicitly cited (for example, the abstract reports 2.22 gFID at 80 epochs while the surrounding text refers to 64 epochs, and the ablation text states an rFID of 19.79 while Table 5 lists 19.69).

Authors’ abstract

The performance of Latent Diffusion Models (LDMs) is critically dependent on the quality of their visual tokenizers. While recent works have explored incorporating Vision Foundation Models (VFMs) into the tokenizers training via distillation, we empirically find this approach inevitably weakens the robustness of learnt representation from original VFM. In this paper, we bypass the distillation by proposing a more direct approach by leveraging the frozen VFM for the LDMs tokenizer, named VFM Variational Autoencoder (VFM-VAE).To fully exploit the potential to leverage frozen VFM for the LDMs tokenizer, we design a new decoder to reconstruct realistic images from the semantic-rich representation of VFM. With the proposed VFM-VAE, we conduct a systematic study on how the representation from different tokenizers impact the representation learning process throughout diffusion training, enabling synergistic benefits of dual-side alignment on both tokenizers and diffusion models. Our effort in tokenizer design and training strategy lead to superior performance and efficiency: our system reaches a gFID (w/o CFG) of 2.22 in merely 80 epochs (a 10$\times$ speedup over prior tokenizers). With continued training to 640 epochs, it further attains a gFID (w/o CFG) of 1.62. These results offer solid evidence for the substantial potential of VFMs to serve as visual tokenizers to accelerate the LDM training progress.