vector-quantize-pytorch 1.28.0__tar.gz → 1.28.2__tar.gz

{vector_quantize_pytorch-1.28.0 → vector_quantize_pytorch-1.28.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.28.0
+Version: 1.28.2
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -859,3 +859,15 @@ assert loss.item() >= 0
     url     = {https://arxiv.org/abs/2509.10140},
 }
 ```
+
+```bibtex
+@misc{chee2024quip2bitquantizationlarge,
+    title   = {QuIP: 2-Bit Quantization of Large Language Models With Guarantees}, 
+    author  = {Jerry Chee and Yaohui Cai and Volodymyr Kuleshov and Christopher De Sa},
+    year    = {2024},
+    eprint  = {2307.13304},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2307.13304}, 
+}
+```

{vector_quantize_pytorch-1.28.0 → vector_quantize_pytorch-1.28.2}/README.md

@@ -815,3 +815,15 @@ assert loss.item() >= 0
     url     = {https://arxiv.org/abs/2509.10140},
 }
 ```
+
+```bibtex
+@misc{chee2024quip2bitquantizationlarge,
+    title   = {QuIP: 2-Bit Quantization of Large Language Models With Guarantees}, 
+    author  = {Jerry Chee and Yaohui Cai and Volodymyr Kuleshov and Christopher De Sa},
+    year    = {2024},
+    eprint  = {2307.13304},
+    archivePrefix = {arXiv},
+    primaryClass = {cs.LG},
+    url     = {https://arxiv.org/abs/2307.13304}, 
+}
+```

{vector_quantize_pytorch-1.28.0 → vector_quantize_pytorch-1.28.2}/pyproject.toml

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

{vector_quantize_pytorch-1.28.0 → vector_quantize_pytorch-1.28.2}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -76,7 +76,8 @@ class FSQ(Module):
         force_quantization_f32 = True,
         preserve_symmetry = False,
         noise_dropout = 0.,
-        bound_hard_clamp = False # for residual fsq, if input is pre-softclamped to the right range
+        bound_hard_clamp = False,                   # for residual fsq, if input is pre-softclamped to the right range
+        orthogonal_rotation = False                 # increase codebook utilization. ensure levels are symmetric! https://arxiv.org/abs/2307.13304v2
     ):
         super().__init__()
 
@@ -132,6 +133,17 @@ class FSQ(Module):
 
         self.bound_hard_clamp = bound_hard_clamp
 
+        self.orthogonal_rotation = orthogonal_rotation
+
+        if orthogonal_rotation:
+            is_symmetric = len(set(levels)) == 1
+            if not is_symmetric:
+                print('orthogonal_rotation is not recommended for FSQ with asymmetric levels (i.e. where the number of bins differ across dimensions)')
+
+            orthogonal_rot = torch.empty(codebook_dim, codebook_dim)
+            nn.init.orthogonal_(orthogonal_rot)
+            self.register_buffer('orthogonal_rot', orthogonal_rot)
+
     def bound(self, z, eps = 1e-3, hard_clamp = False):
         """ Bound `z`, an array of shape (..., d). """
         maybe_tanh = tanh if not hard_clamp else partial(clamp, min = -1., max = 1.)
@@ -219,6 +231,9 @@ class FSQ(Module):
 
         codes = self._indices_to_codes(indices)
 
+        if self.orthogonal_rotation:
+            codes = codes @ self.orthogonal_rot.t()
+
         if self.keep_num_codebooks_dim:
             codes = rearrange(codes, '... c d -> ... (c d)')
 
@@ -253,6 +268,9 @@ class FSQ(Module):
 
         z = rearrange(z, 'b n (c d) -> b n c d', c = self.num_codebooks)
 
+        if self.orthogonal_rotation:
+            z = z @ self.orthogonal_rot
+
         # whether to force quantization step to be full precision or not
 
         force_f32 = self.force_quantization_f32
@@ -275,6 +293,9 @@ class FSQ(Module):
 
             codes = self.maybe_apply_noise(codes)
 
+            if self.orthogonal_rotation:
+                codes = codes @ self.orthogonal_rot.t()
+
             codes = rearrange(codes, 'b n c d -> b n (c d)')
 
             codes = codes.to(orig_dtype)

{vector_quantize_pytorch-1.28.0 → vector_quantize_pytorch-1.28.2}/vector_quantize_pytorch/lookup_free_quantization.py

@@ -116,7 +116,8 @@ class LFQ(Module):
         experimental_softplus_entropy_loss = False,
         entropy_loss_offset = 5.,                   # how much to shift the loss before softplus
         spherical = False,                          # from https://arxiv.org/abs/2406.07548
-        force_quantization_f32 = True               # will force the quantization step to be full precision
+        force_quantization_f32 = True,              # will force the quantization step to be full precision
+        orthogonal_rotation = False                 # increase codebook utilization without aux losses, inspired by https://arxiv.org/abs/2307.13304v2
     ):
         super().__init__()
 
@@ -165,6 +166,15 @@ class LFQ(Module):
         self.spherical = spherical
         self.maybe_l2norm = (lambda t: l2norm(t) * self.codebook_scale) if spherical else identity
 
+        # orthogonal rotation
+
+        self.orthogonal_rotation = orthogonal_rotation
+
+        if orthogonal_rotation:
+            orthogonal_rot = torch.empty(codebook_dim, codebook_dim)
+            nn.init.orthogonal_(orthogonal_rot)
+            self.register_buffer('orthogonal_rot', orthogonal_rot)
+
         # entropy aux loss related weights
 
         assert 0 < frac_per_sample_entropy <= 1.
@@ -234,6 +244,9 @@ class LFQ(Module):
 
         codes = self.maybe_l2norm(codes)
 
+        if self.orthogonal_rotation:
+            codes = codes @ self.orthogonal_rot.t()
+
         codes = rearrange(codes, '... c d -> ... (c d)')
 
         # whether to project codes out to original dimensions
@@ -287,6 +300,9 @@ class LFQ(Module):
 
         x = rearrange(x, 'b n (c d) -> b n c d', c = self.num_codebooks)
 
+        if self.orthogonal_rotation:
+            x = x @ self.orthogonal_rot
+
         # maybe l2norm
 
         x = self.maybe_l2norm(x)
@@ -412,6 +428,11 @@ class LFQ(Module):
             if force_f32:
                 x = x.type(orig_dtype)
 
+        # rotate back if needed
+
+        if self.orthogonal_rotation:
+            x = x @ self.orthogonal_rot.t()
+
         # merge back codebook dim
 
         x = rearrange(x, 'b n c d -> b n (c d)')

{vector_quantize_pytorch-1.28.0 → vector_quantize_pytorch-1.28.2}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -322,7 +322,7 @@ def directional_reparam(src, tgt, noise_variance = 5e-3):
     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)
+    unit_noised_dir = l2norm(noised_dir).detach()
 
     return src + unit_noised_dir * error_dir_norm
 
@@ -860,6 +860,7 @@ class VectorQuantize(Module):
         assert at_most_one_of(straight_through, rotation_trick, directional_reparam)
         self.rotation_trick = rotation_trick
 
+        assert not (directional_reparam and threshold_ema_dead_code == 0), 'periodic dead code replacement should be enabled when differential reparam method is turned on'
         self.directional_reparam = directional_reparam
         self.directional_reparam_variance = directional_reparam_variance