PaperlayerView exact revision on arXiv

Source-linked AI summary

TokenFlow: Unified Image Tokenizer for Multimodal Understanding and Generation

Liao Qu, Huichao Zhang, Yiheng Liu, Xu Wang, Yi Jiang, Yiming Gao, Hu Ye, Daniel K. Du, Zehuan Yuan, Xinglong Wu

arXiv:2412.03069v2cs.CVcs.AI

TL;DR

TokenFlow addresses the challenge of unifying multimodal understanding and image generation, which require different granularities of visual information. Its dual-codebook tokenizer decouples semantic and pixel-level feature learning while aligning them through shared indices, achieving stronger understanding, reconstruction, and autoregressive generation results.

  • Problem

    Existing unified approaches use a single reconstruction-targeted VQ encoder despite understanding and generation requiring different visual granularities, creating a trade-off that can compromise multimodal understanding.

  • Method

    TokenFlow uses dual codebooks to separately learn semantic and pixel-level features while preserving their alignment through a shared mapping and joint quantization.

  • Results

    TokenFlow surpasses LLaVA-1.5 13B by 7.2% on average in multimodal understanding, achieves FID 0.63 at 384×384 reconstruction, and reaches GenEval 0.55 at 256×256 autoregressive generation.

  • Takeaways & Limitations

    TokenFlow provides a unified visual tokenizer with strong multimodal understanding, reconstruction quality, and competitive visual generation using substantially fewer inference steps.

  • Takeaways & Limitations

    Conventional independent top-k sampling in the next-scale generation paradigm can cause image collapse and repetitive patterns, especially with limited inference steps.

Abstract

from arXiv · show

We present TokenFlow, a novel unified image tokenizer that bridges the long-standing gap between multimodal understanding and generation. Prior research attempt to employ a single reconstruction-targeted Vector Quantization (VQ) encoder for unifying these two tasks. We observe that understanding and generation require fundamentally different granularities of visual information. This leads to a critical trade-off, particularly compromising performance in multimodal understanding tasks. TokenFlow addresses this challenge through an innovative dual-codebook architecture that decouples semantic and pixel-level feature learning while maintaining their alignment via a shared mapping mechanism. This design enables direct access to both high-level semantic representations crucial for understanding tasks and fine-grained visual features essential for generation through shared indices. Our extensive experiments demonstrate TokenFlow's superiority across multiple dimensions. Leveraging TokenFlow, we demonstrate for the first time that discrete visual input can surpass LLaVA-1.5 13B in understanding performance, achieving a 7.2\% average improvement. For image reconstruction, we achieve a strong FID score of 0.63 at 384*384 resolution. Moreover, TokenFlow establishes state-of-the-art performance in autoregressive image generation with a GenEval score of 0.55 at 256*256 resolution, achieving comparable results to SDXL.

1. Introduction

TokenFlow targets the persistent divide between multimodal understanding and visual generation by using discrete visual tokens with separate semantic and pixel-level representations. Its dual-codebook design reports strong results across understanding, reconstruction, and autoregressive generation.

  • Motivation: Existing multimodal systems use distinct architectures for understanding and generation, motivating unified approaches.Understanding models align vision encoders with language models, while generation uses diffusion methods or discrete image tokens.
  • Motivation: Discrete-token frameworks offer a simpler route to unifying visual and textual information within one next-token prediction architecture.VQ encoders convert images into discrete tokens that can be processed alongside text and trained end-to-end.
  • Challenge: Reconstruction-targeted VQ encoders trade semantic richness for pixel fidelity, limiting their suitability for multimodal understanding.Understanding requires semantic representations, whereas generation requires spatial and textural detail.
  • TokenFlow: TokenFlow decouples semantic and pixel-level feature learning with dual codebooks aligned through shared indices.A semantic encoder supplies CLIP-style semantic priors, while a pixel encoder captures detailed visual information.
  • Results: 0.63 FID at 384×384 resolution and 0.55 GenEval at 256×256 resolution demonstrate strong reconstruction and autoregressive generation results.The generation result is reported as state of the art in the autoregressive paradigm.
  • Results: 7.2% average improvement over LLaVA-1.5 13B is reported for multimodal understanding with TokenFlow's discrete visual input.The paper presents this as the first case in which discrete visual input surpasses that baseline in understanding performance.

2. Related Work

Related work has pursued unified multimodal understanding and generation with shared or separate tokenizers, but existing designs retain architectural or representational limitations. TokenFlow instead uses dual codebooks with a shared mapping to connect feature levels.

  • VQ image tokenizers: VQ-VAE, VQ-VAE-2, and VQGAN developed reconstruction-based discrete image tokenization through nearest-codebook quantization, hierarchical modeling, and perceptual or adversarial losses.These approaches primarily support autoregressive image generation.
  • Unified multimodal models: Chameleon, EMU3, and Show-o use VQ tokenizers for both understanding and generation but often require multimodal training from scratch and can suffer performance degradation.
  • Unified multimodal models: Janus separates understanding and generation tokenizers, increasing model complexity without fundamentally resolving the modality challenge.
  • TokenFlow: A concurrent unified vision tower constrains low- and high-level representations within one flow, whereas TokenFlow learns their joint mapping through dual codebooks and shared indices.This design is presented as enabling flexible combinations of low- and high-level visual features.

3. Method

TokenFlow unifies multimodal understanding and image generation by jointly learning semantic and pixel-level representations through dual encoders, dual codebooks, and shared mapping. Its unified representation supports both downstream tasks while preserving complementary visual information.

  • Decoder and Training Objective: Semantic and pixel quantized features are independently decoded for semantic alignment and image reconstruction, then concatenated for downstream tasks.Training uses a semantic loss, vector-quantization loss, and pixel reconstruction loss, with pixel reconstruction combining pixel-wise, perceptual, and adversarial terms.
  • Encoder: TokenFlow uses dual semantic and pixel encoders to extract complementary high-level and low-level image features.The semantic encoder is initialized with a pretrained text-aligned vision encoder, while the pixel encoder captures pixel-level features.
  • Quantization: Dual codebooks share a unified mapping, allowing quantization to jointly consider semantic information and pixel details.The shared mapping assigns semantically and visually similar patches to common indices.
  • Quantization: The selected index minimizes dsem,i + wdis · dpix,i, balancing semantic and pixel-level distances during quantization.wdis is the distance-balance weight, and the framework adopts a multi-scale VQ structure to enrich codebook representations.
  • Multi-task Capabilities: The unified representation bridges understanding and generation and remains extensible to additional feature codebooks without architectural modifications.With a codebook size of 131,072, TokenFlow maintains over 95% utilization while achieving its best reconstruction and multimodal understanding performance.
  • Inference Strategy: Multi-step sampling produces more coherent and visually appealing generations than single-pass top-k or top-p sampling.The approach addresses image collapse, repetitive local patterns, and token inconsistency associated with independent sampling in the next-scale paradigm.

4. Experiments

TokenFlow is evaluated across reconstruction, multimodal understanding, and visual generation, with ablations examining codebook size and key architectural choices. The results show strong performance across tasks, while larger codebooks trade off class-conditional generation quality.

  • Multimodal Understanding: 7.2% average improvement over LLaVA-1.5 13B shows TokenFlow-XL with Qwen-2.5-14B surpassing the baseline across multimodal understanding benchmarks.With the same Vicuna 13B backbone, TokenFlow-XL instead exceeds LLaVA-1.5 13B by 1.7% on average.
  • Visual Generation: 0.55 GenEval at 256×256 surpasses Stable Diffusion v2.1 and PixArt-alpha under a 25-step multi-step sampling scheme.The model is reported to require significantly fewer generation steps than competing approaches.
  • Visual Generation: TokenFlow consistently outperforms VAR under identical text-to-image training configurations and dataset settings across all reported benchmark metrics.The comparison further supports the effectiveness of the unified tokenization approach.
  • Codebook Size: Increasing codebook size improves multimodal understanding and reconstruction, but sizes above 32,768 slightly degrade class-conditional generation.The authors adopt 32,768 for text-to-image generation experiments because larger codebooks make autoregressive generation more complex.

5. Conclusion

TokenFlow unifies multimodal understanding and generation with a dual-codebook tokenizer that decouples semantic and pixel-level learning while aligning them through shared mapping. Experiments report strong reconstruction, understanding, and generation performance with low training cost and fewer inference steps.

  • 5. Conclusion: TokenFlow uses dual-codebook, decoupled yet aligned feature learning to address the different visual granularities required by understanding and generation.The shared mapping aligns semantic and pixel-level representations within one tokenizer.
  • 5. Conclusion: The experiments report superior reconstruction, state-of-the-art multimodal understanding, and competitive visual generation with minimal training costs and substantially fewer inference steps.The conclusion positions TokenFlow as a promising foundation tokenizer for vision-language systems.

Generation

The supplementary experiments compare VQ tokenizer representations and reconstruction from quantized semantic features. They use ImageNet-based evaluations to examine visual similarity and the recoverability of original images.

  • Generation: Visual comparisons process 50,000 ImageNet-1k validation images through VQKD, VQGAN, and TokenFlow encoders before nearest-codebook matching.Average pooling produces a 1 × 1 representation, and closest indices are identified using l2 distance.
  • Generation: Reconstructing images from VQKD quantized features retains the original VQKD encoder and quantizer while adding a TokenFlow-style pixel decoder.The added decoder is trained on ImageNet-1K for 100 epochs.

A.2. Tokenizer Training Details

Tokenizer training uses multiple TokenFlow variants with different resolutions and semantic teacher models. Additional analysis compares codebook cluster distributions and reports high utilization from the shared mapping design.

  • A.2. Tokenizer Training Details: TokenFlow-B, -L, and -XL primarily differ in input resolution and semantic teacher model, using CLIP ViT-B/14, ViTamin-XL, and SigLIP-SO400M respectively.The variants share common training hyperparameters.
  • A.2. Tokenizer Training Details: With an 8,192-entry codebook, TokenFlow exhibits a smoother cluster-size distribution than VQKD and VQGAN on 50,000 ImageNet-1k validation images.The shared mapping design is attributed with learning joint semantic and pixel-level distributions.

B.1. Additional Ablation Study

The ablations examine sampling strategy, model size, input features, and decoder finetuning for TokenFlow’s generation and understanding capabilities. Progressive sampling and larger models improve generation, while final-scale features favor multimodal understanding and decoder enhancement preserves fine details.

  • Sampling Strategy: +0.039 GenEval and +0.084 ImageReward result from two-step sampling over single-step sampling.The configuration uses top-k=[1200, 1] and top-p=[0.8, 0].
  • Sampling Strategy: 10.2% and 14.3% improvements over traditional single-step sampling are achieved by the three-step strategy.The sampling space narrows progressively as top-k changes from 1200 to 100 to 1.
  • Model Size: The 7B model significantly outperforms the 1B model under identical sampling and training-data configurations.The 1B model remains weaker even when trained for twice as many epochs.
  • Input Strategy: Final-scale features consistently outperform full-scale and full-scale residual features across multimodal understanding benchmarks.Using semantic features alone does not improve overall understanding performance.
  • Tokenizer Decoder Finetuning: Finetuning a decoder with doubled residual layers and channel dimensions improves reconstruction fidelity while preserving learned visual token mappings.All other components remain frozen during finetuning.

B.2. More Analysis of TokenFlow

Further analyses evaluate TokenFlow’s shared mapping, decoder capacity, codebook utilization, and comparison with semantic teachers. The results indicate strong codebook usage, improved fine-detail reconstruction, and a diminishing teacher gap at higher resolutions.

  • Decoder Capacity: The enhanced decoder preserves fine-grained details more effectively, particularly facial details and textual elements.The comparison contrasts the base pixel decoder with an enhanced decoder at 2× capacity.
  • Joint Distribution Learning: 7161/8192 clusters are non-empty for TokenFlow, compared with 2.5% for VQGAN and 27.1% for VQKD.The models use identical codebook sizes of 8,192 tokens in the comparison.
  • Distance Balancing: The quantization index balances semantic and pixel distances through arg min_i(d_sem,i + w_dis · d_pix,i).Commit and perceptual losses jointly contribute when semantic and pixel distances favor different indices.
  • Semantic Teacher Comparison: The performance gap to continuous semantic teachers narrows from 6.3% at 224×224 to 2.9% at 384×384.The analysis attributes this improvement to more discrete tokens and additional scales supplementing residual features.

B.3. More Visual Generation Results

TokenFlow produces visual generation results across varied styles, subjects, and scenarios. The section presents these qualitative outputs alongside the evaluation setup for GenEval and DPG-Bench.

  • Quantitative Results: GenEval and DPG-Bench are used to report quantitative generation quality in Table 10.GenEval results use GPT-4V rewriting, while released checkpoints support comparisons for DPG-Bench and VAR.
  • Qualitative Results: Additional visual results demonstrate generation across various styles, subjects, and scenarios.These examples are presented in Figure 12 at 256×256 resolution.

C. Limitation and Future Work

The paper identifies remaining gaps in discrete-token understanding and has not yet trained a fully unified model across understanding and generation. It also illustrates TokenFlow’s ability to combine semantic and low-level visual similarity.

  • Limitation: TokenFlow remains 2.9% below its continuous semantic teacher at 384×384 because of vector quantization distillation.Text alignment loss is suggested as one possible improvement.
  • Future Work: The work validates TokenFlow separately on multimodal understanding and visual generation rather than as a fully unified model.Joint training on interleaved vision-language data is identified as a future extension.
  • Shared Mapping: TokenFlow combines semantic and low-level similarity by mapping patches with both properties to shared indices.This allows images with different backgrounds to receive different indices while retaining semantic grouping.
Loading 2412.03069v2…