vector-quantize-pytorch 1.25.2__tar.gz → 1.27.0__tar.gz

{vector_quantize_pytorch-1.25.2 → vector_quantize_pytorch-1.27.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vector-quantize-pytorch
-Version: 1.25.2
+Version: 1.27.0
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch
@@ -36,7 +36,7 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Requires-Python: >=3.9
 Requires-Dist: einops>=0.8.0
 Requires-Dist: einx>=0.3.0
-Requires-Dist: torch>=2.0
+Requires-Dist: torch>=2.4
 Provides-Extra: examples
 Requires-Dist: torchvision; extra == 'examples'
 Requires-Dist: tqdm; extra == 'examples'

{vector_quantize_pytorch-1.25.2 → vector_quantize_pytorch-1.27.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.25.2"
+version = "1.27.0"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }
@@ -23,7 +23,7 @@ classifiers=[
   ]
 
 dependencies = [
-    "torch>=2.0",
+    "torch>=2.4",
     "einops>=0.8.0",
     "einx>=0.3.0",
 ]

vector_quantize_pytorch-1.27.0/tests/test_beam.py

@@ -0,0 +1,64 @@
+import torch
+from vector_quantize_pytorch import VectorQuantize
+
+def test_topk_and_manual_ema_update():
+
+    vq1 = VectorQuantize(
+        dim = 256,
+        codebook_size = 512
+    )
+
+    vq2 = VectorQuantize(
+        dim = 256,
+        codebook_size = 512
+    )
+    
+    vq2.load_state_dict(vq1.state_dict())
+
+    x = torch.randn(1, 1024, 256)
+    mask = torch.randint(0, 2, (1, 1024)).bool()
+
+    vq1.train()
+    quantize1, indices1, commit_loss1 = vq1(x, mask = mask)
+
+    vq2.train()
+    quantize2, indices2, commit_losses = vq2(x, mask = mask, topk = 1, ema_update = False)
+
+    assert quantize2.shape == (1, 1024, 1, 256)
+    assert indices2.shape == (1, 1024, 1)
+    assert commit_losses.shape == (1, 1024, 1)
+
+    top_quantize2 = quantize2[..., 0, :]
+    top_indices2 = indices2[..., 0]
+
+    assert torch.allclose(commit_loss1, commit_losses.sum() / mask.sum())
+    assert torch.equal(indices1, top_indices2)
+    assert torch.allclose(quantize1, top_quantize2)
+
+    assert not torch.allclose(vq1._codebook.embed_avg, vq2._codebook.embed_avg)
+
+    vq2.update_ema_indices(x, top_indices2, mask = mask)
+
+    assert torch.allclose(vq1._codebook.cluster_size, vq2._codebook.cluster_size)
+    assert torch.allclose(vq1._codebook.embed_avg, vq2._codebook.embed_avg)
+    assert torch.allclose(vq1.codebook, vq2.codebook)
+
+def test_beam_search():
+    import torch
+    from vector_quantize_pytorch import ResidualVQ
+
+    residual_vq = ResidualVQ(
+        dim = 256,
+        num_quantizers = 8,      # specify number of quantizers
+        codebook_size = 1024,    # codebook size
+        quantize_dropout = True
+    )
+
+    x = torch.randn(1, 1024, 256)
+
+    for _ in range(5):
+        quantized, indices, commit_loss = residual_vq(x, beam_size = 3)
+
+    assert quantized.shape == (1, 1024, 256)
+    assert indices.shape == (1, 1024, 8)
+    assert commit_loss.shape == (8,)

{vector_quantize_pytorch-1.25.2 → vector_quantize_pytorch-1.27.0}/tests/test_readme.py

@@ -247,13 +247,15 @@ def test_directional_reparam():
     quantized, indices, _ = rq(x)
 
 @pytest.mark.parametrize('preserve_symmetry', (True, False))
+@pytest.mark.parametrize('bound_hard_clamp', (True, False))
 def test_fsq(
-    preserve_symmetry
+    preserve_symmetry,
+    bound_hard_clamp
 ):
     from vector_quantize_pytorch import FSQ
 
     levels = [8,5,5,5] # see 4.1 and A.4.1 in the paper
-    quantizer = FSQ(levels, preserve_symmetry = preserve_symmetry)
+    quantizer = FSQ(levels, preserve_symmetry = preserve_symmetry, bound_hard_clamp = bound_hard_clamp)
 
     x = torch.randn(1, 1024, 4) # 4 since there are 4 levels
     xhat, indices = quantizer(x)

{vector_quantize_pytorch-1.25.2 → vector_quantize_pytorch-1.27.0}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -11,7 +11,7 @@ from typing import List, Tuple
 import torch
 import torch.nn as nn
 from torch.nn import Module
-from torch import tensor, Tensor, int32
+from torch import tensor, Tensor, int32, tanh, atanh, clamp
 from torch.amp import autocast
 
 import einx
@@ -30,6 +30,9 @@ def default(*args):
             return arg
     return None
 
+def identity(t):
+    return t
+
 def maybe(fn):
     @wraps(fn)
     def inner(x, *args, **kwargs):
@@ -73,6 +76,7 @@ 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
     ):
         super().__init__()
 
@@ -121,22 +125,31 @@ class FSQ(Module):
         self.allowed_dtypes = allowed_dtypes
         self.force_quantization_f32 = force_quantization_f32
 
-    def bound(self, z, eps = 1e-3):
+        # allow for a hard clamp
+
+        self.bound_hard_clamp = bound_hard_clamp
+
+    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.)
+        maybe_atanh = atanh if not hard_clamp else identity
+
         half_l = (self._levels - 1) * (1 + eps) / 2
         offset = torch.where(self._levels % 2 == 0, 0.5, 0.0)
-        shift = (offset / half_l).atanh()
-        bounded_z = (z + shift).tanh() * half_l - offset
+        shift = maybe_atanh(offset / half_l)
+        bounded_z = maybe_tanh(z + shift) * half_l - offset
         half_width = self._levels // 2
         return round_ste(bounded_z) / half_width
 
     # symmetry-preserving and noise-approximated quantization, section 3.2 in https://arxiv.org/abs/2411.19842
     
-    def symmetry_preserving_bound(self, z):
+    def symmetry_preserving_bound(self, z, hard_clamp = False):
         """ QL(x) = 2 / (L - 1) * [(L - 1) * (tanh(x) + 1) / 2 + 0.5] - 1 """
+        maybe_tanh = tanh if not hard_clamp else partial(clamp, min = -1., max = 1.)
+
         levels_minus_1 = (self._levels - 1)
         scale = 2. / levels_minus_1
-        bracket = (levels_minus_1 * (z.tanh() + 1) / 2.) + 0.5
+        bracket = (levels_minus_1 * (maybe_tanh(z) + 1) / 2.) + 0.5
         bracket = floor_ste(bracket)
         return scale * bracket - 1.
 
@@ -146,7 +159,7 @@ class FSQ(Module):
         shape, device, noise_dropout, preserve_symmetry = z.shape[0], z.device, self.noise_dropout, self.preserve_symmetry
         bound_fn = self.symmetry_preserving_bound if preserve_symmetry else self.bound
 
-        bounded_z = bound_fn(z)
+        bounded_z = bound_fn(z, hard_clamp = self.bound_hard_clamp)
 
         # determine where to add a random offset elementwise
         # if using noise dropout

{vector_quantize_pytorch-1.25.2 → vector_quantize_pytorch-1.27.0}/vector_quantize_pytorch/residual_fsq.py

@@ -1,11 +1,11 @@
+from __future__ import annotations
+
 import random
 from math import ceil
 from functools import partial
 
-from typing import List
-
 import torch
-from torch import nn
+from torch import nn, tensor
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
 from torch.amp import autocast
@@ -52,14 +52,15 @@ class ResidualFSQ(Module):
     def __init__(
         self,
         *,
-        levels: List[int],
+        levels: list[int],
         num_quantizers,
         dim = None,
         is_channel_first = False,
         quantize_dropout = False,
         quantize_dropout_cutoff_index = 0,
         quantize_dropout_multiple_of = 1,
-        soft_clamp_input_value = None,
+        soft_clamp_input_value: float | list[float] | Tensor | None = None,
+        bound_hard_clamp = True,
         **kwargs
     ):
         super().__init__()
@@ -74,25 +75,24 @@ class ResidualFSQ(Module):
         self.is_channel_first = is_channel_first
         self.num_quantizers = num_quantizers
 
-        # soft clamping the input value
-
-        self.soft_clamp_input_value = soft_clamp_input_value
-
         # layers
 
         self.levels = levels
         self.layers = nn.ModuleList([])
 
-        levels_tensor = torch.Tensor(levels)
+        levels_tensor = tensor(levels)
+        assert (levels_tensor > 1).all()
 
         scales = []
 
         for ind in range(num_quantizers):
-            scales.append((levels_tensor - 1) ** -ind)
+            scales.append(levels_tensor.float() ** -ind)
 
             fsq = FSQ(
                 levels = levels,
                 dim = codebook_dim,
+                preserve_symmetry = True,
+                bound_hard_clamp = bound_hard_clamp,
                 **kwargs
             )
 
@@ -111,6 +111,17 @@ class ResidualFSQ(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
 
+        # soft clamping the input value
+
+        if bound_hard_clamp:
+            assert not exists(soft_clamp_input_value)
+            soft_clamp_input_value = 1 + (1 / (levels_tensor - 1))
+
+        if isinstance(soft_clamp_input_value, (list, float)):
+            soft_clamp_input_value = tensor(soft_clamp_input_value)
+
+        self.register_buffer('soft_clamp_input_value', soft_clamp_input_value, persistent = False)
+
     @property
     def codebooks(self):
         codebooks = [layer.implicit_codebook for layer in self.layers]

{vector_quantize_pytorch-1.25.2 → vector_quantize_pytorch-1.27.0}/vector_quantize_pytorch/residual_vq.py

@@ -6,7 +6,7 @@ from functools import partial, cache
 from itertools import zip_longest
 
 import torch
-from torch import nn, Tensor
+from torch import nn, Tensor, arange, cat
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
 import torch.distributed as dist
@@ -14,6 +14,7 @@ from vector_quantize_pytorch.vector_quantize_pytorch import VectorQuantize
 
 from einops import rearrange, repeat, reduce, pack, unpack
 
+import einx
 from einx import get_at
 
 # helper functions
@@ -36,6 +37,47 @@ def unique(arr):
 def round_up_multiple(num, mult):
     return ceil(num / mult) * mult
 
+# tensor helpers
+
+def pad_at_dim(
+    t,
+    pad: tuple[int, int],
+    dim = -1,
+    value = 0.
+):
+    if pad == (0, 0):
+        return t
+
+    dims_from_right = (- dim - 1) if dim < 0 else (t.ndim - dim - 1)
+    zeros = ((0, 0) * dims_from_right)
+    return F.pad(t, (*zeros, *pad), value = value)
+
+def pack_one(t, pattern):
+    packed, packed_shape = pack([t], pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        return first(unpack(out, packed_shape, inv_pattern))
+
+    return packed, inverse
+
+def batch_select(t, indices, pattern = None):
+
+    if exists(pattern):
+        indices = rearrange(indices, '... k -> (...) k')
+        t, inv_pack = pack_one(t, pattern)
+
+
+    batch_indices = arange(t.shape[0], device = t.device)
+    batch_indices = rearrange(batch_indices, 'b -> b 1')
+
+    out = t[batch_indices, indices]
+
+    if exists(pattern):
+        out = inv_pack(out)
+
+    return out
+
 # distributed helpers
 
 def is_distributed():
@@ -128,6 +170,8 @@ class ResidualVQ(Module):
         accept_image_fmap = False,
         implicit_neural_codebook = False, # QINCo from https://arxiv.org/abs/2401.14732
         mlp_kwargs: dict = dict(),
+        beam_size = None,
+        eval_beam_size = None,
         **vq_kwargs
     ):
         super().__init__()
@@ -183,6 +227,15 @@ 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
 
+        # determine whether is using ema update
+
+        self.vq_is_ema_updating = first(self.layers).ema_update
+
+        # beam size
+
+        self.beam_size = default(beam_size, eval_beam_size)
+        self.eval_beam_size = eval_beam_size
+
         # setting up the MLPs for implicit neural codebooks
         
         self.mlps = None
@@ -295,19 +348,29 @@ class ResidualVQ(Module):
         return_all_codes = False,
         sample_codebook_temp = None,
         freeze_codebook = False,
+        beam_size = None,
         rand_quantize_dropout_fixed_seed = None
     ):
-        num_quant, quant_dropout_multiple_of, return_loss, device = self.num_quantizers, self.quantize_dropout_multiple_of, exists(indices), x.device
+
+        # variables
+
+        input_shape, num_quant, quant_dropout_multiple_of, return_loss, device = x.shape, self.num_quantizers, self.quantize_dropout_multiple_of, exists(indices), x.device
+
+        beam_size = default(beam_size, self.beam_size if self.training else self.eval_beam_size)
+
+        is_beam_search = exists(beam_size) and beam_size > 1
+
+        # projecting in
 
         x = self.project_in(x)
 
         assert not (self.accept_image_fmap and exists(indices))
 
-        quantized_out = 0.
+        quantized_out = torch.zeros_like(x)
         residual = x
 
-        all_losses = []
-        all_indices = []
+        all_losses = torch.empty((0,), device = device, dtype = torch.float32)
+        all_indices = torch.empty((*input_shape[:-1], 0), device = device, dtype = torch.long)
 
         if isinstance(indices, list):
             indices = torch.stack(indices)
@@ -319,7 +382,6 @@ class ResidualVQ(Module):
         should_quantize_dropout = self.training and self.quantize_dropout and not return_loss
 
         # sample a layer index at which to dropout further residual quantization
-        # also prepare null indices and loss
 
         if should_quantize_dropout:
 
@@ -335,9 +397,24 @@ class ResidualVQ(Module):
             if quant_dropout_multiple_of != 1:
                 rand_quantize_dropout_index = round_up_multiple(rand_quantize_dropout_index + 1, quant_dropout_multiple_of) - 1
 
-            null_indices_shape = (x.shape[0], *x.shape[-2:]) if self.accept_image_fmap else tuple(x.shape[:2])
-            null_indices = torch.full(null_indices_shape, -1., device = device, dtype = torch.long)
-            null_loss = torch.full((1,), 0., device = device, dtype = x.dtype)
+        # save all inputs across layers, for use during expiration at end under shared codebook setting, or ema update during beam search
+
+        all_residuals = torch.empty((*input_shape[:-1], 0, input_shape[-1]), dtype = residual.dtype, device = device)
+
+        # maybe prepare beam search
+
+        if is_beam_search:
+            prec_dims = x.shape[:-1]
+
+            search_scores = torch.zeros((*prec_dims, 1), device = device, dtype = x.dtype)
+
+            residual = rearrange(residual, '... d -> ... 1 d')
+            quantized_out = rearrange(quantized_out, '... d -> ... 1 d')
+
+            all_residuals = rearrange(all_residuals, '... l d -> ... 1 l d')
+            all_indices = rearrange(all_indices, '... l -> ... 1 l')
+
+            all_losses = all_losses.reshape(*input_shape[:-1], 1, 0)
 
         # setup the mlps for implicit neural codebook
 
@@ -346,17 +423,13 @@ class ResidualVQ(Module):
         if self.implicit_neural_codebook:
             maybe_code_transforms = (None, *self.mlps)
 
-        # save all inputs across layers, for use during expiration at end under shared codebook setting
-
-        all_residuals = []
-
         # go through the 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)
-                all_losses.append(null_loss)
+                all_indices = pad_at_dim(all_indices, (0, 1), value = -1, dim = -1)
+                all_losses = pad_at_dim(all_losses, (0, 1), value = 0, dim = -1)
                 continue
 
             layer_indices = None
@@ -368,9 +441,9 @@ class ResidualVQ(Module):
             if exists(maybe_mlp):
                 maybe_mlp = partial(maybe_mlp, condition = quantized_out)
 
-            # save for expiration
+            # save the residual input for maybe expiration as well as ema update after beam search
 
-            all_residuals.append(residual)
+            all_residuals = cat((all_residuals, rearrange(residual, '... d -> ... 1 d')), dim = -2)
 
             # vector quantize forward
 
@@ -380,29 +453,111 @@ class ResidualVQ(Module):
                 indices = layer_indices,
                 sample_codebook_temp = sample_codebook_temp,
                 freeze_codebook = freeze_codebook,
-                codebook_transform_fn = maybe_mlp
+                codebook_transform_fn = maybe_mlp,
+                topk = beam_size if is_beam_search else None
             )
 
-            residual = residual - quantized.detach()
-            quantized_out = quantized_out + quantized
+            # cross entropy loss for some old paper
 
             if return_loss:
-                ce_loss = rest[0]
+                ce_loss = first(rest)
                 ce_losses.append(ce_loss)
                 continue
 
             embed_indices, loss = rest
 
-            all_indices.append(embed_indices)
-            all_losses.append(loss)
+            # handle expanding first residual if doing beam search
+
+            if is_beam_search:
+
+                search_scores = einx.add('... j, ... j k -> ... (j k)', search_scores, -loss)
+
+                residual = rearrange(residual, '... j d -> ... j 1 d')
+                quantized_out = rearrange(quantized_out, '... j d -> ... j 1 d')
+
+                all_residuals = repeat(all_residuals, '... j l d -> ... (j k) l d', k = beam_size)
+
+            # core residual vq logic
+
+            residual = residual - quantized.detach()
+            quantized_out = quantized_out + quantized
+
+            # handle sort and topk beams
+
+            if is_beam_search:
+                residual = rearrange(residual, '... j k d -> ... (j k) d')
+                quantized_out = rearrange(quantized_out, '... j k d -> ... (j k) d')
+
+                # broadcat the indices
+
+                all_indices = repeat(all_indices, '... j l -> ... j k l', k = embed_indices.shape[-1])
+                embed_indices = rearrange(embed_indices, '... j k -> ... j k 1')
+
+                all_indices = cat((all_indices, embed_indices), dim = -1)
+                all_indices = rearrange(all_indices, '... j k l -> ... (j k) l')
+
+                # broadcat the losses
+
+                all_losses = repeat(all_losses, '... j l -> ... j k l', k = loss.shape[-1])
+                loss = rearrange(loss, '... -> ... 1')
+
+                all_losses = cat((all_losses, loss), dim = -1)
+                all_losses = rearrange(all_losses, '... j k l -> ... (j k) l')
+
+                # handle sort and selection of highest beam size
+
+                if search_scores.shape[-1] > beam_size:
+                    search_scores, select_indices = search_scores.topk(beam_size, dim = -1)
+
+                    residual = batch_select(residual, select_indices, '* k d')
+                    quantized_out = batch_select(quantized_out, select_indices, '* k d')
+
+                    all_indices = batch_select(all_indices, select_indices, '* k l')
+                    all_losses = batch_select(all_losses, select_indices, '* k l')
+
+                    all_residuals = batch_select(all_residuals, select_indices, '* k l d')
+            else:
+                # aggregate indices and losses
+
+                all_indices = cat((all_indices, rearrange(embed_indices, '... -> ... 1')), dim = -1)
+
+                all_losses = cat((all_losses, rearrange(loss, '... -> ... 1')), dim = -1)
+
+        # handle beam search
+
+        if is_beam_search:
+            top_index = search_scores.argmax(dim = -1, keepdim = True)
+
+            quantized_out = batch_select(quantized_out, top_index, '* k d')
+            all_indices = batch_select(all_indices, top_index, '* k l')
+            all_losses = batch_select(all_losses, top_index, '* k l')
+            all_residuals = batch_select(all_residuals, top_index, '* k l d')
+
+            quantized_out, all_indices, all_losses, all_residuals = [t[..., 0, :] for t in (quantized_out, all_indices, all_losses, all_residuals)]
+
+            # handle commit loss, which should be the average
+
+            if exists(mask):
+                all_losses = einx.where('..., ... l,', mask, all_losses, 0.)
+                all_losses = reduce(all_losses, '... l -> l', 'sum') / mask.sum(dim = -1).clamp_min(1e-4)
+            else:
+                all_losses = reduce(all_losses, '... l -> l', 'mean')
+
+            # handle updating ema
+
+            if self.vq_is_ema_updating:
+                for vq, layer_input, indices in zip(self.layers, all_residuals.unbind(dim = -2), all_indices.unbind(dim = -1)): # in the case of quantize dropout, zip will terminate with the shorter sequence, which should be all_residuals
+                    vq.update_ema_indices(layer_input, indices, mask = mask)
 
         # if shared codebook, update ema only at end
 
-        if self.training and self.shared_codebook:
+        if self.training and self.shared_codebook and not is_beam_search:
             shared_layer = first(self.layers)
             shared_layer._codebook.update_ema()
             shared_layer.update_in_place_optimizer()
-            shared_layer.expire_codes_(torch.cat(all_residuals, dim = -2))
+
+            all_codes_for_expire = rearrange(all_residuals, '... n l d -> ... (n l) d')
+            shared_layer.expire_codes_(all_codes_for_expire)
 
         # project out, if needed
 
@@ -415,8 +570,6 @@ class ResidualVQ(Module):
 
         # stack all losses and indices
 
-        all_losses, all_indices = map(partial(torch.stack, dim = -1), (all_losses, all_indices))
-
         ret = (quantized_out, all_indices, all_losses)
 
         if return_all_codes: