PQ-VAE: Learning Hierarchical Discrete Representations with Progressive Quantization

Abstract

Variational auto-encoders (VAEs) are widely used in generative modeling and representation learning with applications ranging from image generation to data compression. However, conventional VAEs face challenges in balancing the trade-off between compactness and informativeness of the learned latent codes. In this work, we propose Progressive Quantization VAE (PQ-VAE), which aims to learn a progressive sequential structure for data representation that maximizes the mutual information between the latent representations and the original data in a limited description length. The resulting representations provide a global, compact, and hierarchical understanding of the data semantics, making it suitable for high-level tasks while achieving high compression rates. The proposed model offers an effective solution for generative modeling and data compression while enabling improved performance in high-level tasks such as image understanding and generation.

Progressive Quantization

One quantization step. Feature map F_l is mapped to a code vector by T_l, vector-quantized, and mapped back. The residual F_l − F̂_l becomes the input to the next stage, so the codes form a coarse-to-fine hierarchy.

Highlights

Hierarchical codes
Global-to-local sequential structure instead of a spatially flat code grid.
Rate-aware
Maximizes mutual information with the input under a limited description length.
Useful downstream
The same codes support compression, image understanding, and generation.

BibTeX

@inproceedings{huang2024pqvae,
  title     = {PQ-VAE: Learning Hierarchical Discrete Representations with Progressive Quantization},
  author    = {Huang, Lun and Qiu, Qiang and Sapiro, Guillermo},
  booktitle = {Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) Workshops},
  pages     = {7550--7558},
  year      = {2024}
}