Skip to content

Commit 852b98e

Browse files
committed
document the new SimVQ and ResidualSimVQ
1 parent f883a07 commit 852b98e

File tree

2 files changed

+41
-0
lines changed

2 files changed

+41
-0
lines changed

README.md

+41
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,47 @@ indices = quantizer(x)
276276

277277
This repository should also automatically synchronizing the codebooks in a multi-process setting. If somehow it isn't, please open an issue. You can override whether to synchronize codebooks or not by setting `sync_codebook = True | False`
278278

279+
### Sim VQ
280+
281+
<img src="./images/simvq.png" width="400px"></img>
282+
283+
A <a href="https://arxiv.org/abs/2411.02038">new ICLR 2025 paper</a> proposes a scheme where the codebook is frozen, and the codebook is implicitly generated through a linear projection. The author claims this setup leads to less codebook collapse as well as easier convergence. I have found this to perform even better when paired with <a href="https://arxiv.org/abs/2410.06424">rotation trick</a> from Fifty et al., and expanding the linear projection to a small one layer MLP. You can experiment with it as so
284+
285+
```python
286+
import torch
287+
from vector_quantize_pytorch import SimVQ
288+
289+
sim_vq = SimVQ(
290+
dim = 512,
291+
codebook_size = 1024
292+
)
293+
294+
x = torch.randn(1, 1024, 512)
295+
quantized, indices, commit_loss = sim_vq(x)
296+
297+
assert x.shape == quantized.shape
298+
assert torch.allclose(quantized, sim_vq.indices_to_codes(indices), atol = 1e-6)
299+
```
300+
301+
For the residual flavor, just import `ResidualSimVQ` instead
302+
303+
```python
304+
import torch
305+
from vector_quantize_pytorch import SimVQ, ResidualSimVQ
306+
307+
residual_sim_vq = ResidualSimVQ(
308+
dim = 512,
309+
num_quantizers = 4,
310+
codebook_size = 1024
311+
)
312+
313+
x = torch.randn(1, 1024, 512)
314+
quantized, indices, commit_loss = residual_sim_vq(x)
315+
316+
assert x.shape == quantized.shape
317+
assert torch.allclose(quantized, residual_sim_vq.get_output_from_indices(indices), atol = 1e-6)
318+
```
319+
279320
### Finite Scalar Quantization
280321

281322
<img src="./images/fsq.png" width="500px"></img>

images/simvq.png

91 KB
Loading

0 commit comments

Comments
 (0)