vector-quantize-pytorch 1.15.6__tar.gz → 1.16.2__tar.gz

{vector_quantize_pytorch-1.15.6 → vector_quantize_pytorch-1.16.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.15.6
+Version: 1.16.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
@@ -732,3 +732,14 @@ assert loss.item() >= 0
     url     = {https://api.semanticscholar.org/CorpusID:256901024}
 }
 ```
+
+```bibtex
+@article{Huijben2024ResidualQW,
+    title   = {Residual Quantization with Implicit Neural Codebooks},
+    author  = {Iris Huijben and Matthijs Douze and Matthew Muckley and Ruud van Sloun and Jakob Verbeek},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2401.14732},
+    url     = {https://api.semanticscholar.org/CorpusID:267301189}
+}
+```

{vector_quantize_pytorch-1.15.6 → vector_quantize_pytorch-1.16.2}/README.md

@@ -688,3 +688,14 @@ assert loss.item() >= 0
     url     = {https://api.semanticscholar.org/CorpusID:256901024}
 }
 ```
+
+```bibtex
+@article{Huijben2024ResidualQW,
+    title   = {Residual Quantization with Implicit Neural Codebooks},
+    author  = {Iris Huijben and Matthijs Douze and Matthew Muckley and Ruud van Sloun and Jakob Verbeek},
+    journal = {ArXiv},
+    year    = {2024},
+    volume  = {abs/2401.14732},
+    url     = {https://api.semanticscholar.org/CorpusID:267301189}
+}
+```

{vector_quantize_pytorch-1.15.6 → vector_quantize_pytorch-1.16.2}/pyproject.toml

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

{vector_quantize_pytorch-1.15.6 → vector_quantize_pytorch-1.16.2}/tests/test_readme.py

@@ -62,19 +62,31 @@ def test_vq_mask():
     assert (mask_quantized[:, 512:] == 0.).all()
     assert (mask_indices[:, 512:] == -1).all()
 
-def test_residual_vq():
+@pytest.mark.parametrize('implicit_neural_codebook', (True, False))
+@pytest.mark.parametrize('use_cosine_sim', (True, False))
+@pytest.mark.parametrize('train', (True, False))
+def test_residual_vq(
+    implicit_neural_codebook,
+    use_cosine_sim,
+    train
+):
     from vector_quantize_pytorch import ResidualVQ
 
     residual_vq = ResidualVQ(
-        dim = 256,
-        num_quantizers = 8,      # specify number of quantizers
-        codebook_size = 1024,    # codebook size
+        dim = 32,
+        num_quantizers = 8,
+        codebook_size = 128,
+        implicit_neural_codebook = implicit_neural_codebook,
+        use_cosine_sim = use_cosine_sim,
     )
 
-    x = torch.randn(1, 1024, 256)
+    x = torch.randn(1, 256, 32)
 
-    quantized, indices, commit_loss = residual_vq(x)
-    quantized, indices, commit_loss, all_codes = residual_vq(x, return_all_codes = True)
+    residual_vq.train(train)
+
+    quantized, indices, commit_loss = residual_vq(x, freeze_codebook = train and not implicit_neural_codebook)
+    quantized_out = residual_vq.get_output_from_indices(indices)
+    assert torch.allclose(quantized, quantized_out, atol = 1e-6)
 
 def test_residual_vq2():
     from vector_quantize_pytorch import ResidualVQ
@@ -91,7 +103,6 @@ def test_residual_vq2():
     x = torch.randn(1, 1024, 256)
     quantized, indices, commit_loss = residual_vq(x)
 
-
 def test_grouped_residual_vq():
     from vector_quantize_pytorch import GroupedResidualVQ
 

{vector_quantize_pytorch-1.15.6 → vector_quantize_pytorch-1.16.2}/vector_quantize_pytorch/residual_vq.py

@@ -22,6 +22,9 @@ from einx import get_at
 def exists(val):
     return val is not None
 
+def first(it):
+    return it[0]
+
 def default(val, d):
     return val if exists(val) else d
 
@@ -34,6 +37,64 @@ def round_up_multiple(num, mult):
 def is_distributed():
     return dist.is_initialized() and dist.get_world_size() > 1
 
+# the mlp for generating the neural implicit codebook
+# from Huijben et al. https://arxiv.org/abs/2401.14732
+
+class MLP(Module):
+    def __init__(
+        self,
+        dim,
+        dim_hidden = None,
+        depth = 4,             # they used 4 layers in the paper
+        l2norm_output = False
+    ):
+        super().__init__()
+        dim_hidden = default(dim_hidden, dim)
+
+        self.proj_in = nn.Linear(2 * dim, dim)
+
+        layers = ModuleList([])
+
+        for _ in range(depth):
+            layers.append(nn.Sequential(
+                nn.Linear(dim, dim_hidden),
+                nn.SiLU(),
+                nn.Linear(dim_hidden, dim)
+            ))
+
+        self.layers = layers
+        self.l2norm_output = l2norm_output
+
+    def forward(
+        self,
+        codes,
+        *,
+        condition
+    ):
+        one_headed = codes.ndim == 2
+
+        if one_headed:
+            codes = rearrange(codes, 'c d -> 1 c d')
+
+        heads, num_codes, batch, seq_len = codes.shape[0], codes.shape[-2], condition.shape[0], condition.shape[-2]
+
+        codes = repeat(codes, 'h c d -> h b n c d', n = seq_len, b = batch)
+        condition = repeat(condition, 'b n d -> h b n c d', c = num_codes, h = heads)
+
+        x = torch.cat((condition, codes), dim = -1)
+        x = self.proj_in(x)
+
+        for layer in self.layers:
+            x = layer(x) + x
+
+        if self.l2norm_output:
+            x = F.normalize(x, dim = -1)
+
+        if not one_headed:
+            return x
+
+        return rearrange(x, '1 ... -> ...')
+
 # main class
 
 class ResidualVQ(Module):
@@ -50,7 +111,9 @@ class ResidualVQ(Module):
         quantize_dropout_cutoff_index = 0,
         quantize_dropout_multiple_of = 1,
         accept_image_fmap = False,
-        **kwargs
+        implicit_neural_codebook = False, # QINCo from https://arxiv.org/abs/2401.14732
+        mlp_kwargs: dict = dict(),
+        **vq_kwargs
     ):
         super().__init__()
         assert heads == 1, 'residual vq is not compatible with multi-headed codes'
@@ -65,7 +128,16 @@ class ResidualVQ(Module):
         self.num_quantizers = num_quantizers
 
         self.accept_image_fmap = accept_image_fmap
-        self.layers = ModuleList([VectorQuantize(dim = codebook_dim, codebook_dim = codebook_dim, accept_image_fmap = accept_image_fmap, **kwargs) for _ in range(num_quantizers)])
+
+        self.implicit_neural_codebook = implicit_neural_codebook
+
+        if implicit_neural_codebook:
+            vq_kwargs.update(
+                learnable_codebook = True,
+                ema_update = False
+            )
+
+        self.layers = ModuleList([VectorQuantize(dim = codebook_dim, codebook_dim = codebook_dim, accept_image_fmap = accept_image_fmap, **vq_kwargs) for _ in range(num_quantizers)])
 
         assert all([not vq.has_projections for vq in self.layers])
 
@@ -76,6 +148,12 @@ class ResidualVQ(Module):
         self.quantize_dropout_cutoff_index = quantize_dropout_cutoff_index
         self.quantize_dropout_multiple_of = quantize_dropout_multiple_of  # encodec paper proposes structured dropout, believe this was set to 4
 
+        # setting up the MLPs for implicit neural codebooks
+
+        self.mlps = ModuleList([MLP(dim = codebook_dim, l2norm_output = first(self.layers).use_cosine_sim, **mlp_kwargs) for _ in range(num_quantizers - 1)])
+
+        # sharing codebook logic
+
         if not shared_codebook:
             return
 
@@ -120,7 +198,31 @@ class ResidualVQ(Module):
         mask = indices == -1.
         indices = indices.masked_fill(mask, 0) # have it fetch a dummy code to be masked out later
 
-        all_codes = get_at('q [c] d, b n q -> q b n d', self.codebooks, indices)
+        if not self.implicit_neural_codebook:
+            # gather all the codes
+
+            all_codes = get_at('q [c] d, b n q -> q b n d', self.codebooks, indices)
+
+        else:
+            # else if using implicit neural codebook, codes will need to be derived layer by layer
+
+            code_transform_mlps = (None, *self.mlps)
+
+            all_codes = []
+            quantized_out = 0.
+
+            for codes, indices, maybe_transform_mlp in zip(self.codebooks, indices.unbind(dim = -1), code_transform_mlps):
+
+                if exists(maybe_transform_mlp):
+                    codes = maybe_transform_mlp(codes, condition = quantized_out)
+                    layer_codes = get_at('b n [c] d, b n -> b n d', codes, indices)
+                else:
+                    layer_codes = get_at('[c] d, b n -> b n d', codes, indices)
+
+                all_codes.append(layer_codes)
+                quantized_out += layer_codes
+
+            all_codes = torch.stack(all_codes)
 
         # mask out any codes that were dropout-ed
 
@@ -195,9 +297,16 @@ class ResidualVQ(Module):
             null_indices = torch.full(null_indices_shape, -1., device = device, dtype = torch.long)
             null_loss = torch.full((1,), 0., device = device, dtype = x.dtype)
 
+        # setup the mlps for implicit neural codebook
+
+        maybe_code_transforms = (None,) * len(self.layers)
+
+        if self.implicit_neural_codebook:
+            maybe_code_transforms = (None, *self.mlps)
+
         # go through the layers
 
-        for quantizer_index, layer in enumerate(self.layers):
+        for quantizer_index, (vq, maybe_mlp) in enumerate(zip(self.layers, maybe_code_transforms)):
 
             if should_quantize_dropout and quantizer_index > rand_quantize_dropout_index:
                 all_indices.append(null_indices)
@@ -208,12 +317,20 @@ class ResidualVQ(Module):
             if return_loss:
                 layer_indices = indices[..., quantizer_index]
 
-            quantized, *rest = layer(
+            # setup the transform code function to be passed into VectorQuantize forward
+
+            if exists(maybe_mlp):
+                maybe_mlp = partial(maybe_mlp, condition = quantized_out)
+
+            # vector quantize forward
+
+            quantized, *rest = vq(
                 residual,
                 mask = mask,
                 indices = layer_indices,
                 sample_codebook_temp = sample_codebook_temp,
-                freeze_codebook = freeze_codebook
+                freeze_codebook = freeze_codebook,
+                codebook_transform_fn = maybe_mlp
             )
 
             residual = residual - quantized.detach()

{vector_quantize_pytorch-1.15.6 → vector_quantize_pytorch-1.16.2}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from functools import partial, cache
 from collections import namedtuple
 
@@ -59,10 +61,13 @@ def ema_inplace(old, new, decay):
         old.mul_(decay).add_(new * (1 - decay))
 
 def pack_one(t, pattern):
-    return pack([t], pattern)
+    packed, ps = pack([t], pattern)
+
+    def unpack_one(to_unpack, unpack_pattern = None):
+        unpacked, = unpack(to_unpack, ps, default(unpack_pattern, pattern))
+        return unpacked
 
-def unpack_one(t, ps, pattern):
-    return unpack(t, ps, pattern)[0]
+    return packed, unpack_one
 
 def lens_to_mask(lens, max_length):
     seq = torch.arange(max_length, device = lens.device)
@@ -240,12 +245,6 @@ def kmeans(
 
     return means, bins
 
-def batched_embedding(indices, embeds):
-    batch, dim = indices.shape[1], embeds.shape[-1]
-    indices = repeat(indices, 'h b n -> h b n d', d = dim)
-    embeds = repeat(embeds, 'h c d -> h b c d', b = batch)
-    return embeds.gather(2, indices)
-
 # distributed helpers
 
 @cache
@@ -465,7 +464,8 @@ class EuclideanCodebook(Module):
         x,
         sample_codebook_temp = None,
         mask = None,
-        freeze_codebook = False
+        freeze_codebook = False,
+        codebook_transform_fn: Callable | None = None
     ):
         needs_codebook_dim = x.ndim < 4
         sample_codebook_temp = default(sample_codebook_temp, self.sample_codebook_temp)
@@ -476,7 +476,7 @@ class EuclideanCodebook(Module):
             x = rearrange(x, '... -> 1 ...')
 
         dtype = x.dtype
-        flatten, ps = pack_one(x, 'h * d')
+        flatten, unpack_one = pack_one(x, 'h * d')
 
         if exists(mask):
             mask = repeat(mask, 'b n -> c (b h n)', c = flatten.shape[0], h = flatten.shape[-2] // (mask.shape[0] * mask.shape[1]))
@@ -486,24 +486,51 @@ class EuclideanCodebook(Module):
         if self.affine_param:
             self.update_affine(flatten, self.embed, mask = mask)
 
-        embed = self.embed if self.learnable_codebook else self.embed.detach()
+        # affine params
 
         if self.affine_param:
             codebook_std = self.codebook_variance.clamp(min = 1e-5).sqrt()
             batch_std = self.batch_variance.clamp(min = 1e-5).sqrt()
             embed = (embed - self.codebook_mean) * (batch_std / codebook_std) + self.batch_mean
 
-        dist = -cdist(flatten, embed)
+        # get maybe learnable codes
+
+        embed = self.embed if self.learnable_codebook else self.embed.detach()
+
+        # handle maybe implicit neural codebook
+        # and calculate distance
+
+        if exists(codebook_transform_fn):
+            transformed_embed = codebook_transform_fn(embed)
+            transformed_embed = rearrange(transformed_embed, 'h b n c d -> h (b n) c d')
+            broadcastable_input = rearrange(flatten, '... d -> ... 1 d')
+
+            dist = -F.pairwise_distance(broadcastable_input, transformed_embed)
+        else:
+            dist = -cdist(flatten, embed)
+
+        # sample or argmax depending on temperature
 
         embed_ind, embed_onehot = self.gumbel_sample(dist, dim = -1, temperature = sample_codebook_temp, training = self.training)
 
-        embed_ind = unpack_one(embed_ind, ps, 'h *')
+        embed_ind = unpack_one(embed_ind, 'h *')
+
+        if exists(codebook_transform_fn):
+            transformed_embed = unpack_one(transformed_embed, 'h * c d')
 
         if self.training:
-            unpacked_onehot = unpack_one(embed_onehot, ps, 'h * c')
-            quantize = einsum('h b n c, h c d -> h b n d', unpacked_onehot, embed)
+            unpacked_onehot = unpack_one(embed_onehot, 'h * c')
+
+            if exists(codebook_transform_fn):
+                quantize = einsum('h b n c, h b n c d -> h b n d', unpacked_onehot, transformed_embed)
+            else:
+                quantize = einsum('h b n c, h c d -> h b n d', unpacked_onehot, embed)
+
         else:
-            quantize = batched_embedding(embed_ind, embed)
+            if exists(codebook_transform_fn):
+                quantize = einx.get_at('h b n [c] d, h b n -> h b n d', transformed_embed, embed_ind)
+            else:
+                quantize = einx.get_at('h [c] d, h b n -> h b n d', embed, embed_ind)
 
         if self.training and self.ema_update and not freeze_codebook:
 
@@ -533,7 +560,7 @@ class EuclideanCodebook(Module):
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
 
-        dist = unpack_one(dist, ps, 'h * d')
+        dist = unpack_one(dist, 'h * d')
 
         return quantize, embed_ind, dist
 
@@ -650,7 +677,8 @@ class CosineSimCodebook(Module):
         x,
         sample_codebook_temp = None,
         mask = None,
-        freeze_codebook = False
+        freeze_codebook = False,
+        codebook_transform_fn: Callable | None = None
     ):
         needs_codebook_dim = x.ndim < 4
         sample_codebook_temp = default(sample_codebook_temp, self.sample_codebook_temp)
@@ -662,7 +690,7 @@ class CosineSimCodebook(Module):
 
         dtype = x.dtype
 
-        flatten, ps = pack_one(x, 'h * d')
+        flatten, unpack_one = pack_one(x, 'h * d')
 
         if exists(mask):
             mask = repeat(mask, 'b n -> c (b h n)', c = flatten.shape[0], h = flatten.shape[-2] // (mask.shape[0] * mask.shape[1]))
@@ -671,16 +699,37 @@ class CosineSimCodebook(Module):
 
         embed = self.embed if self.learnable_codebook else self.embed.detach()
 
-        dist = einsum('h n d, h c d -> h n c', flatten, embed)
+        # handle maybe implicit neural codebook
+        # and compute cosine sim distance
+
+        if exists(codebook_transform_fn):
+            transformed_embed = codebook_transform_fn(embed)
+            transformed_embed = rearrange(transformed_embed, 'h b n c d -> h (b n) c d')
+            transformed_embed = l2norm(transformed_embed)
+
+            dist = einsum('h n d, h n c d -> h n c', flatten, transformed_embed)
+        else:
+            dist = einsum('h n d, h c d -> h n c', flatten, embed)
 
         embed_ind, embed_onehot = self.gumbel_sample(dist, dim = -1, temperature = sample_codebook_temp, training = self.training)
-        embed_ind = unpack_one(embed_ind, ps, 'h *')
+        embed_ind = unpack_one(embed_ind, 'h *')
+
+        if exists(codebook_transform_fn):
+            transformed_embed = unpack_one(transformed_embed, 'h * c d')
 
         if self.training:
-            unpacked_onehot = unpack_one(embed_onehot, ps, 'h * c')
-            quantize = einsum('h b n c, h c d -> h b n d', unpacked_onehot, embed)
+            unpacked_onehot = unpack_one(embed_onehot, 'h * c')
+
+            if exists(codebook_transform_fn):
+                quantize = einsum('h b n c, h b n c d -> h b n d', unpacked_onehot, transformed_embed)
+            else:
+                quantize = einsum('h b n c, h c d -> h b n d', unpacked_onehot, embed)
+
         else:
-            quantize = batched_embedding(embed_ind, embed)
+            if exists(codebook_transform_fn):
+                quantize = einx.get_at('h b n [c] d, h b n -> h b n d', transformed_embed, embed_ind)
+            else:
+                quantize = einx.get_at('h [c] d, h b n -> h b n d', embed, embed_ind)
 
         if self.training and self.ema_update and not freeze_codebook:
             if exists(mask):
@@ -702,13 +751,13 @@ class CosineSimCodebook(Module):
             embed_normalized = self.embed_avg / rearrange(cluster_size, '... -> ... 1')
             embed_normalized = l2norm(embed_normalized)
 
-            self.embed.data.copy_(l2norm(embed_normalized))
+            self.embed.data.copy_(embed_normalized)
             self.expire_codes_(x)
 
         if needs_codebook_dim:
             quantize, embed_ind = map(lambda t: rearrange(t, '1 ... -> ...'), (quantize, embed_ind))
 
-        dist = unpack_one(dist, ps, 'h * d')
+        dist = unpack_one(dist, 'h * d')
         return quantize, embed_ind, dist
 
 # main class
@@ -845,6 +894,7 @@ class VectorQuantize(Module):
                 affine_param_codebook_decay = affine_param_codebook_decay,
             )
 
+        self.use_cosine_sim = use_cosine_sim
         self._codebook = codebook_class(**codebook_kwargs)
 
         self.in_place_codebook_optimizer = in_place_codebook_optimizer(self._codebook.parameters()) if exists(in_place_codebook_optimizer) else None
@@ -882,7 +932,7 @@ class VectorQuantize(Module):
         if not is_multiheaded:
             codes = codebook[indices]
         else:
-            indices, ps = pack_one(indices, 'b * h')
+            indices, unpack_one = pack_one(indices, 'b * h')
             indices = rearrange(indices, 'b n h -> b h n')
 
             indices = repeat(indices, 'b h n -> b h n d', d = codebook.shape[-1])
@@ -890,7 +940,7 @@ class VectorQuantize(Module):
 
             codes = codebook.gather(2, indices)
             codes = rearrange(codes, 'b h n d -> b n (h d)')
-            codes = unpack_one(codes, ps, 'b * d')
+            codes = unpack_one(codes, 'b * d')
 
         if not self.channel_last:
             codes = rearrange(codes, 'b ... d -> b d ...')
@@ -910,6 +960,7 @@ class VectorQuantize(Module):
         sample_codebook_temp = None,
         freeze_codebook = False,
         return_loss_breakdown = False,
+        codebook_transform_fn: Callable | None = None
     ):
         orig_input = x
 
@@ -962,7 +1013,8 @@ class VectorQuantize(Module):
         codebook_forward_kwargs = dict(
             sample_codebook_temp = sample_codebook_temp,
             mask = mask,
-            freeze_codebook = freeze_codebook
+            freeze_codebook = freeze_codebook,
+            codebook_transform_fn = codebook_transform_fn
         )
 
         # quantize