Source-linked AI summary
UniTok: A Unified Tokenizer for Visual Generation and Understanding
Chuofan Ma, Yi Jiang, Junfeng Wu, Jihan Yang, Xin Yu, Zehuan Yuan, Bingyue Peng, Xiaojuan Qi
TL;DR
Unified multimodal models lack a visual tokenizer that simultaneously supports generation and understanding, with discrete token capacity emerging as the central bottleneck. UniTok uses multi-codebook quantization and attention projection to expand that capacity, achieving strong ImageNet reconstruction and zero-shot accuracy while supporting native visual generation in MLLMs. It also improves classifier-free-guidance-free generation, although the reported one-epoch training schedule limits semantic representation learning.
Problem
Existing visual tokenizers specialize in generation or understanding, and combining their objectives has produced convergence and understanding problems attributed to limited discrete-token capacity.
Method
UniTok combines multi-codebook quantization with attention projection to enlarge discrete vocabulary and bottleneck dimension for unified visual generation and understanding.
Results
UniTok achieves 0.38 reconstruction FID and 78.6% zero-shot accuracy on ImageNet, enables native visual generation in MLLMs, and reduces cfg-free generation FID from 14.6 to 2.5.
Takeaways & Limitations
Reconstruction and semantic supervision do not inherently conflict when the discrete feature space is sufficiently expressive.
Takeaways & Limitations
UniTok is trained for only one epoch, which is insufficient for CLIP-based semantic representation learning and may limit understanding performance.
Abstract
from arXiv · showhide
Visual generative and understanding models typically rely on distinct tokenizers to process images, presenting a key challenge for unifying them within a single framework. Recent studies attempt to address this by connecting the training of VQVAE (for autoregressive generation) and CLIP (for understanding) to build a unified tokenizer. However, directly combining these training objectives has been observed to cause severe loss conflicts. In this paper, we show that reconstruction and semantic supervision do not inherently conflict. Instead, the underlying bottleneck stems from limited representational capacity of discrete token space. Building on these insights, we introduce UniTok, a unified tokenizer featuring a novel multi-codebook quantization mechanism that effectively scales up the vocabulary size and bottleneck dimension. In terms of final performance, UniTok sets a new record of 0.38 rFID and 78.6% zero-shot accuracy on ImageNet. Besides, UniTok can be seamlessly integrated into MLLMs to unlock native visual generation capability, without compromising the understanding performance. Additionally, we show that UniTok favors cfg-free generation, reducing gFID from 14.6 to 2.5 on ImageNet 256$\times$256 benchmark. GitHub: https://github.com/FoundationVision/UniTok.
1 Introduction
Unified multimodal models need a visual tokenizer that supports both generation and understanding, but CLIP and VQVAE each favor one capability. UniTok attributes the gap primarily to limited discrete-token capacity and addresses it with multi-codebook quantization and attention projection.
- CLIP supports multimodal understanding, whereas discrete VQVAE tokens suit autoregressive generation but struggle to capture essential semantics.
- Combining CLIP supervision with VQVAE training causes severe convergence issues and trails the CLIP baseline in multimodal understanding.
- Ablations show the unified tokenizer’s understanding gap mainly arises from vector quantization, especially token factorization and discretization.These operations reduce the expressiveness of visual tokens and expose limited discrete-token capacity as the primary bottleneck.
- Multi-codebook quantization partitions tokens across small sub-codebooks, exponentially expanding vocabulary size while avoiding large-monolithic-codebook optimization problems.Attention modules replace traditional linear projection layers for token factorization.
- 0.38 reconstruction FID and 78.6% zero-shot accuracy are achieved by UniTok on ImageNet at 256×256 resolution.UniTok also enables native visual generation in MLLMs while maintaining understanding performance, and reduces cfg-free generation FID from 14.6 to 2.5.
2 Related Work
Prior visual tokenizers generally specialize in generation or understanding, leaving unified vision-language modeling with a mismatch between discrete language-compatible tokens and effective visual representations. Related approaches use vector quantization, continuous tokenizers, diffusion models, or CLIP discretization, but reported methods can impair understanding or increase complexity.
- Image Tokenization for Generation: Vector-quantized tokenizers provide compact discrete latent spaces compatible with autoregressive and masked generative models.VQVAE introduced mapping continuous tokens to nearest neighbors in a learnable codebook, while VQGAN extended the approach with perceptual loss.
- Image Tokenization for Understanding: Pretrained CLIP is a common understanding tokenizer because language alignment supports multimodal tasks, but its continuous features complicate uniform vision-text modeling.Discretizing CLIP tokens or using VQVAE encoders has been observed to substantially impair MLLM understanding performance.
- Unified Vision-Language Models: Unified vision-language models have used continuous visual tokenizers with pretrained diffusion models, increasing model complexity and disconnecting visual sampling.
3 Method
UniTok unifies visual generation and understanding by jointly applying reconstruction and semantic supervision, then addressing the discrete-token bottleneck with multi-codebook quantization and attention projection. The analysis attributes the main understanding gap to token factorization and discretization rather than reconstruction supervision.
- Unified supervision: UniTok jointly trains a tokenizer to reconstruct images and align discrete latent features with text captions.The reconstruction objective combines pixel, perceptual, discriminator, and vector-quantization losses, while CLIP-style contrastive supervision adds semantic alignment.
- Quantization bottleneck: Factorizing 768-d CLIP tokens through 16-d channels significantly compromises token expressiveness and degrades downstream VQA performance.The experiment excludes vector quantization and reconstruction supervision, isolating channel compression as the tested change.
- Quantization bottleneck: A 16k codebook causes an average accuracy drop of 2.1 in VQA tasks when discretizing factorized tokens.The result supports information loss from mapping continuous factorized tokens to a comparatively small codebook.
- Quantization bottleneck: Reconstruction supervision has negligible impact on downstream understanding performance, indicating limited discrete-token capacity is the primary bottleneck.Joint training can reduce ImageNet classification accuracy and reconstruction FID, but those degradations diminish after quantization improves.
- UniTok design: Multi-codebook quantization splits each latent token into chunks and quantizes them with independent sub-codebooks, expanding vocabulary and latent dimension.With four 16k sub-codebooks, the theoretical vocabulary grows from 2^14 to 2^56 combinations, while latent dimension increases from 16-d to 64-d.
- UniTok design: Attention projection replaces simplified linear or convolutional factorization with multi-head attention-based channel compression to preserve richer semantics.The method uses average pooling across attention heads instead of concatenating their features.
- Unified MLLM integration: UniTok reuses its code embeddings in a unified MLLM, projecting them into the language-model token space with an MLP projector.This integration uses discrete vision and language sequences under a universal next-token prediction loss.
4 Experiments
Experiments evaluate UniTok as a tokenizer for reconstruction, generation, multimodal understanding, and unified MLLM generation. Results show strong ImageNet tokenizer metrics, improved CFG-free generation, competitive VQA and text-to-image performance, and consistent benefits from multi-codebook quantization.
- Tokenizer Comparison: 0.38 rFID and 78.6% zero-shot accuracy are achieved on ImageNet at 256×256 resolution.The rFID uses 16× downsampling; pretrained CLIP initialization boosts zero-shot accuracy from 70.8% to 78.6%.
- Class-Conditional Image Generation: UniTok reduces CFG-free ImageNet generation FID by 12.11 versus the VQGAN baseline under the same generator setup.The result suggests UniTok learns a more structured code distribution that benefits autoregressive modeling.
- Unified Understanding and Generation: UniTok improves unified MLLM understanding over discrete-tokenizer baselines, including 7.2% on VQAv2, 3.3% on TextVQA, and 112 MME-Perception points.It also narrows the gap with MLLMs using continuous visual tokenizers.
- Unified Understanding and Generation: UniTok outperforms most unified MLLMs and is competitive with diffusion models on GenEval while using the same text-to-image data as Liquid.The comparison covers unified and domain-specific visual generation methods.
- Unified Understanding and Generation: 7.46 FID is reported for UniTok at 256 resolution on MJHQ-30K, compared with 5.47 for Liquid at 512 resolution.FID across resolutions is less comparable because higher-resolution images can capture more fine-grained detail.
- Ablation Studies: Increasing the number of sub-codebooks consistently improves reconstruction FID and classification accuracy at constant global codebook size.MCQ also outperforms residual quantization under the 64-d bottleneck setting, with better reconstruction and classification performance.
5 Limitations and Conclusion
UniTok addresses unified visual tokenization by targeting limited discrete-token representational power, and its ablation indicates that discriminative and generative representation learning do not inherently conflict. The method achieves strong downstream generation and understanding performance, but its one-epoch training schedule limits CLIP-based semantic learning.
- Conclusion: UniTok identifies limited representational power of discrete tokens as the main challenge in unifying visual generation and understanding.The paper introduces multi-codebook quantization and attention projection to address this limitation.
- Conclusion: UniTok excels in downstream visual generation and understanding tasks while serving as a unified tokenizer for multimodal models.The paper positions unified visual tokenization as a cornerstone of unified multimodal large language models.
- Conclusion: The ablation study finds that discriminative and generative representation learning does not inherently conflict.The paper instead attributes the unification challenge mainly to limited discrete-token capacity.
- Limitations: UniTok is trained for only one epoch, which is insufficient for CLIP-based semantic representation learning.The authors believe extending the training schedule could further improve tokenizer performance, especially understanding performance.
A Attention Projection Modules
The attention projection adaptations modify the traditional multi-head attention module to support channel compression and expansion during token factorization.
- Attention Projection Modules: The modified attention blocks enable channel compression and expansion in token factorization.Figure 6 describes these adaptations as changes to the traditional multi-head attention module.
- Attention Projection Modules: In the attention module, C denotes the channel dimension, c the per-head channel dimension, and h the number of heads.The dimensions satisfy C = h × c.
B More Generation Results
UniTok consistently delivers superior generation performance on GenAI-Bench across the reported basic- and advanced-prompt evaluations.
- More Generation Results: UniTok consistently delivers superior generation performance on GenAI-Bench.The results are reported separately for basic prompts and advanced prompts.
- More Generation Results: Table 10 compares visual generation methods on GenAI-Bench using basic prompts.The table is the reported evaluation for the basic-prompt setting.
- More Generation Results: Table 11 compares visual generation methods on GenAI-Bench using advanced prompts.The table is the reported evaluation for the advanced-prompt setting.