Soft convex quantization: revisiting Vector Quantization with convex optimization

Tanmay Gautam, Reid Pryzant, Ziyi Yang, Chenguang Zhu, Somayeh Sojoudi
Proceedings of the 6th Annual Learning for Dynamics & Control Conference, PMLR 242:273-285, 2024.

Abstract

Vector Quantization (VQ) is a well-known technique in deep learning for extracting informative discrete latent representations. VQ-embedded models have shown impressive results in a range of applications including image and speech generation. VQ operates as a parametric K-means algorithm that quantizes inputs using a single codebook vector in the forward pass. While powerful, this technique faces practical challenges including codebook collapse, non-differentiability and lossy compression. To mitigate the aforementioned issues, we propose Soft Convex Quantization (SCQ) as a direct substitute for VQ. SCQ works like a differentiable convex optimization (DCO) layer: in the forward pass, we solve for the optimal convex combination of codebook vectors to quantize the inputs. In the backward pass, we leverage differentiability through the optimality conditions of the forward solution. We then introduce a scalable relaxation of the SCQ optimization and demonstrate its efficacy on the CIFAR-10, GTSRB and LSUN datasets. We train powerful SCQ autoencoder models that significantly outperform matched VQ architectures, observing an order of magnitude better image reconstruction and codebook usage with comparable quantization runtime.

Cite this Paper

BibTeX
@InProceedings{pmlr-v242-gautam24a, title = {Soft convex quantization: revisiting {V}ector {Q}uantization with convex optimization}, author = {Gautam, Tanmay and Pryzant, Reid and Yang, Ziyi and Zhu, Chenguang and Sojoudi, Somayeh}, booktitle = {Proceedings of the 6th Annual Learning for Dynamics & Control Conference}, pages = {273--285}, year = {2024}, editor = {Abate, Alessandro and Cannon, Mark and Margellos, Kostas and Papachristodoulou, Antonis}, volume = {242}, series = {Proceedings of Machine Learning Research}, month = {15--17 Jul}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v242/gautam24a/gautam24a.pdf}, url = {https://proceedings.mlr.press/v242/gautam24a.html}, abstract = {Vector Quantization (VQ) is a well-known technique in deep learning for extracting informative discrete latent representations. VQ-embedded models have shown impressive results in a range of applications including image and speech generation. VQ operates as a parametric K-means algorithm that quantizes inputs using a single codebook vector in the forward pass. While powerful, this technique faces practical challenges including codebook collapse, non-differentiability and lossy compression. To mitigate the aforementioned issues, we propose Soft Convex Quantization (SCQ) as a direct substitute for VQ. SCQ works like a differentiable convex optimization (DCO) layer: in the forward pass, we solve for the optimal convex combination of codebook vectors to quantize the inputs. In the backward pass, we leverage differentiability through the optimality conditions of the forward solution. We then introduce a scalable relaxation of the SCQ optimization and demonstrate its efficacy on the CIFAR-10, GTSRB and LSUN datasets. We train powerful SCQ autoencoder models that significantly outperform matched VQ architectures, observing an order of magnitude better image reconstruction and codebook usage with comparable quantization runtime.} }
Endnote
%0 Conference Paper %T Soft convex quantization: revisiting Vector Quantization with convex optimization %A Tanmay Gautam %A Reid Pryzant %A Ziyi Yang %A Chenguang Zhu %A Somayeh Sojoudi %B Proceedings of the 6th Annual Learning for Dynamics & Control Conference %C Proceedings of Machine Learning Research %D 2024 %E Alessandro Abate %E Mark Cannon %E Kostas Margellos %E Antonis Papachristodoulou %F pmlr-v242-gautam24a %I PMLR %P 273--285 %U https://proceedings.mlr.press/v242/gautam24a.html %V 242 %X Vector Quantization (VQ) is a well-known technique in deep learning for extracting informative discrete latent representations. VQ-embedded models have shown impressive results in a range of applications including image and speech generation. VQ operates as a parametric K-means algorithm that quantizes inputs using a single codebook vector in the forward pass. While powerful, this technique faces practical challenges including codebook collapse, non-differentiability and lossy compression. To mitigate the aforementioned issues, we propose Soft Convex Quantization (SCQ) as a direct substitute for VQ. SCQ works like a differentiable convex optimization (DCO) layer: in the forward pass, we solve for the optimal convex combination of codebook vectors to quantize the inputs. In the backward pass, we leverage differentiability through the optimality conditions of the forward solution. We then introduce a scalable relaxation of the SCQ optimization and demonstrate its efficacy on the CIFAR-10, GTSRB and LSUN datasets. We train powerful SCQ autoencoder models that significantly outperform matched VQ architectures, observing an order of magnitude better image reconstruction and codebook usage with comparable quantization runtime.
APA
Gautam, T., Pryzant, R., Yang, Z., Zhu, C. & Sojoudi, S.. (2024). Soft convex quantization: revisiting Vector Quantization with convex optimization. Proceedings of the 6th Annual Learning for Dynamics & Control Conference, in Proceedings of Machine Learning Research 242:273-285 Available from https://proceedings.mlr.press/v242/gautam24a.html.

Related Material