Spaces:
Running on Zero
Running on Zero
| # Finite Scalar Quantization from https://arxiv.org/pdf/2309.15505 | |
| # quantizes each dimension independently by bounding to 0, num_bins then rounding to nearest integer | |
| # prevents token collapse and no auxiliary losses necessary | |
| import torch | |
| import torch.nn as nn | |
| from einops import rearrange | |
| class FiniteScalarQuantizer(nn.Module): | |
| def __init__(self, latent_dim=5, num_bins=4): | |
| super().__init__() | |
| self.num_bins = num_bins # D | |
| self.levels_np = torch.tensor(latent_dim * [num_bins]) | |
| self.codebook_size = num_bins**latent_dim # L^D | |
| # fsq basis [L^0, L^1, ..., L^(L-1)] for converting between indices and latents | |
| self.register_buffer('basis', (num_bins**torch.arange(latent_dim, dtype=torch.long))) | |
| def scale_and_shift(self, z): | |
| # scale and shift z from [-1, 1] to [0, num_bins - 1] | |
| return 0.5 * (z + 1) * (self.num_bins - 1) | |
| def unscale_and_unshift(self, z): | |
| # unscale and unshift z from [0, num_bins - 1] to [-1, 1] | |
| return 2 * z / (self.num_bins - 1) - 1 | |
| def forward(self, z): | |
| # z: [B, T, P, L] | |
| # apply 0.5 * (tanh(z) + 1) to go from z range to [0, num_bins - 1] | |
| tanh_z = torch.tanh(z) | |
| bounded_z = self.scale_and_shift(tanh_z) | |
| # round to nearest integer | |
| rounded_z = torch.round(bounded_z) | |
| # stopgrad for straight-through gradient bypassing quantization | |
| quantized_z = bounded_z + (rounded_z - bounded_z).detach() | |
| # normalize back to [-1, 1] | |
| quantized_z = self.unscale_and_unshift(quantized_z) | |
| return quantized_z | |
| def get_codebook_usage(self, quantized_z): | |
| unique_bins = torch.unique(quantized_z).shape[0] | |
| return unique_bins / self.num_bins | |
| def get_indices_from_latents(self, latents, dim=-1): | |
| # to get fsq indices, for each dimension, we get the index and add it (so multiply by L then sum) | |
| # for each dimension of each latent, get sum of (value * L^current_dim) along latent dim which is the index of that latent in the codebook | |
| # codebook size = L^latent_dim | |
| # latents: [*, L] | |
| # go from [-1, 1] to [0, num_bins - 1] in each dimension | |
| digits = torch.round(self.scale_and_shift(latents)).clamp(0, self.num_bins-1) | |
| # get indices for each latent by summing (value * L^current_dim_idx) along latent dim | |
| indices = torch.sum(digits * self.basis.to(latents.device), dim=dim).long() # [*] | |
| return indices | |
| def get_latents_from_indices(self, indices, dim=-1): | |
| # indices: [*] | |
| # recover each entry of latent in range [0, num_bins - 1] by repeatedly dividing by L^current_dim and taking mod | |
| digits = (indices.unsqueeze(-1) // self.basis) % self.num_bins # [*, L] | |
| # go from [0, num_bins - 1] to [-1, 1] in each dimension | |
| latents = self.unscale_and_unshift(digits) # [*, L] | |
| return latents | |