MQ-VAE: Training Vector-Quantized Networks via Meta Learning

Peijia Qin, Jianguo Zhang

Keywords: Vector-Quantized Networks, Bi-level Optimization, Image Generation

TL;DR: We introduce MQ-VAE, a bi-level optimization-based vector-quantization framework to improve the traditional vector quantization framework.

Abstract:

Deep neural networks with discrete latent variables are particularly well-suited for tasks that naturally involve sequences of discrete symbols.
The vector-quantized variational auto-encoder (VQ-VAE) has made significant progress in this area by leveraging vector quantization.
However, while much effort has been put into maximizing codebook utilization, this does not always result in better performance.
Additional challenges include quantization errors in the VQ layer and the lack of direct integration of task loss into the codebook objective.
To address these issues, we propose Meta-Quantized Variational Auto-Encoder (MQ-VAE), a bi-level optimization-based vector quantization framework inspired by meta-learning.
In MQ-VAE, the codebook and encoder-decoder pair are optimized at different levels, with the codebook treated as hyperparameters optimized via hyper-gradient descent.
This approach effectively tackles these challenges within a unified framework.
The evaluation of MQ-VAE on two computer vision tasks demonstrates its superiority over existing methods and ablation baselines.
Code is available at https://anonymous.4open.science/r/MQVAE-B52C.

Primary Area: unsupervised, self-supervised, semi-supervised, and supervised representation learning

Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.

Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2025/AuthorGuide.

Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors’ identity.

No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.

Submission Number: 9824