GaussianToken: An Effective Image Tokenizer with 2D Gaussian Splatting

TL;DR

GaussianToken addresses the bottleneck of discrete image tokenizers by introducing 2D Gaussian Splatting to enrich the discrete latent space without increasing token count. The method combines a Gaussian Embedding module with 2D Gaussian quantization and splatting, integrated into a VQ-VAE–style pipeline, and trained with reconstruction, commitment, and GAN losses. It demonstrates improved image reconstruction quality on CIFAR, Mini-ImageNet, and ImageNet-1K, with faster convergence and better utilization of the codebook compared to strong VQ-based baselines. The work expands the representational capacity of discrete tokens and lays groundwork for future exploration of downstream generation tasks.

Abstract

Effective image tokenization is crucial for both multi-modal understanding and generation tasks due to the necessity of the alignment with discrete text data. To this end, existing approaches utilize vector quantization (VQ) to project pixels onto a discrete codebook and reconstruct images from the discrete representation. However, compared with the continuous latent space, the limited discrete codebook space significantly restrict the representational ability of these image tokenizers. In this paper, we propose GaussianToken: An Effective Image Tokenizer with 2D Gaussian Splatting as a solution. We first represent the encoded samples as multiple flexible featured 2D Gaussians characterized by positions, rotation angles, scaling factors, and feature coefficients. We adopt the standard quantization for the Gaussian features and then concatenate the quantization results with the other intrinsic Gaussian parameters before the corresponding splatting operation and the subsequent decoding module. In general, GaussianToken integrates the local influence of 2D Gaussian distribution into the discrete space and thus enhances the representation capability of the image tokenizer. Competitive reconstruction performances on CIFAR, Mini-ImageNet, and ImageNet-1K demonstrate the effectiveness of our framework. Our code is available at: https://github.com/ChrisDong-THU/GaussianToken.

Figures (5)

• Figure 1: Motivation of GaussianToken. The Variational AutoEncoder (VAE) adopts a continuous space to represent image features, which fails to align with discrete modalities. In addition, Vector Quantization (VQ) approaches employ discrete codebooks for nearest-neighbor matching of image features, while the representation space is restricted to the codebook size. Differently, GaussianToken performs local adaptive learning on discrete features with continuous Gaussian parameters, enlarging the representation space under the discretized basis.
• Figure 2: Illustration of the proposed Gaussian Embedding module. The information flow traverses through $B$ iterative blocks (marked in red arrow lines), with each block sequentially encompassing three primary operations: self-attention, cross-attention, and refinement, alternately updating the featured 2D Gaussian anchors under the guidance of visual information.
• Figure 3: The overall framework of GaussianToken. We process image features using the Gaussian Embedder, obtaining intrinsic spatial parameters (position, rotation angle, and scaling factor) as well as feature coefficients. We perform nearest neighbor matching with a discrete codebook for the feature coefficients to achieve the quantization process and concatenate the quantized results with the spatial parameters. Subsequently, we adopt a 2D splatting module and an image decoder to generate the corresponding reconstructed image. Finally, we impose quality constraints on the reconstructed image through an additional discriminator. The overall constraints include the image reconstruction loss, the quantization-related VQ loss and commitment loss, as well as the GAN loss for image quality.
• Figure 4: Frequency of vector usage in the codebook. The red dashed line represents the 20% frequency cumulative line.
• Figure 5: Reconstruction comparison across different tokenizers. Compared to the baseline VQGAN esser2021taming and LlamaGen sun2024autoregressive, our proposed model substantially enhances reconstruction quality (highlighted in the white box) while utilizing the same number of image tokens (256 tokens for a 256x256 resolution).