add ability to soft clamp the value before lookup free quantization step

Author: lucidrains

pyproject.toml

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

vector_quantize_pytorch/lookup_free_quantization.py

@@ -65,7 +65,8 @@ def __init__(
         codebook_scale = 1.,                        # for residual LFQ, codebook scaled down by 2x at each layer
         frac_per_sample_entropy = 1.,               # make less than 1. to only use a random fraction of the probs for per sample entropy
         use_code_agnostic_commit_loss = False,
-        projection_has_bias = True
+        projection_has_bias = True,
+        soft_clamp_input_value = None
     ):
         super().__init__()
 
@@ -114,6 +115,11 @@ def __init__(
         self.commitment_loss_weight = commitment_loss_weight
         self.use_code_agnostic_commit_loss = use_code_agnostic_commit_loss
 
+        # 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.
+
         # for no auxiliary loss, during inference
 
         self.register_buffer('mask', 2 ** torch.arange(codebook_dim - 1, -1, -1))
@@ -195,6 +201,12 @@ def forward(
 
         x = self.project_in(x)
 
+        # maybe soft clamp
+
+        if exists(self.soft_clamp_input_value):
+            clamp_value = self.soft_clamp_input_value
+            x = (x / clamp_value).tanh() * clamp_value
+
         # split out number of codebooks
 
         x = rearrange(x, 'b n (c d) -> b n c d', c = self.num_codebooks)