Image Tokenizer Needs Post-Training

TL;DR

<3-5 sentence high-level summary> The paper investigates why image tokenizers optimized for reconstruction underperform during generation, due to distribution differences between training-time ground-truth tokens and sampling-time tokens. It introduces RobusTok, a two-stage tokenizer framework: main-training with latent perturbations to build a robust latent space, and post-training to align the decoder with a well-trained generator using a preservation ratio and SDEdit for diffusion. A novel pFID metric is proposed to better predict generation quality from tokenizer robustness, and experiments across autoregressive and diffusion models demonstrate broad applicability and significant gains (e.g., ~1.36 gFID with ~400M generators). The work highlights robustness of the latent space as a key design criterion for tokenizers, enabling more reliable and scalable image generation with discrete latents.

Abstract

Recent image generative models typically capture the image distribution in a pre-constructed latent space, relying on a frozen image tokenizer. However, there exists a significant discrepancy between the reconstruction and generation distribution, where current tokenizers only prioritize the reconstruction task that happens before generative training without considering the generation errors during sampling. In this paper, we comprehensively analyze the reason for this discrepancy in a discrete latent space, and, from which, we propose a novel tokenizer training scheme including both main-training and post-training, focusing on improving latent space construction and decoding respectively. During the main training, a latent perturbation strategy is proposed to simulate sampling noises, \ie, the unexpected tokens generated in generative inference. Specifically, we propose a plug-and-play tokenizer training scheme, which significantly enhances the robustness of tokenizer, thus boosting the generation quality and convergence speed, and a novel tokenizer evaluation metric, \ie, pFID, which successfully correlates the tokenizer performance to generation quality. During post-training, we further optimize the tokenizer decoder regarding a well-trained generative model to mitigate the distribution difference between generated and reconstructed tokens. With a $\sim$400M generator, a discrete tokenizer trained with our proposed main training achieves a notable 1.60 gFID and further obtains 1.36 gFID with the additional post-training. Further experiments are conducted to broadly validate the effectiveness of our post-training strategy on off-the-shelf discrete and continuous tokenizers, coupled with autoregressive and diffusion-based generators.

Figures (24)

• Figure 1: (a) Discrepancy between reconstruction and generation task imposes a latent token distribution difference between them. Specifically, reconstruction always rely on true tokens whereas generation task always sample out-of-distribution (OOD) tokens. To resolve this problem, we propose RobusTok (b) to enhance the robustness of tokenizer during main-trainig by latent perturbation, and (c) align the generated latent space with its target image in post-training stage.
• Figure 2: RobusTok 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 3: 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 4: Generated images under different $\sigma$ for (left) autoregressive and (right) diffusion model.
• Figure 6: Maximum Mean Discrepency (MMD) between generated and $\alpha$-perturbed latent.