The Rotation Trick replaces the Straight-Through Estimator (STE) with a geometric rotation of the gradient. It preserves angular information between the encoder output and codebook vector, leading to higher codebook utilization.
I have also benchmarked it on audio. Documented everything from paper to mathematics behind in \docs in github-repo
The implementation uses matrix-free Householder reflection (O(d) memory). PyPI: pip install vqrotate Benchmark: https://github.com/develepo/vqrotate-benchmark
I also explored an adaptive scaling extension. It achieved lower MSE but introduced assignment instability, an interesting trade-off that I would love to hear thoughts on.
Would love any feedback, especially on Residual VQ extensions.