Vector quantization approaches (VQ-VAE, VQ-GAN) learn discrete neural representations of images, but these representations are inherently position-dependent: codes are spatially arranged and contextually entangled, requiring autoregressive or diffusion-based priors to model their dependencies at sample time. In this work, we ask whether positional information is necessary for discrete representations of spatially aligned data. We propose the permutation-invariant vector-quantized autoencoder (PI-VQ), in which latent codes are constrained to carry no positional information. We find that this constraint encourages codes to capture global, semantic features, and enables direct interpolation between images without a learned prior. To address the reduced information capacity of permutation-invariant representations, we introduce matching quantization, a vector quantization algorithm based on optimal bipartite matching that increases effective bottleneck capacity by 3.5× relative to naive nearest-neighbour quantization. The compositional structure of the learned codes further enables interpolation-based sampling, allowing synthesis of novel images in a single forward pass. We evaluate PI-VQ on CelebA, CelebA-HQ and FFHQ, obtaining competitive precision, density and coverage metrics for images synthesised with our approach. We discuss the trade-offs inherent to position-free representations, including separability and interpretability of the latent codes, pointing to numerous directions for future work.
Figures & tables
Method
Discrete?
O(1) -time sampling?
Interpolation?
BigGAN [ 1 ]
-
✔
-
VDVAE [ 2 ]
-
-
-
DDGAN [ 3 ]
-
✔
-
StyleGAN-XL [ 4 ]
-
✔
✔
LDM-4 [ 5 ]
-
-
-
ADM-IP [ 6 ]
-
-
-
Table 1: Comparison of the capabilities of our method with the state-of-the-art. “ O(1) -time sampling” indicates methods which can sample an image with a single forward-pass of a neural network.
Figure 1: Visualisation of (approximately) smooth interpolations on CelebA 64x64 (source images to left and right). The discrete and permutation-invariant nature of the latents allows us to generate multiple, equally plausible paths between two images.
Figure 2: The encoder and decoder components of the permutation-invariant autoencoder (quantization step omitted for clarity). Positional embeddings are not applied to the latent codes before decoding: as a result, the decoder is invariant to permutation of the latent codes.
Figure 3: Matching quantization: Our novel approach to vector quantization (right) ensures that no two embedding vectors are mapped to the same codebook element in a given image, effectively minimising redundancy in the discrete representation.
Figure 4: Information capacity (in bits) in terms of representation length L of three approaches with codebook size K=4096 : standard VQ-VAE with no permutation invariance; PI-VQ with nearest neighbour quantization ( Kimg=49 ); PI-VQ with proposed matching quantization.
Quantization method
Kimg (max) ↑
Capacity (bits) ↑
Nearest
49
614
Matching (ours)
512
2221
Table 2: Upper bound of latent bottleneck information capacity: comparison between nearest-neighbour and matching quantization for models trained on CelebA-HQ. Kimg is included as the maximum per-image codebook usage observed over the entire dataset. Capacity is rounded to the nearest bit.
Method
FID ↓
P ↑
R ↑
D ↑
C ↑
FFHQ 256x256
StyleGAN-XL [ 4 ]
2.2
0.80
0.39
0.86
0.73
StyleGAN2 [ 55 ]
3.8
0.69
0.40
1.12
0.80
VQ-GAN (UT) [ 13 ]
7.1
0.69
0.48
1.06
0.77
VQ-GAN (TT) [ 12 ]
9.6
0.64
0.29
0.89
0.59
VDVAE [ 2 ]
28.5
0.59
0.20
0 .80
0.50
Table 3: Comparison with state-of-the-art in FID, precision (P), recall (R), density (D), and coverage (C) scores for image quality. Best result for each dataset and metric is in bold. *Averages are taken over only the results where both SOTA and our results are available for comparison.
Figure 5: Logistic regression accuracy with error bars (5-way cross-validated) for predicting ground-truth FFHQ annotations based on learned permutation-invariant representations. For each attribute, we compare against the baseline accuracy (always predicting the majority class).
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
Jaeyung Kim, YoungJoon Yoo
Department of Artificial Intelligence, Chung-Ang University, Seoul, Republic of Korea · SNUAILAB, Seoul, Republic of Korea
Vector Quantized Variational Autoencoders (VQ-VAEs) are fundamental to modern generative modeling, yet they often suffer from training instability and "codebook collapse" due to the inherent coupling of representation learning and discrete codebook optimization. In this paper, we propose VP-VAE (Vector Perturbation VAE), a novel paradigm that decouples representation learning from discretization by eliminating the need for an explicit codebook during training. Our key insight is that, from the neural network's viewpoint, performing quantization primarily manifests as injecting a structured perturbation in latent space. Accordingly, VP-VAE replaces the non-differentiable quantizer with distribution-consistent and scale-adaptive latent perturbations generated via Metropolis--Hastings sampling. This design enables stable training without a codebook while making the model robust to inference-time quantization error. Moreover, under the assumption of approximately uniform latent variables, we derive FSP (Finite Scalar Perturbation), a lightweight variant of VP-VAE that provides a unified theoretical explanation and a practical improvement for FSQ-style fixed quantizers. Extensive experiments on image and audio benchmarks demonstrate that VP-VAE and FSP improve reconstruction fidelity and achieve substantially more balanced token usage, while avoiding the instability inherent to coupled codebook training.
The effectiveness of modern visual representation learning and autoregressive models critically depends on vector quantization (VQ), which discretizes continuous feature representations using a learnable codebook. Despite its widespread use, existing VQ methods often suffer from training instability and codebook collapse, arising from gradient mismatch induced by the straight-through estimator and the under-utilization of code vectors. In this work, we show that both issues can be traced to a fundamental mismatch between the distributions of feature vectors and code vectors, leading to inefficient representation and information loss. Building on this observation, we propose a distributional matching framework for vector quantization. We introduce principled criteria for desirable VQ behavior and demonstrate through theoretical analysis and empirical evaluation that aligning feature and code vector distributions provides a unifying mechanism for mitigating training instability and codebook collapse. We instantiate this framework using a Wasserstein-based objective with an efficient closed-form under a mild Gaussian approximation, and further show that a nonparametric alternative based on maximum mean discrepancy yields comparable performance. Extensive experiments on visual tokenization benchmarks support the effectiveness and robustness of the proposed approach.
Xianghong Fang, Litao Guo, Hengchao Chen +8
University of Toronto · The Hong Kong University of Science and Technology · Lehigh University +2