Image Tokenizer Needs Post-Training

TL;DR

RobusTok argues that image tokenizers fail downstream generation because they are trained to reconstruct clean tokens, but at inference time the autoregressive or diffusion generator hands them sampling-error-corrupted tokens. The fix is a two-stage recipe: main-training with annealed latent perturbation (randomly swap tokens with top-δ codebook neighbors during reconstruction) to make the latent space robust, and a lightweight post-training that fine-tunes only the decoder against real generator outputs, using a “preservation ratio” to interpolate between reconstructed and generated tokens for paired supervision. The paper also introduces pFID — reconstruction FID measured under codebook-neighbor perturbation — and shows it correlates with downstream gFID where vanilla rFID does not. On ImageNet 256×256 with a 400M RAR generator, RobusTok reaches gFID 1.60, then 1.36 after post-training — a new SOTA under 500M parameters.

Key claims

Reconstruction FID does not predict generation FID across tokenizers, but pFID (FID under codebook-neighbor perturbation, averaged over a swept perturbation rate and strength grid) does — within fixed codebook size, pFID-gFID is strongly correlated while rFID-gFID is not [§5.2, Fig. 5].
Smaller codebooks yield better generation at matched reconstruction, and semantic tokenizers (DINOv2-aligned) reconstruct and generate better than non-semantic ones — codebook size and semantic regularization are the two pre-existing axes that already correlate with gFID [§5.2 “General observations”].
The main-training latent-perturbation recipe: within a batch, perturb a fraction β of images; for each selected image, replace a fraction α of tokens with one of their top-δ nearest neighbors in the codebook; anneal α from a high value toward 0.5 over training [§4.2].
Annealing matters: no annealing causes mode collapse; annealing to zero produces an over-deterministic decoder that loses Fig. 7’s latent flexibility; annealing to half is the sweet spot — gFID 4.62 vs 5.32 (anneal-to-0) vs >7 (no anneal) on the baseline ablation [§5.3, Table 5].
Post-training freezes encoder + quantizer and fine-tunes only the decoder against generator outputs, using a “preservation ratio” γ that replaces a fraction of generated tokens with ground-truth tokens to maintain image–latent pairing for supervision [§4.3].
A “more data ≠ better” finding for post-training: scaling from 0.2M → 1M post-training samples barely moves gFID (1.60 → 1.59); training stability dominates over data volume [§5.3, Table 4].
Headline: 400M RAR + RobusTok reaches gFID 1.60; adding post-training pushes it to 1.36 with IS 300.2 — new SOTA under 500M parameters [Table 2].
Post-training generalizes across tokenizer families: applied off-the-shelf to MAETok (continuous, diffusion), LlamaGen (discrete, AR), and GigaTok (discrete, AR), it consistently improves gFID — MAETok 1.87 → 1.68, LlamaGen 3.80 → 3.51, GigaTok 3.84 → 3.68 [Table 3].
Visualization of codebook usage frequency shows perturbation-trained tokenizers concentrate usage on a small set of “key” reusable tokens with the rest serving as supportive fine-detail tokens, while non-perturbed tokenizers spread usage uniformly — the perturbation explicitly induces a coarse-to-fine token hierarchy [§5.3, Fig. 7].

Method

RobusTok is a ViT-encoder/ViT-decoder VQ tokenizer trained in two stages. Main-training patches the input image, concatenates a set of learnable tokens, runs the ViT encoder, applies vector quantization, then injects a plug-and-play latent perturbation: for a fraction β of images in the batch (the authors use 0.5 annealed to ≈0.25), randomly select α·H·W tokens and replace each with one of its top-δ nearest neighbors in the learned codebook. δ is annealed from 0.5 toward 0.5/2 over training. The decoder reconstructs from these perturbed latents while a DINOv2 semantic-regularization branch keeps the codebook’s clean-side meaningful. The intuition: in AR sampling, errors take the form of substituting a token with a codebook-near alternative, so training-time perturbation should mimic that exact distribution.

Post-training mitigates the residual gap between synthetic perturbations and real generator errors. The encoder and quantizer are frozen; only the decoder is fine-tuned, paired with the pretrained discriminator under the same adversarial-loss combination. The problem is that real generated latents have no ground-truth image pairing. The fix is a preservation ratio γ: during AR sampling, replace γ·H·W of the generated tokens with ground-truth tokens from the source image, smoothly interpolating between full reconstruction (γ=1) and full generation (γ=0). This is teacher-forcing applied at the image–latent pair level. For diffusion-based tokenizers, the same mechanism is realized via SDEdit-style partial-noise re-injection. The decoder learns to reconstruct from this interpolated mixture, transferring robustness from synthetic codebook-neighbor noise to actual generator-induced distribution shift.

Results

The strongest result is the rFID/pFID/gFID decoupling. The paper benchmarks 11 open-sourced tokenizers across 4 codebook sizes and plots rFID-vs-gFID and pFID-vs-gFID. rFID is uncorrelated with gFID; pFID tracks gFID closely within each codebook size [Fig. 5]. This is the metric contribution and it explains the paper’s title — reconstruction-only evaluation has been measuring the wrong thing.

The systematic comparison at ImageNet 256×256 [Table 2]: RAR (the prior strong AR baseline) reaches gFID 1.70 / IS 299.5 / rFID 2.28 / pFID 5.03. RAR + RobusTok-L reaches gFID 1.60 / IS 305.8 / rFID 1.02 / pFID 2.28 — i.e. simultaneously better reconstruction and better generation than RAR. RAR + RobusTok-L + post-training pushes to gFID 1.36 / IS 300.2 with the same rFID/pFID. This is the new SOTA-under-500M-parameters number.

Broad applicability of post-training [Table 3]: applied to MAETok (continuous, diffusion-side), LlamaGen (discrete, AR-side), GigaTok (discrete, AR-side), and both RobusTok-B/L, gFID improves consistently — MAETok 1.87→1.68, LlamaGen 3.80→3.51, GigaTok 3.84→3.68, RobusTok-B 1.83→1.60, RobusTok-L 1.60→1.36. The post-training recipe is tokenizer-agnostic and works for both AR and diffusion downstreams.

Ablations [Table 5]: the perturbation parameters matter. Latent perturbation alone (large α) hurts reconstruction (rFID 3.97) and only marginally helps gFID; the annealing schedule plus codebook-size expansion (4096) recovers a strong rFID 1.02 and a strong pFID 2.28. Post-training adds the final FID step.

Why it’s interesting

RobusTok sits cleanly inside the Diffusion training efficiency thread but pushes a different lever than the page’s existing entries. The two natural contrasts:

Against What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion (PAE): both papers explicitly reject “optimize rFID” and argue for a different objective in tokenizer training. PAE’s lever is manifold geometry shaped by VFM priors. RobusTok’s lever is robustness to sampling-error perturbation. Both report tokenizer-side gFID gains on ImageNet 256×256 and both critique rFID as the wrong metric. PAE’s perturbation-based regularization is in fact a close cousin of RobusTok’s latent perturbation — but PAE applies it to enforce local continuity of the latent manifold, while RobusTok applies it to make the decoder robust to codebook-neighbor substitutions specifically. They likely compose.
Against Making Reconstruction FID Predictive of Diffusion Generation FID (which directly attacks the rFID-vs-gFID gap): RobusTok’s pFID is another concrete proposal for a tokenizer-side metric that actually predicts downstream generation. The two metrics should be compared head-to-head — pFID is cheaper (no generator training) but tied to a codebook-neighbor noise model; the other line of work fits gFID directly. Worth noting that filing both makes “what’s the right tokenizer-evaluation metric” a real open question on the wiki, not a settled one.

The post-training contribution is also worth flagging in its own right: it is a generator-conditional fine-tune of an existing tokenizer’s decoder, recovering ~0.3 gFID for free on top of any pretrained tokenizer (LlamaGen, MAETok, GigaTok all gain). This is unusual — tokenizers are normally treated as frozen-once-trained — and it implies any latent-diffusion or AR pipeline shipping today can claw back a measurable FID delta with a few epochs of decoder-only fine-tuning against the trained generator’s outputs.