OptVQ: Add a New VQ Method via Optimal Transport

[zbr17]

We are excited to propose a pull request (PR) for the integration of OptVQ, a new approach to Vector Quantization (VQ) that elegantly tackles the notorious issue of training instability, especially index collapse, by employing optimal transport theory in place of the conventional nearest neighbor search. This method, as described in the research paper "Preventing Local Pitfalls in Vector Quantization via Optimal Transport," remarkably boosts codebook utilization, reaching nearly 100% efficiency in most scenarios without the need for special initialization or additional regularization techniques.

We allow users to activate this feature by simply setting optimal_transport = True within the VectorQuantize class. Users can also reach the original code at the GitHub repository zbr17/OptVQ. For those who are curious about the visual impact of OptVQ compared to standard VQ in a 2D context, we have included a notebook that demonstrates this comparison. Additionally, we provide an example of training a VQ-VAE model with OptVQ seamlessly integrated into the framework, showcasing its practical application.

We hope that this PR will be well-received and contribute to the community's efforts in advancing vector quantization techniques. Thank you for considering our proposal, and we look forward to any feedback or suggestions you may have.

[zbr17]

Are there any additional changes or clarifications we need to provide? We would appreciate any feedback or updates. Thank you! @lucidrains