New method improves digital image coding with full code use

Transporting Quantiles to Codebook: Scalable Vector Quantization without Codebook Collapse

Machine Learning

Summary

Digital tools often compress images using a set of code pieces, but sometimes these codes aren’t used efficiently, limiting quality. The authors propose FlowVQ, which learns to group image data better by smoothly moving simple starting codes into a structured set that covers the data well. This approach avoids the common problem where some codes don’t get used and helps capture images more precisely. They tested it on image generation tasks and showed it scales well while improving how much of the code set is used.

What this means in practice

  • For image generation developers: Improve image generation systems by using a new quantization method that fully exploits codebook capacity for higher quality outputs.
  • For signal processing engineers: Design compression algorithms with robust discrete representations that avoid wasted code usage leading to more efficient encoding.

Authors

Yuwei Zeng, Zekun Shi

Abstract

Effective discrete representation learning with VQ-VAE serves as a fundamental component of modern autoregressive generative models. However, it is widely known to suffer from codebook collapse and inefficient code utilization. To address this issue, we propose FlowVQ, a simple yet effective alternative to standard vector quantization by formulating neural quantization as a joint modeling and clustering process. Specifically, we learn the latent space using a lightweight flow matching model alongside training, corresponding to a probability flow ODE that transports an often fixed well-behaved source distribution while preserving probability mass. Based on this, we define the codebook C_1 by transporting even-quantile partitioned prior codes C_0 from the source space through the learned flow, and obtaining the Voronoi centroids under the induced distribution, leading to a structured discretization. We evaluated on image tokenization and generation, as well as on two tasks requiring high-precision tokenization, demonstrating scalability and improved expressivity of the codebook with full utilization.