MergeVQ: Un Framework Unificato per la Generazione e la Rappresentazione Visiva con Fusione e Quantizzazione di Token Disaccoppiati

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

April 1, 2025

Autori: Siyuan Li, Luyuan Zhang, Zedong Wang, Juanxi Tian, Cheng Tan, Zicheng Liu, Chang Yu, Qingsong Xie, Haonan Lu, Haoqian Wang, Zhen Lei

cs.AI

Abstract

Il Masked Image Modeling (MIM) con Vector Quantization (VQ) ha ottenuto un grande successo sia nel pre-training auto-supervisionato che nella generazione di immagini. Tuttavia, la maggior parte dei metodi esistenti fatica a bilanciare il compromesso nello spazio latente condiviso tra qualità della generazione, apprendimento di rappresentazioni ed efficienza. Per spingere al limite questo paradigma, proponiamo MergeVQ, che incorpora tecniche di fusione di token nei modelli generativi basati su VQ per colmare il divario tra generazione di immagini e apprendimento di rappresentazioni visive in un'architettura unificata. Durante il pre-training, MergeVQ disaccoppia la semantica top-k dallo spazio latente tramite il modulo di fusione di token dopo i blocchi di self-attention nell'encoder, per poi applicare la Look-up Free Quantization (LFQ) e l'allineamento globale, e recupera i dettagli fini attraverso il cross-attention nel decoder per la ricostruzione. Per la generazione in seconda fase, introduciamo MergeAR, che esegue la compressione della KV Cache per una predizione efficiente in ordine raster. Esperimenti estensivi su ImageNet verificano che MergeVQ, come modello generativo AR, raggiunge prestazioni competitive sia nei compiti di apprendimento di rappresentazioni visive che di generazione di immagini, mantenendo al contempo un'efficienza favorevole dei token e una velocità di inferenza. Il codice e il modello saranno disponibili su https://apexgen-x.github.io/MergeVQ.

English

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 shared latent space for generation quality vs. representation learning and efficiency. To push the limits of this paradigm, we propose MergeVQ, which 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 the 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. The code and model will be available at https://apexgen-x.github.io/MergeVQ.