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

TL;DR

ArcVQ-VAE introduces a spherical angular-margin prior with regularization and margin loss to enhance latent representation diversity and codebook utilization in vector quantized autoencoders.

Contribution

It proposes a novel spherical angular-margin prior framework that improves latent space discrimination and diversity in VQ-VAE models.

Findings

Achieves competitive image reconstruction and generation performance.

Promotes more discriminative and uniformly dispersed latent representations.

Improves codebook utilization and representation diversity.

Abstract

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 improving effective latent-space coverage and leading to improved codebook utilization. Experimental results on standard image reconstruction and generation tasks show that ArcVQ-VAE achieves competitive performance against baseline models in terms of reconstruction accuracy, representation diversity, and sample quality. The code is available at: https://github.com/goals4292/ArcVQ-VAE