Residual Quantization VAE (RQ-VAE)

Updated 22 January 2026

RQ-VAE is a hierarchical autoencoder that stacks residual quantizers to overcome standard VQ limitations and enhance latent representation.
It employs a unified codebook across multiple quantization stages, drastically reducing code sequence length while maintaining tractable rate-distortion performance.
RQ-VAE has been effectively applied in image, text, and 3D motion synthesis, achieving state-of-the-art reconstruction fidelity and sampling speed.

Residual Quantization Variational Autoencoder (RQ-VAE) is a hierarchical vector quantized autoencoder framework designed to address the limitations of standard vector quantization (VQ-VAE) in learning flexible, highly efficient discrete representations for generative modeling tasks. Unlike single-layer VQ-VAEs, RQ-VAE achieves exponentially greater latent expressivity with fixed-size codebooks by stacking residual quantizers, enabling high-fidelity compression—including dramatic reductions in code sequence length—while maintaining tractable rate-distortion and robust codebook utilization (Lee et al., 2022).

1. Mathematical Framework and Quantization Principle

The defining feature of RQ-VAE is its residual quantization of feature maps. Let $z \in \mathbb{R}^n$ denote a continuous latent vector (e.g., one cell of the encoder’s spatial feature map). Given codebook $C = \{e(k)\in\mathbb{R}^n : k=1...K\}$ and quantization depth $D$ , residual quantization proceeds as:

$r_0 = z$ , $\hat{z}^{(0)} = 0$
For each stage $d=1,\dots,D$ $d = 1, \dots, D$ :
- $r_d = z - \hat{z}^{(d-1)}$
- $k_d = \arg\min_{k\in\{1...K\}}\|r_d - e(k)\|^2$
- $\hat{z}_d = e(k_d)$
- $\hat{z}^{(d)} = \hat{z}^{(d-1)} + \hat{z}_d$

At the feature-map level, for input $X\in \mathbb{R}^{H_0 \times W_0 \times 3}$ and encoder $E$ , the encoded feature map $Z = E(X) \in \mathbb{R}^{H \times W \times n_z}$ (with $H=W=H_0/f$ for downsampling factor $f$ ) is quantized independently at each spatial location to a stack of codes $M_{h,w} = (k_1,\dots,k_D) \in \{1,...,K\}^D$ , yielding reconstructed features $\hat{Z}_{h,w} = \sum_{d=1}^D e(M_{h,w,d})$ (Lee et al., 2022).

The RQ-VAE objective generalizes classical VQ-VAE losses to depth $D$ :

Reconstruction loss: $L_{\text{rec}} = \|X - G(\hat{Z}^{(D)})\|^2_2$
Commitment loss: $L_{\text{commit}} = \sum_{d=1}^D \|Z - \mathrm{sg}[\hat{Z}^{(d)}]\|^2_2$
Total loss: $L_{\text{RQ-VAE}} = L_{\text{rec}} + \beta L_{\text{commit}}$

where $\mathrm{sg}$ is the stop-gradient operator, and $G$ is the decoder.

Codebook learning is performed via exponential moving average (EMA) updates of assigned features, and a straight-through estimator enables differentiable quantization, setting $\partial \hat{z}/\partial z = 1$ at the embedding (Lee et al., 2022). Optional adversarial and perceptual objectives (e.g., patch-GAN, deep feature loss) can further enhance output sharpness.

2. Architecture, Codebooks, and Hierarchical Structure

A typical RQ-VAE architecture consists of:

Encoder $E$ : stack of convolutional residual blocks, with aggressive spatial downsampling ( $f=32$ for $256\times 256$ inputs, producing $8\times8$ feature maps). Hidden channel dimension $n_z\approx256$ .
Decoder $G$ : symmetric upsampling/residual block network.
Codebook $C$ : a single shared codebook of size $K$ (e.g., $K=16,384$ in high-res image experiments), used jointly by all $D$ quantization stages. Only one set of embeddings need be maintained; this ensures every codeword is used at every residual stage.
Quantization depth $D$ : Empirically $D=4$ suffices for lossless $8\times8$ compression; increasing $D$ (see Table 1) further reduces distortion, but increases AR decoding cost.

Quantization Depth $D$	Codebook Size $K$	Feature Map	rFID
1 (VQ-GAN, 8x8)	$128$k	$8\times8$	17.1
4 (RQ-VAE, 8x8)	$16$k	$8\times8$	4.7
8	$16$k	$8\times8$	2.7
16	$16$k	$8\times8$	1.8

This exponential latent space ( $K^D$ ) enables high-rate, high-fidelity quantization even at heavily reduced spatial resolutions, supporting ultra-fast autoregressive modeling (Lee et al., 2022).

Variants such as HR-VQVAE (Adiban et al., 2022), MoSa's hierarchical RQ-VAE (Liu et al., 3 Nov 2025), and stochastic RSQ-VAE (Takida et al., 2023) generalize this principle by (i) supporting per-layer codebooks, (ii) allowing different residual scales, and (iii) employing Bayesian entropy-based regularization to mitigate codebook collapse.

3. Rate–Distortion Analysis and Codebook Utilization

The rate-distortion advantage of RQ-VAE arises from the hierarchical composition of quantizers. A single $K$ -codebook VQ partitions $\mathbb{R}^n$ into $K$ cells, subject to the distance-minimizing error $D(K)$ . In contrast, an $L$ -layer residual-quantized model with depth $D$ achieves $K^D$ unique code sums, substantially increasing representational power without an exponential codebook size (Lee et al., 2022).

For image encoding:

VQ-VAE (16x16x1, $K=16,384$ ): rFID $\approx 4.3$
RQ-VAE (8x8x4, $K=16,384$ ): rFID $\approx 4.7$ , matching the fidelity but reducing the autoregressive sequence length from $256$ to $64$.
Further increasing $D$ to $8$ or $16$ pushes rFID as low as $2.7$ and $1.8$, respectively.

Ablations consistently demonstrate that, for fixed $K$ , increasing $D$ is vastly more effective at reducing distortion than growing the codebook size (Lee et al., 2022).

Later extensions provide additional codebook usage guarantees. HR-VQVAE (Adiban et al., 2022) employs hierarchical codebook conditioning and per-layer contrastive targets, while HQ-VAE (Takida et al., 2023) incorporates variational entropy penalties to maximize code utilization and eliminate code collapse even at large $D$ .

4. Encoding and Decoding Algorithms

The full RQ-VAE encoding process proceeds as:

Encode:
- For each spatial location $(h,w)$ $(h, w)$ in $Z = E(X)$ $Z = E (X)$ :
  - Initialize $r \leftarrow Z_{h,w}$ , $\hat{z}^{(0)} \leftarrow 0$ .
  - For $d=1\ldots D$ :
  - $k_d \leftarrow \arg\min_k \|r - e(k)\|^2$
  - $\hat{z}_d \leftarrow e(k_d)$
  - $\hat{z}^{(d)} \leftarrow \hat{z}^{(d-1)} + \hat{z}_d$
  - $r \leftarrow r - \hat{z}_d$
  - Store code indices $M_{h,w,d} = k_d$
Decode:
- For all $(h,w)$ , sum the codebook embeddings for each stack of codes to form $\hat{Z}_{h,w} = \sum_{d=1}^D e(M_{h,w,d})$
- Image reconstruction $\hat{X} = G(\hat{Z})$

Gradient flow:

Losses ( $L_{\text{rec}}, L_{\text{commit}}$ ) propagate gradients through $G$ and $E$ .
The straight-through estimator routes gradients from codeword assignments to the encoder outputs.
Codebook $C$ updates occur via EMA of encoder outputs (Lee et al., 2022).

Extensions in MoSa (Liu et al., 3 Nov 2025) introduce scale-varying downsampling/upsampling at each quantization stage, with reconstructed latents formed by upsampling the dequantized embeddings back to full length prior to computing the residual. This supports arbitrarily coarse-to-fine token hierarchies, beneficial for tasks such as temporal motion synthesis.

5. Applications, Empirical Results, and Comparisons

RQ-VAE frameworks are deployed in multiple domains:

Image generation: RQ-VAE enables efficient AR modeling by compressing $256\times256$ images to $8\times8\times D$ discrete codes. When coupled with an RQ-Transformer, this yields FID and IS statistics superior to VQ-GAN and competitive methods, with 4-7× faster sampling (Lee et al., 2022).
Class-conditional generation (ImageNet 256 $^2$ ): RQ-Transformer ($1.4$B params, no rejection) achieves FID=$11.6$, IS=$112.4$ (vs. VQ-GAN $15.8$/74.3). With rejection sampling, FID=$4.45$, IS=$326$ (state-of-the-art AR).
Text-conditional generation: On CC-3M, RQ-Transformer (650M params) yields FID=$12.3$, CLIP-sim=$0.26$; VQ-GAN yields $28.9$/0.20 (Lee et al., 2022).
3D motion synthesis: MoSa’s RQ-VAE compresses 64-frame motion windows into 10 hierarchical code groups at variable scales, achieving state-of-the-art FID ($0.06$ on Motion-X, vs. $0.20$ for MoMask) and $27\%$ lower inference time (Liu et al., 3 Nov 2025).
General discrete representation learning: Hierarchical and stochastic variants (HR-VQVAE (Adiban et al., 2022), HQ-VAE (Takida et al., 2023)) outperform VQ-VAE/VQ-VAE-2 on reconstruction MSE and FID across FFHQ, ImageNet, CIFAR-10, MNIST, and audio datasets.

6. Variants, Theoretical Advances, and Mitigation of Codebook Collapse

Several extensions sharpen the RQ-VAE paradigm:

HR-VQVAE (Adiban et al., 2022): Implements multi-layer residual VQ with per-layer codebooks and hierarchical linkage. Layer $l+1$ 's codebook eligibility is indexed by codes at layer $l$ , restricting the effective code dictionary per location and avoiding collapse.
MoSa RQ-VAE (Liu et al., 3 Nov 2025): Utilizes nonlinear scale schedules and MTPS (Multi-scale Token Preservation Strategy) to align AR inference with hierarchical quantization groups, allowing transformer models to sample all tokens at each scale in parallel with $Q=10$ total AR steps.
HQ-VAE (RSQ-VAE instance) (Takida et al., 2023): Recasts RQ-VAE as a Bayesian model, introducing stochastic (Gaussian-categorical) code assignment, entropy regularization for codebook usage, and joint learning of quantization strengths. This stochastically annealed training cures codebook and layer collapse, eliminates heuristic penalties (no $\beta$ , EMA, or resets), and converges to full codebook usage with improved RMSE/LPIPS/SSIM.

A plausible implication is that variational residual quantization (e.g., HQ-VAE) should increasingly supplant heuristic-hardened, deterministic approaches for settings demanding both deep hierarchies and stable code allocation.

7. Significance and Outlook

The RQ-VAE family provides an expressive and efficient method for learning high-rate discrete representations with tractable quantization complexity, substantially increasing downstream generative modeling capacity and speed. Its flexible architecture supports domain generalization (images, motions, audio) and is compatible with both deterministic and variational Bayesian training. By enabling discrete representations whose cardinality grows exponentially in quantization depth (without codebook inflation), RQ-VAE achieves state-of-the-art rates on challenging benchmarks and is robust to codebook collapse, a frequent failure mode in non-residual models (Lee et al., 2022, Adiban et al., 2022, Takida et al., 2023, Liu et al., 3 Nov 2025). This positions RQ-VAE and its variants as fundamental to the evolution of highly scalable discrete latent variable models for modern deep generative applications.