HR-VQVAE: Hierarchical Residual VQVAE

• The paper introduces HR-VQVAE, which integrates hierarchical quantization with residual learning to capture multi-level image features while mitigating codebook collapse.
• It employs a tree-structured codebook and contrastive objective to enforce non-redundant, complementary latent representations, significantly improving reconstruction fidelity on datasets like FFHQ and ImageNet.
• The model's design scales exponentially in representational capacity while retaining linear decoding complexity, offering rapid inference and robust performance compared to existing VQ-VAE variants.

Hierarchical Residual Learning Vector Quantized Variational Autoencoder (HR-VQVAE) is a discrete latent variable generative model that integrates hierarchical vector quantization with residual learning to efficiently capture multi-level details in image data. By employing a novel contrastive objective and a tree-structured codebook hierarchy, HR-VQVAE addresses the limitations of classical VQ-VAE architectures, such as codebook collapse and redundancy, enabling high-quality data reconstruction and generation while maintaining computational efficiency.

1. Architecture and Hierarchical Residual Quantization

HR-VQVAE extends the classic VQ-VAE by introducing a hierarchical structure in which multiple quantization layers are arranged sequentially, each responsible for encoding residual information not captured by preceding layers.

• The input image $\mathbf{x}$ is first mapped to a continuous latent embedding $\boldsymbol{\xi}^0 = E(\mathbf{x})$ via an encoder $E$.
• The model employs $n$ quantization layers, indexed by $i$, each equipped with a codebook structure that is hierarchically linked:
  • Layer 1: A single codebook with $m$ codewords.
  • Layer 2: $m$ codebooks (one per parent codeword from Layer 1), each of size $m$ ($m^2$ codewords in total).
  • Layer $i$: $m^{i-1}$ codebooks, each of size $m$, leading to $m^i$ possible codewords.
• At every spatial position $(h, w)$, the codebook selection for layer $i$ is deterministically chosen according to the sequence of indices picked by preceding layers, inducing a tree-like quantization path.

Quantization at each layer operates on the residual: $$\boldsymbol{\xi}^i_{hw} = \boldsymbol{\xi}^{i-1}_{hw} - \mathbf{e}^{i}_{k}$$ where

$$k = \arg \min_j \| \boldsymbol{\xi}^{i-1}_{hw} - \mathbf{e}^i_j \|_2$$

with $\mathbf{e}_j^i$ selected from the codebook identified by the ancestor code indices.

The combined quantized representation used for decoding is the sum across all layers: $\mathbf{e}_C = \sum_{i=1}^n \mathbf{e}^i$ This disciplined residual structure forces each layer to learn representations that are complementary, non-redundant, and hierarchically organized according to the spatial and semantic granularity of the input.

2. Objective Function and Learning

The training objective is designed to promote complementary code usage and to avoid redundancy between layers. It consists of three principal terms and per-layer regularization:

$$\begin{align*} \mathcal{L}\big(\mathbf{x}, \mathcal{D}(\mathbf{e}_C)\big) =& \underbrace{\lVert \mathbf{x} - \mathcal{D}(\mathbf{e}_C) \rVert^2_2}_{\text{reconstruction}} + \underbrace{\lVert \textrm{sg}[\boldsymbol{\xi}^{0}] - \mathbf{e}_C \rVert^2_2 + \beta_0 \lVert \textrm{sg}[\mathbf{e}_C] - \boldsymbol{\xi}^0 \rVert^2_2}_{\text{VQ commitment, codebook alignment}} \ &+ \sum_{i=1}^n \underbrace{\left[ \lVert \textrm{sg}[\boldsymbol{\xi}^{i-1}] - \mathbf{e}^i \rVert^2_2 + \beta_i \lVert \textrm{sg}[\mathbf{e}^i] - \boldsymbol{\xi}^{i-1} \rVert^2_2 \right]}_{\text{per-layer VQ regularization}} \end{align*}$$

where $\mathcal{D}(\cdot)$ is the decoder, $\textrm{sg}[\cdot]$ is the stop-gradient operator, and $\beta_0$, $\beta_i$ are hyperparameters. The crucial aspect is that the total embedding for decoding is always the sum of all quantized outputs, ensuring each layer focuses on increasingly finer residuals.

Unlike in VQ-VAE-2, where each layer incurs a parallel reconstruction loss (possibly leading to encoding redundancy), the HR-VQVAE loss penalizes overlapping information, ensuring strict complementarity.

3. Efficient Hierarchical Decoding and Codebook Utilization

The design yields significant efficiency benefits:

• For $n$ hierarchical layers and codebook size $m$, the space of representable codes is $m^n$, yet only $n \times m$ codewords are searched at inference due to the single active path constraint per sample.
• Decoding is thus reduced from exponential to linear complexity in the number of layers, supporting rapid inference and making HR-VQVAE scalable for high-load or large-scale applications.
• The hierarchical linkage prevents codebook and layer "collapse," a phenomenon in which increasing codebook size in flat VQ-VAEs leads to most codes being underutilized or never assigned, degrading model capacity.

Empirical results confirm that HR-VQVAE remains robust to increasing codebook sizes, continuing to leverage the full codebook capacity without collapse, unlike VQVAE and VQVAE-2.

4. Empirical Evaluation: Reconstruction, Generation, and Speed

HR-VQVAE is quantitatively benchmarked on diverse image datasets (FFHQ, ImageNet, CIFAR10, MNIST):

Model	FFHQ (FID/MSE)	ImageNet	CIFAR10	MNIST
VQVAE	2.86/0.00298	3.66/0.00055	21.65/0.00092	7.9/0.00041
VQVAE-2	1.92/0.00195	2.94/0.00039	18.03/0.00068	6.7/0.00025
HR-VQVAE	1.26/0.00163	2.28/0.00027	18.11/0.00041	6.1/0.00011

• Reconstruction Quality: HR-VQVAE outperforms VQVAE and VQVAE-2 both in MSE and FID for all datasets, with improvements becoming more pronounced with additional hierarchy depth.
• Sample Generation: HR-VQVAE achieves lower FID than prior methods, with generated samples exhibiting greater fidelity and diversity.
• Decoding Speed: For reconstructing 10,000 samples, HR-VQVAE is over an order of magnitude faster than VQVAE-2. For example, on FFHQ, HR-VQVAE decodes in 0.84s versus 9.34s for VQVAE-2, with further speed-up as hierarchy depth increases.

5. Advantages and Theoretical Implications

The HR-VQVAE architecture confers several advantages over previous vector quantization VAEs:

• Scalability: Exponential code capacity with only linear inference cost.
• No Codebook Collapse: Maintains high codebook usage even as capacity increases.
• Progressive Representation: Reconstructions using partial sums across layers are increasingly detailed; each layer enhances structural fidelity.
• Interpretability: Layerwise latent codes correspond to progressively finer structure (coarse-to-fine) in the reconstructed data.

A plausible implication is that HR-VQVAE's local hierarchical quantization could be extended to modalities beyond images, for instance, video or structured data, by leveraging analogous residual trees in broader contexts.

Feature	VQVAE	VQVAE-2	HR-VQVAE
Codebook Usage	Prone to collapse	Prone to collapse	Robust and full utilization
Hierarchical Linkage	Absent	Weak	Strict residual, local search
Maximum Codebook Size	Limited	Limited	Large, no collapse
Reconstruction Fidelity	Moderate	Higher, saturates	Best and improves with depth
Decoding Speed	Medium	Slow	Very fast

6. Practical Considerations and Applications

Due to the linear search complexity, the method is suitable for deployment in environments with high-throughput inference requirements. HR-VQVAE's ability to add layers for increased fidelity without redundancy or computational penalty enables dynamic quality control.

The architecture provides an effective framework for both image reconstruction and generative modeling. It is particularly advantageous in tasks demanding scalable, interpretable, and efficient discrete representations, such as large-scale image generation, image compression, and systems requiring rapid sample reconstruction.

7. Relationship to Related Models and Future Directions

HR-VQVAE provides a sharp contrast to models such as VQVAE-2, which use multi-layer but weakly linked latent codes, often leading to information redundancy and under-utilized capacity; in HR-VQVAE, strict residual learning and path-dependent codebook selection enforce efficient code usage and improved generalization.

Future research may investigate further generalizations of the HR-VQVAE paradigm, for instance, combining the hierarchical quantization tree with stochastic or Bayesian assignment (as in HQ-VAE (Takida et al., 2023)), or integrating attention mechanisms for complex data modalities.

---

HR-VQVAE thus represents a rigorous advance in discrete representation learning for images, combining hierarchical quantization and residual learning, and achieves state-of-the-art performance in both qualitative and quantitative metrics through an architecture that is simple to scale, efficiently trainable, and highly effective in practice (Adiban et al., 2022).