MergeVQ: A Unified Framework for Visual Generation and Representation with Token Merging and Quantization

Abstract

Masked Image Modeling (MIM) with Vector Quantization (VQ) has achieved great success in both self-supervised pre-training and image generation. However, most existing methods struggle to address the trade-off in the shared latent space for generation quality vs. representation learning and efficiency. To push the limits of this paradigm, we propose MergeVQ that incorporates token merging techniques into VQ-based generative models to bridge the gap between image generation and visual representation learning in a unified architecture. During pre-training, MergeVQ decouples top-k semantics from latent space with the token merge module after self-attention blocks in the encoder for subsequent Look-up Free Quantization (LFQ) and global alignment, and recovers their fine-grained details through cross-attention in the decoder for reconstruction. As for second-stage generation, we introduce MergeAR which performs KV Cache compression for efficient raster-order prediction. Extensive experiments on ImageNet verify that MergeVQ as an AR generative model achieves competitive performance in both visual representation learning and image generation tasks, while maintaining favorable token efficiency and inference speed.

Introduction

Vector Quantization (VQ) has emerged as a widely adopted technique in visual tasks. Nevertheless, existing methods often encounter challenges that impede their effectiveness. In response, we present MergeVQ, a revolutionary solution designed to overcome these limitations. Our approach recognizes the complementary nature of representation learning and generation, and it leverages this synergy by decoupling and recovering semantics during the encoding and reconstruction processes.

The contributions of MergeVQ are multifaceted. Firstly, it effectively decouples semantic tokens and the source matrix, enabling a more refined and efficient representation of visual data. Secondly, it addresses the longstanding challenges associated with VQ training, leading to improved stability and performance. Finally, MergeVQ achieves state-of-the-art results in both visual generation and representation learning tasks, demonstrating its superiority over existing methods.

Learning Paradigm

3.1 MergeVQ Framework

The MergeVQ framework is a comprehensive solution that encompasses token merge encoding, quantization, token recovery, and reconstruction. Each component plays a crucial role in enabling efficient and effective visual generation and representation learning.

Token Merge Encoding: Given an image \(X \in \mathbb{R}^{H \times W \times 3}\), we employ a two - stage downsampling encoder \(\varepsilon\). First, a CNN layer \(\varepsilon_{1}\) extracts features, producing a feature map \(Z \in \mathbb{R}^{\frac{H}{f} \times \frac{W}{f} \times d}\), where \(f\) is the downsampling factor and \(d\) denotes the channel number. The feature \(Z\) is then flattened into an \(L\) - length token sequence \(Z_{L} \in \mathbb{R}^{L \times d}\) as \(Z_{L}=\mathcal{E}_{1}(X)\). Subsequently, we employ attention with the merging operation, denoted as \(\varepsilon_{2}\), for second - stage extraction. During this process, we obtain a shorter sequence \(Z_{K} \in \mathbb{R}^{K \times d}\) along with its source matrix \(S \in \mathbb{R}^{K \times L}\) that preserves the spatial relationships of the sequence. This process is expressed as \(Z_{K}, S=\mathcal{E}_{2}(Z_{L})\). To ensure that \(Z_{K}\) possesses rich semantics, we concurrently impose global alignment on \(Z_{K}\).

Quantization: We utilize \(LFQ\) as MergeVQ’s quantization module to minimize the loss of details. Specifically, the codebook is reduced to an integer set and can be denoted as: \(C = \times_{i = 1}^{N} \{-1, 1\}\) with \(|C| = 2^{N}\). Thus, the quantization can be summarized as follows: \(z_{K i} = \text{sign}(z_{K i}) = -1 \cdot \mathbb{I}(z_{K i} < 0) + \mathbb{I}(z_{K i} > 0)\), where \(z_{K i}\) denotes the \(i\) - th vector in \(K\) semantic tokens \(Z_{K}\). Then the index of the quantized feature \(z_{m i}\) is formulated as \(\text{Index}(z_{K i}) = \sum_{j = 1}^{N} 2^{k - 1} \cdot \mathbb{I}(z_{K i j} > 0)\). Finally, we obtain the quantized semantic tokens, denoted as \(Z_{K q}\): \(Z_{K q}=\mathcal{Q}(Z_{K}, \mathcal{C})\).

Token Recovery and Reconstruction: We first perform token - level reconstruction with the recovery module \(R(\cdot, \cdot)\) and source matrix \(S\), which yields a new \(L\) - length \(\hat{Z}_{L}\) as \(\hat{Z}_{L}=\mathcal{R}(Z_{q K}, S)\). This sequence is then fed into the decoder \(D\) for pixel - level reconstruction, which can be described as \(\hat{X}=\mathcal{D}(\hat{Z}_{L})\).

3.2 Harmonize Reconstruction and Representation

Inspired by Masked Image Modeling in representation learning, we employ Token Merge and Source Recovery to seamlessly integrate representation learning into the MergeVQ framework. Additionally, we impose global alignment constraints to further enhance the performance and stability of the model.

Efficient Generation

4.1 MergeAR with KV Cache Compression

To achieve efficient autoregressive generation, we introduce MergeAR. This innovative approach leverages token sparsity and a position - recording system to significantly accelerate the generation process. During training, we sample a merge ratio \(r\), introduce a Merge Instruction Token \(M\), and construct a causal mask. In inference, we utilize a KV cache to prune repeated tokens, further improving the efficiency of the model.

4.2 Randomized Auto - regressive with Source Recovery

MergeVQ can be effectively implemented using the RandAR generative framework. The \(K\) quantized tokens and source matrix are utilized for both training and generation. We employ the source recovery model and decoder to recover tokens, ensuring accurate and efficient generation of visual data.

Experiments

5.1 Implementation Details

We offer three distinct versions of MergeVQ to cater to different application scenarios: MergeVQ (G) for pure generation, MergeVQ (G + R) for both generation and representation, and MergeVQ (R) for representation learning only. Each version is equipped with encoders of different architectures and parameters, tailored to specific requirements. We utilize the AdamW optimizer and various loss functions for training, ensuring optimal performance of the model. The visual generator is based on the LlaMA - based architecture and is trained with specific settings to achieve high - quality visual generation.

5.2 Self - supervised Pre - training

To evaluate the performance of self - supervised pre - trained models, we conducted linear probing and end - to - end fine - tuning experiments on ImageNet - 1K. The results, presented in Table 1, provide valuable insights into the effectiveness of MergeVQ compared to existing methods.

5.3 Image Generation

We conducted comprehensive image generation experiments to assess the performance of MergeVQ. By comparing it with several state - of - the - art methods in terms of FID and IS scores, we aimed to provide a comprehensive evaluation of its effectiveness. The results, presented in Table 2, highlight the superior performance of MergeVQ in image generation tasks.

Contributions

Please feel free to raise issues or submit pull requests to contribute to our codebase.

Citation

    @article{li2025mergevq,
      title={MergeVQ: A Unified Framework for Visual Generation and Representation with Token Merging and Quantization},
    }