UniTok: A Unified Tokenizer for Visual Generation and Understanding

2025-02-27

AI Summary

Visual generation and understanding tasks suffer from representational misalignment, hindering unified modeling. This paper introduces UniTok—the first discrete visual tokenizer jointly optimized for both generation and understanding. Its core innovation is a multi-codebook vector quantization architecture that significantly expands the capacity of the discrete latent space without compromising training stability, thereby overcoming the representational bottleneck inherent in conventional single-codebook approaches. Through end-to-end joint optimization and cross-task consistency regularization, UniTok achieves, for the first time, superior performance of a unified discrete tokenizer over task-specific continuous ones: it attains an rFID of 0.38 on ImageNet (outperforming SD-VAE’s 0.87) and zero-shot classification accuracy of 78.6% (exceeding CLIP’s 76.2%). These results effectively bridge the fundamental gap between generative fidelity and semantic understanding.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Economics, Online Markets and Human Computation: Economic ramifications for generative AI infrastructure and applications

Abstract

The representation disparity between visual generation and understanding imposes a critical gap in integrating these capabilities into a single framework. To bridge this gap, we introduce UniTok, a discrete visual tokenizer that encodes fine-grained details for generation while also capturing high-level semantics for understanding. Despite recent studies have shown that these objectives could induce loss conflicts in training, we reveal that the underlying bottleneck stems from limited representational capacity of discrete tokens. We address this by introducing multi-codebook quantization, which divides vector quantization with several independent sub-codebooks to expand the latent feature space, while avoiding training instability caused by overlarge codebooks. Our method significantly raises the upper limit of unified discrete tokenizers to match or even surpass domain-specific continuous tokenizers. For instance, UniTok achieves a remarkable rFID of 0.38 (versus 0.87 for SD-VAE) and a zero-shot accuracy of 78.6% (versus 76.2% for CLIP) on ImageNet. Our code is available at https://github.com/FoundationVision/UniTok.

Problem

• Bridging visual generation and understanding gap
• Enhancing discrete tokenizer representational capacity
• Avoiding training instability with multi-codebook quantization

Innovation

• Unified discrete visual tokenizer
• Multi-codebook quantization technique
• Enhanced latent feature space capacity