vector-quantize-pytorch 1.23.4__tar.gz → 1.24.0__tar.gz

{vector_quantize_pytorch-1.23.4 → vector_quantize_pytorch-1.24.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.23.4
+Version: 1.24.0
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -823,3 +823,15 @@ assert loss.item() >= 0
     url     = {https://api.semanticscholar.org/CorpusID:274423406}
 }
 ```
+
+```bibtex
+@misc{vali2025diveqdifferentiablevectorquantization,
+    title   = {DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick}, 
+    author  = {Mohammad Hassan Vali and Tom Bäckström and Arno Solin},
+    year    = {2025},
+    eprint  = {2509.26469},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2509.26469}, 
+}
+```

{vector_quantize_pytorch-1.23.4 → vector_quantize_pytorch-1.24.0}/README.md

@@ -779,3 +779,15 @@ assert loss.item() >= 0
     url     = {https://api.semanticscholar.org/CorpusID:274423406}
 }
 ```
+
+```bibtex
+@misc{vali2025diveqdifferentiablevectorquantization,
+    title   = {DiVeQ: Differentiable Vector Quantization Using the Reparameterization Trick}, 
+    author  = {Mohammad Hassan Vali and Tom Bäckström and Arno Solin},
+    year    = {2025},
+    eprint  = {2509.26469},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2509.26469}, 
+}
+```

{vector_quantize_pytorch-1.23.4 → vector_quantize_pytorch-1.24.0}/examples/autoencoder.py

@@ -75,8 +75,9 @@ torch.random.manual_seed(seed)
 
 model = SimpleVQAutoEncoder(
     codebook_size = num_codes,
-    rotation_trick = True,
-    straight_through = False
+    rotation_trick = False,
+    straight_through = False,
+    directional_reparam = True
 ).to(device)
 
 opt = torch.optim.AdamW(model.parameters(), lr=lr)

{vector_quantize_pytorch-1.23.4 → vector_quantize_pytorch-1.24.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.23.4"
+version = "1.24.0"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{vector_quantize_pytorch-1.23.4 → vector_quantize_pytorch-1.24.0}/tests/test_readme.py

@@ -224,6 +224,28 @@ def test_tiger():
 
     assert torch.allclose(quantized, quantized_out, atol = 1e-5)
 
+def test_directional_reparam():
+    from vector_quantize_pytorch import VectorQuantize, ResidualVQ
+
+    vq = VectorQuantize(
+        dim = 256,
+        codebook_size = 512,                # codebook size
+        directional_reparam = True
+    )
+
+    x = torch.randn(1, 1024, 256).requires_grad_()
+
+    quantized, indices, _ = vq(x)
+
+    rq = ResidualVQ(
+        dim = 256,
+        num_quantizers = 8,
+        codebook_size = 128,
+        directional_reparam = True
+    )
+
+    quantized, indices, _ = rq(x)
+
 @pytest.mark.parametrize('preserve_symmetry', (True, False))
 def test_fsq(
     preserve_symmetry

{vector_quantize_pytorch-1.23.4 → vector_quantize_pytorch-1.24.0}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from math import sqrt
 from functools import partial, cache
 from collections import namedtuple
 
@@ -28,6 +29,9 @@ def noop(*args, **kwargs):
 def identity(t):
     return t
 
+def at_most_one_of(*bools):
+    return sum([*map(int, bools)]) <= 1
+
 def l2norm(t, dim = -1,  eps = 1e-6):
     return F.normalize(t, p = 2, dim = dim, eps = eps)
 
@@ -264,6 +268,11 @@ def kmeans(
 
     return means, bins
 
+# straight through
+
+def straight_through(src, tgt):
+    return src + (tgt + src).detach()
+
 # rotation trick related
 
 def efficient_rotation_trick_transform(u, q, e):
@@ -279,7 +288,7 @@ def efficient_rotation_trick_transform(u, q, e):
         2 * (e @ rearrange(u, 'b d -> b d 1').detach() @ rearrange(q, 'b d -> b 1 d').detach())
     )
 
-    return rearrange(out, '... 1 -> ...')
+    return rearrange(out, '... 1 d -> ... d')
 
 def rotate_to(src, tgt):
     # rotation trick STE (https://arxiv.org/abs/2410.06424) to get gradients through VQ layer.
@@ -299,6 +308,18 @@ def rotate_to(src, tgt):
 
     return inverse(rotated)
 
+# directional reparam related
+# figure 1. https://openreview.net/forum?id=KRVnpTbx7R
+
+def directional_reparam(src, tgt, noise_variance = 5e-3):
+    error_dir = tgt - src
+    error_dir_norm = error_dir.norm(dim = -1, keepdim = True)
+
+    noised_dir = error_dir + sqrt(noise_variance) * torch.randn_like(error_dir)
+    unit_noised_dir = l2norm(noised_dir)
+
+    return src + unit_noised_dir * error_dir_norm
+
 # distributed helpers
 
 @cache
@@ -898,12 +919,14 @@ class VectorQuantize(Module):
         stochastic_sample_codes = False,
         sample_codebook_temp = 1.,
         straight_through = False,
-        rotation_trick = None,  # Propagate grads through VQ layer w/ rotation trick: https://arxiv.org/abs/2410.06424 by @cfifty
+        rotation_trick = None,       # propagate grads through VQ layer w/ rotation trick: https://arxiv.org/abs/2410.06424 by @cfifty
+        directional_reparam = False, # add the difference between the nearest code and input vector, with some noise on the direction
+        directional_reparam_variance = 5e-3,
         sync_codebook = None,
         sync_affine_param = False,
-        ema_update = True,
+        ema_update = None,
         manual_ema_update = False,
-        learnable_codebook = False,
+        learnable_codebook = None,
         in_place_codebook_optimizer: Callable[..., Optimizer] = None, # Optimizer used to update the codebook embedding if using learnable_codebook
         manual_in_place_optimizer_update = False,
         affine_param = False,
@@ -913,7 +936,14 @@ class VectorQuantize(Module):
         return_zeros_for_masked_padding = True
     ):
         super().__init__()
-        rotation_trick = default(rotation_trick, dim > 1) # only use rotation trick if feature dimension greater than 1
+
+        # defaults
+
+        ema_update = default(ema_update, not directional_reparam)
+        learnable_codebook = default(learnable_codebook, directional_reparam)
+        rotation_trick = default(rotation_trick, not directional_reparam and dim > 1) # only use rotation trick if feature dimension greater than 1
+
+        # basic variables
 
         self.dim = dim
         self.heads = heads
@@ -935,7 +965,7 @@ class VectorQuantize(Module):
 
         self.eps = eps
 
-        self.has_commitment_loss = commitment_weight > 0.
+        self.has_commitment_loss = commitment_weight > 0. and not directional_reparam
         self.commitment_weight = commitment_weight
         self.commitment_use_cross_entropy_loss = commitment_use_cross_entropy_loss # whether to use cross entropy loss to codebook as commitment loss
 
@@ -953,9 +983,12 @@ class VectorQuantize(Module):
         self.codebook_diversity_temperature = codebook_diversity_temperature
         self.codebook_diversity_loss_weight = codebook_diversity_loss_weight
 
-        assert not (straight_through and rotation_trick)
+        assert at_most_one_of(straight_through, rotation_trick, directional_reparam)
         self.rotation_trick = rotation_trick
 
+        self.directional_reparam = directional_reparam
+        self.directional_reparam_variance = directional_reparam_variance
+
         assert not (ema_update and learnable_codebook), 'learnable codebook not compatible with EMA update'
 
         assert 0 <= sync_update_v <= 1.
@@ -1213,9 +1246,11 @@ class VectorQuantize(Module):
             if input_requires_grad:
                 if self.rotation_trick:
                     quantize = rotate_to(x, quantize)
+                elif self.directional_reparam:
+                    quantize = directional_reparam(x, quantize, self.directional_reparam_variance)
                 else:
                     # standard STE to get gradients through VQ layer.
-                    quantize = x + (quantize - x).detach()
+                    quantize = straight_through(x, quantize)
 
             if self.sync_update_v > 0.:
                 # (21) in https://minyoungg.github.io/vqtorch/assets/draft_050523.pdf