simvq

Author: lucidrains

examples/autoencoder.py

@@ -75,8 +75,9 @@ def iterate_dataset(data_loader):
 torch.random.manual_seed(seed)
 
 model = SimpleVQAutoEncoder(
-    codebook_size=num_codes,
-    rotation_trick=rotation_trick
+    codebook_size = num_codes,
+    rotation_trick = True,
+    straight_through = False
 ).to(device)
 
 opt = torch.optim.AdamW(model.parameters(), lr=lr)

examples/autoencoder_sim_vq.py

@@ -0,0 +1,84 @@
+# FashionMnist VQ experiment with various settings.
+# From https://github.com/minyoungg/vqtorch/blob/main/examples/autoencoder.py
+
+from tqdm.auto import trange
+
+import torch
+import torch.nn as nn
+from torchvision import datasets, transforms
+from torch.utils.data import DataLoader
+
+from vector_quantize_pytorch import SimVQ, Sequential
+
+lr = 3e-4
+train_iter = 10000
+num_codes = 256
+seed = 1234
+rotation_trick = True
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def SimVQAutoEncoder(**vq_kwargs):
+    return Sequential(
+        nn.Conv2d(1, 16, kernel_size=3, stride=1, padding=1),
+        nn.MaxPool2d(kernel_size=2, stride=2),
+        nn.GELU(),
+        nn.Conv2d(16, 64, kernel_size=3, stride=1, padding=1),
+        nn.MaxPool2d(kernel_size=2, stride=2),
+        SimVQ(dim=64, accept_image_fmap = True, **vq_kwargs),
+        nn.Upsample(scale_factor=2, mode="nearest"),
+        nn.Conv2d(64, 16, kernel_size=3, stride=1, padding=1),
+        nn.GELU(),
+        nn.Upsample(scale_factor=2, mode="nearest"),
+        nn.Conv2d(16, 1, kernel_size=3, stride=1, padding=1),
+    )
+
+def train(model, train_loader, train_iterations=1000, alpha=10):
+    def iterate_dataset(data_loader):
+        data_iter = iter(data_loader)
+        while True:
+            try:
+                x, y = next(data_iter)
+            except StopIteration:
+                data_iter = iter(data_loader)
+                x, y = next(data_iter)
+            yield x.to(device), y.to(device)
+
+    for _ in (pbar := trange(train_iterations)):
+        opt.zero_grad()
+        x, _ = next(iterate_dataset(train_loader))
+
+        out, indices, cmt_loss = model(x)
+        out = out.clamp(-1., 1.)
+
+        rec_loss = (out - x).abs().mean()
+        (rec_loss + alpha * cmt_loss).backward()
+
+        opt.step()
+
+        pbar.set_description(
+            f"rec loss: {rec_loss.item():.3f} | "
+            + f"cmt loss: {cmt_loss.item():.3f} | "
+            + f"active %: {indices.unique().numel() / num_codes * 100:.3f}"
+        )
+
+transform = transforms.Compose(
+    [transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))]
+)
+
+train_dataset = DataLoader(
+    datasets.FashionMNIST(
+        root="~/data/fashion_mnist", train=True, download=True, transform=transform
+    ),
+    batch_size=256,
+    shuffle=True,
+)
+
+print("baseline")
+torch.random.manual_seed(seed)
+
+model = SimVQAutoEncoder(
+    codebook_size=num_codes,
+).to(device)
+
+opt = torch.optim.AdamW(model.parameters(), lr=lr)
+train(model, train_dataset, train_iterations=train_iter)

vector_quantize_pytorch/__init__.py

@@ -6,5 +6,6 @@
 from vector_quantize_pytorch.residual_lfq import ResidualLFQ, GroupedResidualLFQ
 from vector_quantize_pytorch.residual_fsq import ResidualFSQ, GroupedResidualFSQ
 from vector_quantize_pytorch.latent_quantization import LatentQuantize
+from vector_quantize_pytorch.sim_vq import SimVQ
 
 from vector_quantize_pytorch.utils import Sequential

vector_quantize_pytorch/sim_vq.py

@@ -0,0 +1,136 @@
+from typing import Callable
+
+import torch
+from torch import nn
+from torch.nn import Module
+import torch.nn.functional as F
+
+from einx import get_at
+from einops import einsum, rearrange, repeat, reduce, pack, unpack
+
+# helper functions
+
+def exists(v):
+    return v is not None
+
+def identity(t):
+    return t
+
+def default(v, d):
+    return v if exists(v) else d
+
+def pack_one(t, pattern):
+    packed, packed_shape = pack([t], pattern)
+
+    def inverse(out, inv_pattern = None):
+        inv_pattern = default(inv_pattern, pattern)
+        out, = unpack(out, packed_shape, inv_pattern)
+        return out
+
+    return packed, inverse
+
+def l2norm(t, dim = -1):
+    return F.normalize(t, dim = dim)
+
+def safe_div(num, den, eps = 1e-6):
+    return num / den.clamp(min = eps)
+
+def efficient_rotation_trick_transform(u, q, e):
+    """
+    4.2 in https://arxiv.org/abs/2410.06424
+    """
+    e = rearrange(e, 'b d -> b 1 d')
+    w = l2norm(u + q, dim = 1).detach()
+
+    return (
+        e -
+        2 * (e @ rearrange(w, 'b d -> b d 1') @ rearrange(w, 'b d -> b 1 d')) +
+        2 * (e @ rearrange(u, 'b d -> b d 1').detach() @ rearrange(q, 'b d -> b 1 d').detach())
+    )
+
+# class
+
+class SimVQ(Module):
+    def __init__(
+        self,
+        dim,
+        codebook_size,
+        init_fn: Callable = identity,
+        accept_image_fmap = False
+    ):
+        super().__init__()
+        self.accept_image_fmap = accept_image_fmap
+
+        codebook = torch.randn(codebook_size, dim)
+        codebook = init_fn(codebook)
+
+        # the codebook is actually implicit from a linear layer from frozen gaussian or uniform
+
+        self.codebook_to_codes = nn.Linear(dim, dim, bias = False)
+        self.register_buffer('codebook', codebook)
+
+    def forward(
+        self,
+        x
+    ):
+        if self.accept_image_fmap:
+            x = rearrange(x, 'b d h w -> b h w d')
+            x, inverse_pack = pack_one(x, 'b * d')
+
+        implicit_codebook = self.codebook_to_codes(self.codebook)
+
+        with torch.no_grad():
+            dist = torch.cdist(x, implicit_codebook)
+            indices = dist.argmin(dim = -1)
+
+        # select codes
+
+        quantized = get_at('[c] d, b n -> b n d', implicit_codebook, indices)
+
+        # commit loss
+
+        commit_loss = (F.pairwise_distance(x, quantized.detach()) ** 2).mean()
+
+        # straight through
+
+        x, inverse = pack_one(x, '* d')
+        quantized, _ = pack_one(quantized, '* d')
+
+        norm_x = x.norm(dim = -1, keepdim = True)
+        norm_quantize = quantized.norm(dim = -1, keepdim = True)
+
+        rot_quantize = efficient_rotation_trick_transform(
+            safe_div(x, norm_x),
+            safe_div(quantized, norm_quantize),
+            x
+        ).squeeze()
+
+        quantized = rot_quantize * safe_div(norm_quantize, norm_x).detach()
+
+        x, quantized = inverse(x), inverse(quantized)
+
+        # quantized = (quantized - x).detach() + x
+
+        if self.accept_image_fmap:
+            quantized = inverse_pack(quantized)
+            quantized = rearrange(quantized, 'b h w d-> b d h w')
+
+            indices = inverse_pack(indices, 'b *')
+
+        return quantized, indices, commit_loss
+
+# main
+
+if __name__ == '__main__':
+
+    x = torch.randn(1, 512, 32, 32)
+
+    sim_vq = SimVQ(
+        dim = 512,
+        codebook_size = 1024,
+        accept_image_fmap = True
+    )
+
+    quantized, indices, commit_loss = sim_vq(x)
+
+    assert x.shape == quantized.shape

vector_quantize_pytorch/utils.py

@@ -12,6 +12,7 @@
 from vector_quantize_pytorch.residual_lfq import ResidualLFQ, GroupedResidualLFQ
 from vector_quantize_pytorch.residual_fsq import ResidualFSQ, GroupedResidualFSQ
 from vector_quantize_pytorch.latent_quantization import LatentQuantize
+from vector_quantize_pytorch.sim_vq import SimVQ
 
 QUANTIZE_KLASSES = (
     VectorQuantize,
@@ -20,6 +21,7 @@
     RandomProjectionQuantizer,
     FSQ,
     LFQ,
+    SimVQ,
     ResidualLFQ,
     GroupedResidualLFQ,
     ResidualFSQ,