improv

Author: lucidrains

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.27.5"
+version = "1.27.6"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_readme.py

@@ -74,11 +74,13 @@ def test_vq_mask():
 @pytest.mark.parametrize('implicit_neural_codebook, use_cosine_sim', ((True, False), (False, True), (False, False)))
 @pytest.mark.parametrize('train', (True, False))
 @pytest.mark.parametrize('shared_codebook', (True, False))
+@pytest.mark.parametrize('quant_grad_frac', (0., 0.1))
 def test_residual_vq(
     implicit_neural_codebook,
     use_cosine_sim,
     train,
-    shared_codebook
+    shared_codebook,
+    quant_grad_frac
 ):
     from vector_quantize_pytorch import ResidualVQ
 
@@ -88,7 +90,8 @@ def test_residual_vq(
         codebook_size = 128,
         implicit_neural_codebook = implicit_neural_codebook,
         use_cosine_sim = use_cosine_sim,
-        shared_codebook = shared_codebook
+        shared_codebook = shared_codebook,
+        quant_grad_frac = quant_grad_frac
     )
 
     x = torch.randn(1, 256, 32)

vector_quantize_pytorch/residual_vq.py

@@ -38,6 +38,14 @@ def unique(arr):
 def round_up_multiple(num, mult):
     return ceil(num / mult) * mult
 
+def frac_gradient(t, frac):
+    detached = t.detach()
+
+    if frac <= 0:
+        return detached
+
+    return frac * t + (1. - frac) * detached
+
 # tensor helpers
 
 def pad_at_dim(
@@ -174,6 +182,7 @@ def __init__(
         beam_size = None,
         eval_beam_size = None,
         beam_score_quantizer_weights: list[float] | None = None,
+        quant_grad_frac = 0.,
         **vq_kwargs
     ):
         super().__init__()
@@ -233,6 +242,10 @@ def __init__(
 
         self.vq_is_ema_updating = first(self.layers).ema_update
 
+        # gradient related - how much gradients to allow up the residual path (previous layers influence the layer above)
+
+        self.quant_grad_frac = quant_grad_frac
+
         # beam size
 
         assert not (exists(eval_beam_size) and not exists(beam_size))
@@ -493,7 +506,7 @@ def forward(
 
             # core residual vq logic
 
-            residual = residual - quantized.detach()
+            residual = residual - frac_gradient(quantized, self.quant_grad_frac)
             quantized_out = quantized_out + quantized
 
             # handle sort and topk beams