Towards Scalable Pre-training of Visual Tokenizers for Generation

TL;DR

VTP (from MiniMax) names a problem and proposes a recipe. The problem: reconstruction-only VAE pre-training does not benefit from scaling — better rFID hurts downstream generation gFID, so investing more compute in the tokenizer is wasted. The recipe: jointly optimize four losses on a ViT autoencoder — image–text contrastive (CLIP), self-supervised distillation + masked image modeling (DINOv2-style), pixel reconstruction (L1 + LPIPS), and a second-stage GAN fine-tune on the decoder only. The result is the first visual tokenizer with clear scaling laws for generation: gFID improves 65.8% as pre-training FLOPs grow 10×, while a matched reconstruction-only AE plateaus immediately. The final VTP-L hits 78.2% zero-shot ImageNet accuracy and 0.36 rFID, with 4.1× faster downstream DiT convergence than VA-VAE.

Key claims

Reconstruction-only ViT autoencoder pre-training exhibits a scaling paradox: rFID improves from 2.0 to 0.5 with more compute, but gFID slightly worsens from 55.04 to 58.56 [§4.3, Fig. 4].
Adding either CLIP or SSL (DINOv2-style) to the reconstruction objective reverses this paradox — understanding, reconstruction, and generation all improve concurrently as pre-training scales [§4.3, Fig. 5].
The combined CLIP+SSL+AE objective beats either single representation auxiliary, hitting gFID 27.8 with 74.9% linear-probe accuracy at fixed compute [§4.3, Fig. 6].
Generative performance scales with tokenizer encoder size from 20M → 300M parameters (gFID 31.28 → 26.12), while a matched reconstruction-only AE stays flat at ~57 across the same range [§4.4, Fig. 7b].
Generative performance scales with tokenizer pre-training data from 100K → 100M images (gFID 47.59 → 27.45), while the reconstruction-only AE moves only 58.37 → 56.71 over the same data sweep — downstream DiT FLOPs held constant [§4.5, Fig. 7a].
ViT autoencoder f16d64 at ViT-L scale matches the CNN LDM baseline on reconstruction PSNR (31.28 vs 30.63) and gFID (53.51 vs 59.53) under a two-stage train (joint L1+LPIPS first, GAN fine-tune of decoder only second) [§4.2, Table 1].
Final VTP-L-d64: 78.2% zero-shot accuracy, 85.7% linear probe, 0.36 rFID, gFID 2.81 at 80ep without CFG — competitive with VA-VAE-d32 (gFID 4.29) and RAE-d768 (gFID 4.28) under the same generation training budget [§5, Table 2].
4.1× faster downstream DiT convergence than VA-VAE, with a higher absolute ceiling — distillation-from-foundation-models is bottlenecked by the foundation model, joint-pretraining is not [§5, Fig. 9].
A small reconstruction-loss weight produces better generation — overweighting reconstruction during joint pre-training degrades the latent space’s semantic content [§4.1, implementation details].
CLIP and SSL objectives need very different batch sizes (16k–32k vs 4k). VTP samples the SSL/recon mini-batch as a random subset of the larger CLIP batch [§3.6].

Method

VTP is a symmetric Vision Transformer encoder–decoder with a bottleneck of dimension 64 (an ablation at d=256 confirms the recipe). Four losses are applied jointly on the encoder’s bottleneck features:

Pixel reconstruction — L1 + LPIPS through a 4-layer ViT-L pixel decoder followed by a pixel-shuffle layer.
Self-supervised learning — masked image modeling + DINO-style self-distillation against an EMA teacher, following DINOv2/iBOT.
Cross-modal contrastive — CLIP-style image–text alignment against a 12-layer, 768-dim text encoder.
(Stage 2 only) GAN fine-tune on the frozen encoder + frozen quantizer, decoder-only — sidesteps GAN-vs-ViT instability noted as a key empirical problem.

QKNorm is used for stability. Data: a filtered DataComp-1B subset (277M pairs) for tokenizer pre-training, then ImageNet for downstream LightningDiT-B training (80 epochs, FID-10k, no CFG). The bottleneck dimension is 64 (matching UniTok) to balance semantic capacity with reconstruction fidelity. The CLIP batch uses the full minibatch; the SSL/recon mini-batches are random subsets — necessary because CLIP wants 16k+ while SSL/recon are best at ~4k.

Results

The scaling story is the headline. Three experiments — compute, parameters, data — all show the same shape: reconstruction-only AE plateaus immediately at ~gFID 56–58, while VTP keeps improving. The compute-scaling figure (Fig. 7) shows VTP at 10× more FLOPs reaching ~24 gFID where the AE never moves off 57 [§4.4, §4.5].

Absolute numbers at ImageNet-256, LightningDiT-B, 80ep, no CFG [Table 2]: VTP-S 5.46, VTP-B 3.88, VTP-L 2.81. Comparison points filed previously: VA-VAE-d32 4.29, RAE-d768 4.28 (under the same generation model for fair comparison), SD-VAE 7.13. The VTP-L number also wins on understanding (78.2% zero-shot, 85.7% LP) where VA-VAE/SD-VAE/UniTok are non-competitive.

Convergence speed [Fig. 9]: VTP reaches VA-VAE’s terminal gFID at 1/4 the downstream DiT steps and continues improving past it. Distillation-from-foundation-models hits a lower ceiling than from-scratch joint pre-training.

Why it’s interesting

This is the third design point in the Representation Autoencoders cluster and it disagrees with both:

RAE (Diffusion Transformers with Representation Autoencoders, Scaling Text-to-Image Diffusion Transformers with Representation Autoencoders, Improved Baselines with Representation Autoencoders) freezes a pretrained DINOv2 / SigLIP-2 encoder and retrains only the decoder, paying with high-dim latents and DiT architectural changes.
FAE (One Layer Is Enough: Adapting Pretrained Visual Encoders for Image Generation) freezes the same pretrained encoders but compresses their features down to SD-VAE-shaped latents via a single self-attention layer, recovering shape compatibility.
VTP rejects the frozen-encoder premise entirely and trains the encoder from scratch with CLIP+SSL+reconstruction simultaneously. The latent is d=64 (SD-VAE-compatible-ish), so downstream LightningDiT specs need no modification — closer to FAE’s engineering position than RAE’s. But unlike FAE, VTP doesn’t start from a foundation model — it discovers one as a side effect of joint training.

The claim worth testing against the existing wiki is the “performance ceiling” argument. VTP-L (gFID 2.81 @ 80ep, no CFG) outperforms VA-VAE-d32 (4.29 @ 80ep, no CFG) and RAE-d768 trained with the same generation model (4.28). RAE-d768 with the full AutoGuidance + DiTDH-XL recipe at 800ep hits 1.13 with CFG, so the comparison is not at the same training budget; the head-to-head between VTP-L scaled to 800ep+CFG and RAE-d768 is not in the paper.

For Jiaming’s note: the “ViT as the VAE” choice plus heavy auxiliary losses is precisely VTP’s setup. The Sand AI reference is presumably MAGI-1 — which used a ViT VAE but with a thinner loss menu. VTP’s contribution is showing that the heavier loss menu is necessary for scaling: without CLIP/SSL, both ViT and CNN AEs plateau.

This is also a counterexample to Making Reconstruction FID Predictive of Diffusion Generation FID / Image Tokenizer Needs Post-Training‘s framing — those papers try to make rFID a better predictor of gFID, while VTP’s data argues that the rFID–gFID relationship is non-monotone when you train with reconstruction only, so the right fix is to change the training objective, not the metric. The two positions might be complementary: VTP changes the objective; pFID-style metrics could then be used to evaluate the new objective.

The “seams in patches” worry from Jiaming’s note isn’t addressed mechanically in the paper, but VTP-L’s rFID of 0.36 and the 2-stage decoder-only GAN fine-tune are the relevant levers. The MAGI-1 / Sand AI lineage of ViT-VAEs has a documented seams problem; VTP’s joint LPIPS + GAN fine-tune may or may not fully resolve it — the paper reports aggregate FID, not seam-localized failures.