UniTok: A Unified Tokenizer for Visual Generation and Understanding

TL;DR

UniTok is a discrete visual tokenizer designed to serve both generation (VQVAE-style reconstruction) and understanding (CLIP-style image-text alignment) in a single unified MLLM. The paper’s central diagnosis: when prior unified tokenizers like VILA-U underperform domain-specific tokenizers, the cause is not the often-blamed “loss conflict” between reconstruction and contrastive objectives but the limited representational capacity of discrete tokens — token factorization (e.g. 768→16 dim projection) and quantization with small codebooks (4k–16k) each cost significant downstream accuracy on VQA tasks. The fix is multi-codebook quantization (MCQ): evenly chunk the latent into N sub-vectors and discretize each with its own independent sub-codebook, scaling the effective vocabulary multiplicatively (N=8 sub-codebooks of 4096 entries each ⇒ $4096^8$ combinations) while keeping each individual codebook small enough to optimize. Combined with attention-based factorization, UniTok hits 0.38 rFID (vs SD-VAE’s 0.87) and 78.6% ImageNet zero-shot (vs CLIP’s 76.2%), and the downstream Llama-2-7B unified MLLM beats VILA-U on VQAv2 (76.8 vs 75.3), MME (1448 vs 1336), and MJHQ FID (7.46 vs 12.81).

Key claims

The unified-tokenizer underperformance gap relative to CLIP is not primarily due to reconstruction-vs-contrastive loss conflict; it stems from token factorization and discretization, which compromise representational capacity [§3.2].
Adding only the 768→16→768 factorization projection (no quantization, no reconstruction loss) on top of CLIP causes severe VQA degradation — direct evidence that dimensionality compression alone, before any discretization, costs downstream accuracy [§3.2, Fig. 3].
Adding 16k-codebook discretization on top of the factorized features causes an additional ~2.1 average VQA accuracy drop [§3.2].
Multi-codebook quantization (MCQ) evenly splits the latent into N chunks and discretizes each with an independent sub-codebook, exponentially expanding effective vocabulary ( $V^N$ combinations) while keeping each sub-codebook at a tractable size — sidestepping the dead-code/utilization problem of monolithic large codebooks [§3.3, Eq. 3].
Attention factorization (causal multi-head attention modules replacing linear/conv projection layers) preserves more semantic information than over-simplified linear factorization [§3.3, Fig. 4].
UniTok reaches 0.38 rFID and 78.6% ImageNet zero-shot at 256² with 256 tokens per image, beating both unified baselines (VILA-U: 1.80 rFID / 73.3%) and the SD-VAE continuous tokenizer (0.87 rFID) [Table 1].
With multi-codebook quantization in place, joint reconstruction + contrastive training is essentially conflict-free — the jointly trained tokenizer matches the contrastive-only tokenizer on understanding and slightly beats the reconstruction-only tokenizer on generation [Table 5].
Increasing sub-codebook count from 1 to 8 (at fixed global codebook size) monotonically improves both rFID and ImageNet zero-shot accuracy [Table 6].
The unified MLLM (Llama-2-7B + UniTok) outperforms VILA-U on VQAv2 (76.8 vs 75.3), MME (1448 vs 1336), MM-Vet (33.9 vs 27.7), and MJHQ-30K FID (7.46 vs 12.81) [Table 2, Table 4].
For MLLM input efficiency, the N sub-codebook codes per spatial token are merged into one input token via MLP projection; for output, a depth-transformer head (RQ-Transformer-style) autoregressively predicts the N codes per spatial position [§3.4].

Method

UniTok instantiates a hybrid ViT/CNN encoder (ViTamin-L/16) and jointly optimizes a VQVAE-style reconstruction loss (pixel L2 + LPIPS perceptual + StyleGAN-style discriminator + VQ commitment) and a CLIP-style image-text contrastive loss, weighted equally. The architectural pieces that make unified training actually work are:

Multi-codebook quantization (MCQ). Given a per-spatial-position latent $z \in \mathbb{R}^d$ , split evenly into N chunks $\{z_1, \dots, z_N\}$ , each of dim $d/N$ . Each chunk is independently looked up against its own codebook $C_k$ of size $V$ . UniTok uses N=8 sub-codebooks, V=4096 entries each, with per-chunk latent dim 8 (global 64-d). Theoretical vocabulary is $4096^8$ combinations vs the standard 16k for a single codebook, and individual codebooks stay small enough to avoid the dead-code problem reported at >16k entries.
Attention factorization. Instead of a linear or convolutional layer to project the encoder’s 768-d output down to the 64-d quantized space (and back up), UniTok uses causal multi-head attention modules. Causality is preserved so the resulting codes remain compatible with autoregressive MLLM generation.
Unified MLLM integration. Following Liquid’s recipe, the discrete codebook embeddings are projected via MLP into the MLLM’s token space. To avoid blowing up sequence length by N×, each spatial position’s N codes are merged into a single MLLM input token; on the output side, a depth-transformer head (à la RQ-Transformer / VILA-U) autoregressively decodes the N sub-codes from a single hidden state. The MLLM is Llama-2-7B, pretrained on 10M DCLM text + 30M MJ-style synthetic + 30M re-captioned COYO/Laion, finetuned on 1.5M T2I + 1.5M MM-IT data from Mini-Gemini.

The tokenizer is trained for one epoch on DataComp-1B (1.28B image-text pairs) at 256² with global batch 16k, LR 1e-3 (discriminator 2e-4). Two init regimes are reported: random init (default) and CLIP-pretrained init.

Results

Reconstruction: 0.38 rFID on ImageNet 256² at 256 tokens, beating VQ-GAN (4.98), RQ-VAE (1.30), VAR (0.90 with 680 tokens), TokenFlow (1.37), VILA-U (1.80), and even SD-VAE (0.87, continuous) [Table 1].
Zero-shot: 78.6% ImageNet top-1 with CLIP-init, 70.5% with random init — for context, CLIP is 76.2% and SigLIP 80.5% [Table 1]. The 70.5%→78.6% gap is attributed to the 1-epoch DataComp-1B schedule being too short to fully train CLIP from scratch.
Downstream MLLM understanding: VQAv2 76.8 / GQA 61.1 / TextVQA 51.6 / POPE 83.2 / MME 1448 / MM-Vet 33.9. Beats VILA-U on every metric; closes most of the gap to continuous-token MLLMs like Janus (VQAv2 77.3) [Table 2].
Downstream MLLM generation: GenAI-Bench advanced-prompts overall 0.67 (vs VILA-U 0.64, Liquid 0.65, DALL-E 3 0.70 with diffusion); MJHQ-30K FID 7.46 at 256² (vs VILA-U 12.81, Liquid 5.47 at 512²) [Table 3, Table 4].
Ablation: jointly trained tokenizer (rFID 0.72, VQAv2 69.14, MME 1333) ≈ contrastive-only tokenizer on understanding (VQAv2 68.95, MME 1373) and ≥ reconstruction-only tokenizer on generation [Table 5].

Why it’s interesting

UniTok directly addresses the central design question for any unified-tokenizer effort: how to scale codebook capacity without dead codes or training instability — and gives a clean recipe (chunk + per-chunk small codebook) backed by a controlled ablation against the standard “just make the codebook bigger” approach. Boxiao’s framing in the Slack post — “alternative strategy to RQ that evenly chunks the continuous latent and learns a sub-codebook for each chunk” — is exactly right: residual quantization stacks codebooks sequentially over residuals while MCQ stacks them in parallel over orthogonal latent chunks (it’s the visual-tokenizer equivalent of product quantization from the ANN literature, e.g. TurboQuant: Online Vector Quantization with Near-optimal Distortion Rate‘s PQ baselines). The two recipes are not mutually exclusive — VAR / RQ-VAE-style depth-transformer heads can sit on top of MCQ codes (UniTok does this for autoregressive generation), which is what makes this directly composable with existing tokenizer pipelines.

Complementary to the wiki’s most recent unified-tokenization data point Improved Baselines with Representation Autoencoders (RAEv2, Saining Xie’s team): RAEv2 keeps the latent continuous (frozen DINOv3 encoder + multi-layer aggregation + REPA) and gets 1.06 gFID on ImageNet-256 in 80 epochs, while UniTok keeps the latent discrete and tries to recover continuous-tokenizer quality via product-style quantization scaling. The two papers stake out the opposite ends of the unified-tokenizer design space — continuous-with-pretrained-encoder vs discrete-with-engineered-codebook — and both beat their respective baselines, so the right axis to compare them on is probably downstream AR-MLLM compatibility (which UniTok demonstrates with Llama-2-7B) vs continuous-DiT integration (which RAEv2 demonstrates). Also worth flagging against OneVision-Encoder: Codec-Aligned Sparsity as a Foundational Principle for Multimodal Intelligence‘s “codec-aligned sparsity” recipe and What Matters for Diffusion-Friendly Latent Manifold? Prior-Aligned Autoencoders for Latent Diffusion‘s prior-aligned autoencoders, which approach the same “what should the visual latent look like” question from the latent-geometry side rather than the discretization side.