commit loss weighting for sim vq

Author: lucidrains

pyproject.toml

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

vector_quantize_pytorch/sim_vq.py

@@ -44,6 +44,7 @@ def __init__(
         accept_image_fmap = False,
         rotation_trick = True,  # works even better with rotation trick turned on, with no straight through and the commit loss from input to quantize
         input_to_quantize_commit_loss_weight = 0.25,
+        commitment_weight = 1.,
         frozen_codebook_dim = None # frozen codebook dim could have different dimensions than projection
     ):
         super().__init__()
@@ -74,6 +75,10 @@ def __init__(
 
         self.input_to_quantize_commit_loss_weight = input_to_quantize_commit_loss_weight
 
+        # total commitment loss weight
+
+        self.commitment_weight = commitment_weight
+
     @property
     def codebook(self):
         return self.code_transform(self.frozen_codebook)
@@ -132,7 +137,7 @@ def forward(
 
             indices = inverse_pack(indices, 'b *')
 
-        return quantized, indices, commit_loss
+        return quantized, indices, commit_loss * self.commitment_weight
 
 # main