residual LFQ can accept a soft clamp input value and appropriately scale down across residual layers

Author: lucidrains

pyproject.toml

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

vector_quantize_pytorch/lookup_free_quantization.py

@@ -152,7 +152,7 @@ def __init__(
         # whether to soft clamp the input value from -value to value
 
         self.soft_clamp_input_value = soft_clamp_input_value
-        assert not exists(soft_clamp_input_value) or soft_clamp_input_value >= 1.
+        assert not exists(soft_clamp_input_value) or soft_clamp_input_value >= codebook_scale
 
         # for no auxiliary loss, during inference
 

vector_quantize_pytorch/residual_lfq.py

@@ -39,6 +39,7 @@ def __init__(
         quantize_dropout = False,
         quantize_dropout_cutoff_index = 0,
         quantize_dropout_multiple_of = 1,
+        soft_clamp_input_value = None,
         **kwargs
     ):
         super().__init__()
@@ -59,11 +60,15 @@ def __init__(
             lfq = LFQ(
                 dim = codebook_dim,
                 codebook_scale = codebook_scale,
+                soft_clamp_input_value = soft_clamp_input_value,
                 **kwargs
             )
 
             self.layers.append(lfq)
 
+            if exists(soft_clamp_input_value):
+                soft_clamp_input_value *= 0.5
+
         assert all([not lfq.has_projections for lfq in self.layers])
 
         self.quantize_dropout = quantize_dropout and num_quantizers > 1