ArcVQ-VAE: A Spherical Vector Quantization Framework with ArcCosine Additive Margin

Vector Quantized Variational Autoencoder (VQ-VAE) has become a fundamental framework for learning discrete representations in image modeling. However, VQ-VAE models must tokenize entire images using a finite set of codebook vectors, and this capacity limitation restricts their ability to capture rich and diverse representations. In this paper, we propose ArcCosine Additive Margin VQ-VAE (ArcVQ-VAE), a novel vector quantization framework that introduces a spherical angular-margin prior (SAMP) for the codebook of a conventional VQ-VAE. The proposed SAMP consists of Ball-Bounded Norm Regularization, which constrains all codebook vectors within a time-dependent Euclidean ball, and ArcCosine Additive Margin Loss, which encourages greater angular separability among latent vectors. This formulation promotes more discriminative and uniformly dispersed latent representations within the constrained space, thereby enabling codebook vectors to capture richer information and leading to improved codebook utilization. Experimental results on standard image reconstruction and generation tasks show that ArcVQ-VAE outperforms baseline models in terms of reconstruction accuracy, representation diversity, and sample quality.

Many image AI systems learn by converting an image into a collection of smaller internal patterns. These patterns help the model remember and recreate important visual information. However, because the number of available patterns is limited, some patterns may be used too often while others may be rarely used. This can make it harder for the model to capture the full variety of information in images. In this paper, we propose ArcVQ-VAE, a method that helps an image model use its internal patterns more effectively. The main idea is to encourage these patterns to be more clearly separated from one another while keeping them within a controlled space. As a result, the model can learn more diverse and informative image representations. Experiments on image reconstruction and generation show that our method improves image quality, representation diversity, and the use of the model’s available pattern set.