Robust Latent Matters: Boosting Image Generation with Sampling Error Synthesis

TL;DR

This work addresses the gap between tokenizer reconstruction metrics and autoregressive image generation quality by analyzing sampling errors in discrete latent spaces. It proposes a novel latent perturbation framework and a tokenizer-specific robustness metric, Perturbed FID ($pFID$), to guide tokenizer training and evaluation. The core contribution, RobustTok, is a plug-and-play tokenizer trained with latent perturbations and annealing, supervised by semantically meaningful signals from DINO, which yields improved $gFID$ scores and faster convergence on ImageNet 256×256 across multiple tokenizers and AR backbones. The approach offers a practical path to robust discrete latent spaces for generation, reducing the reliance on expensive generator training for tokenizer ranking and enabling more reliable, high-quality image synthesis in autoregressive settings.

Abstract

Recent image generation schemes typically capture image distribution in a pre-constructed latent space relying on a frozen image tokenizer. Though the performance of tokenizer plays an essential role to the successful generation, its current evaluation metrics (e.g. rFID) fail to precisely assess the tokenizer and correlate its performance to the generation quality (e.g. gFID). In this paper, we comprehensively analyze the reason for the discrepancy of reconstruction and generation qualities in a discrete latent space, and, from which, we propose a novel plug-and-play tokenizer training scheme to facilitate latent space construction. Specifically, a latent perturbation approach is proposed to simulate sampling noises, i.e., the unexpected tokens sampled, from the generative process. With the latent perturbation, we further propose (1) a novel tokenizer evaluation metric, i.e., pFID, which successfully correlates the tokenizer performance to generation quality and (2) a plug-and-play tokenizer training scheme, which significantly enhances the robustness of tokenizer thus boosting the generation quality and convergence speed. Extensive benchmarking are conducted with 11 advanced discrete image tokenizers with 2 autoregressive generation models to validate our approach. The tokenizer trained with our proposed latent perturbation achieve a notable 1.60 gFID with classifier-free guidance (CFG) and 3.45 gFID without CFG with a $\sim$400M generator. Code: https://github.com/lxa9867/ImageFolder.

Figures (13)

• Figure 1: (a) Traditional image generation scheme: a discrete image tokenizer is fisrt trained with reconstruction target where visual decoder is fed with clean image tokens. After that, an AR/LLM is trained with clean tokens under teacher forcing. However, the during subsequent AR prediction (inference), unexpected tokens can be sampled from the learned distribution and challenge the robustness of frozen visual decoder. (b) RobustTok leverages latent perturbation to enhance the robustness of tokenizer, thus boosting the image generation quality.
• Figure 2: Visualization of (a) traditional tokenizer, (b) semantic tokenizer, and (c) our RobusTok in reconstruction task with Latent Perturbation. Non-semantic tokenizer leads to distorted reconstructions when perturbations are introduced while our method shows promising robustness to those perturbations.
• Figure 3: RobustTok overview. We adopt vision transformer as our encoder $\mathcal{E}$ and decoder $\mathcal{D}$. $\beta$ of data in one batch will process our Latent Perturbation, which will be randomly replaced by top-$\delta$ neighbor from codebook with probability $\alpha$. A frozen DINO encoder is utilized to supervise our latent space.
• Figure 4: Comparison of rFID-gFID and pFID-gFID curves of differnt tokenizers under LlamaGen-B training setting. $K$ denotes codebook size. Each point represents a method in \ref{['tab:token_setting']}.
• Figure 5: T-SNE visualization of latent space of tokenizer trained with and without latent perturbation. Colors and thresholds represent the frequency of tokens being used during inference without perturbation.