VQ-VAE

1. 1

   En generativ model der ved hjælp af vektorkvantisering lærer en diskret kodebog af latente repræsentationer. Encoderen producerer en kontinuert vektor, som kvantiseres til den nærmeste kodebogsvektor, og decoderen genererer data ud fra denne diskrete kode. Træning bruger en straight-through estimator til at propagere gradienter gennem det ikke-differentierbare kvantiseringstrin.

   • VQ-VAE har vist imponerende resultater inden for billedgenerering ved at lære en diskret kodebog på 512 indgange. — van den Oord et al., 2017
   • VQ-VAE bruges som grundlag for DALL-E's dVAE-komponent. — DALL-E paper, 2021

Bruges i generative modeller hvor diskrete repræsentationer er vigtige, fx i DALL-E (som en VQ-VAE komponent) og i lydgenerering som WaveNet. Vektorkvantiseringslaget tvinger encoderen til at producere en embedding, der matches til nærmeste kodebogsvektor.

Loss = ||x - decoder(z_q)||^2 + ||sg[z_e] - e||^2 + β||z_e - sg[e]||^2

class VectorQuantizer(nn.Module):
    def __init__(self, num_embeddings, embedding_dim):
        super().__init__()
        self.embedding = nn.Embedding(num_embeddings, embedding_dim)
        self.embedding.weight.data.uniform_(-1/num_embeddings, 1/num_embeddings)

    def forward(self, z):
        z_flat = z.permute(0,2,3,1).reshape(-1, z.size(1))
        distances = torch.cdist(z_flat, self.embedding.weight)
        indices = distances.argmin(1)
        z_q = self.embedding(indices).reshape(z.shape)
        return z_q + (z - z_q).detach()

Simpel PyTorch-implementering af vektorkvantiseringslaget i VQ-VAE.

VQ-VAE står for Vector Quantized Variational Autoencoder, en kombination af VAE (Variational Autoencoder) og vektorkvantisering fra signalbehandling.

• Neural Discrete Representation Learning