turn on a new finding that dramatically improves vector quantization

README.md

@@ -570,3 +570,14 @@ $ accelerate launch train.py
     url     = {https://api.semanticscholar.org/CorpusID:265018962}
 }
 ```
+
+```bibtex
+@article{Fifty2024Restructuring,
+    title   = {Restructuring Vector Quantization with the Rotation Trick},
+    author  = {Christopher Fifty, Ronald G. Junkins, Dennis Duan, Aniketh Iyengar, Jerry W. Liu, Ehsan Amid, Sebastian Thrun, Christopher Ré},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2410.06424},
+    url     = {https://api.semanticscholar.org/CorpusID:273229218}
+}
+```

audiolm_pytorch/soundstream.py

@@ -463,6 +463,7 @@ def __init__(
         rq_quantize_dropout_multiple_of = 1,
         rq_groups = 1,
         rq_stochastic_sample_codes = False,
+        rq_rotation_trick = True,
         rq_kwargs: dict = {},
         use_lookup_free_quantizer = False,              # proposed in https://arxiv.org/abs/2310.05737, adapted for residual quantization
         use_finite_scalar_quantizer = False,            # proposed in https://arxiv.org/abs/2309.15505, adapted for residual quantization
@@ -600,6 +601,7 @@ def __init__(
                 quantize_dropout = True,
                 quantize_dropout_cutoff_index = quantize_dropout_cutoff_index,
                 stochastic_sample_codes = rq_stochastic_sample_codes,
+                rotation_trick = rq_rotation_trick,
                 **rq_kwargs
             )
 

audiolm_pytorch/version.py

@@ -1 +1 @@
-__version__ = '2.0.7'
+__version__ = '2.1.0'

setup.py

@@ -36,7 +36,7 @@
     'torchaudio',
     'transformers',
     'tqdm',
-    'vector-quantize-pytorch>=1.12.5'
+    'vector-quantize-pytorch>=1.18.1'
   ],
   classifiers=[
     'Development Status :: 4 - Beta',