VQ-VAE: Discrete Latent Spaces

VQ-VAE: Discrete Latent Spaces

Standard variational autoencoders (VAEs) force the latent space to be continuous, typically a vector of Gaussian-distributed variables. While this works well for images and smooth signals, many data types in the pharmaceutical and biotech sciences are fundamentally discrete : molecular structures encoded as SMILES strings, amino acid sequences, nucleotide bases, and categorical drug descriptors. A continuous latent vector has no natural way to represent "the third codon in this sequence," which motivates the Vector Quantised Variational Autoencoder (VQ-VAE) . The core innovation of VQ-VAE is the codebook (also called an embedding space): a learnable dictionary of K discrete vectors, each of dimension D . During encoding, the encoder network produces a continuous output, but instead of feeding it directly to the decoder, each encoder output vector is snapped (quantised) to its nearest neighbour in the codebook. The decoder then receives only these quantised vectors. Conceptually, the latent space becomes a finite alphabet of symbols — much like the 20 amino acids or the characters of a SMILES string — rather than an infinite continuum. Training VQ-VAEs requires care because the nearest-neighbour lookup is a discrete, non-differentiable operation. The model therefore uses three separate loss terms. First, a standard reconstruction loss trains the encoder and decoder through the quantised vectors. Second, a codebook loss moves codebook vectors toward the encoder outputs that selected them (implemented as a stopping gradient on the encoder side). Third, a commitment loss prevents the encoder outputs from drifting wildly between codebook entries by penalising the distance between the encoder output and its chosen code. A common practical issue is codebook collapse , where most codes go unused; techniques such as EMA updates and codebook restarting help keep all entries active. For drug discovery applications, the discrete latent space is a major advantage. Each molecule can be represented as a compact sequence of codebook indices, which can be modelled with autoregressive sequence models (such as Transformers) to generate novel candidates. This two-stage design — VQ-VAE first, sequence model second — underlies successful molecular generative frameworks and makes latent space interpolation and structure–activity exploration far more interpretable than in continuous VAEs. The following code sketch defines a minimal VQ-VAE layer that performs vector quantisation with straight-through gradient estimation: import torch import torch.nn as nn import torch.nn.functional as F class VectorQuantizer(nn.Module): def __init__(self, num_codes, code_dim, beta=0.25): super().__init__() self.codebook = nn.Embedding(num_codes, code_dim) self.codebook.weight.data.uniform_(-1/num_codes, 1/num_codes) self.beta = beta # commitment loss weight def forward(self, z): # z shape: (batch, latent_dim) -> flatten for distance calc flat = z.reshape(-1, z.size(-1)) # Squared Euclidean distances to every codebook entry dists = (flat.pow(2).sum(1, keepdim=True) - 2 * flat @ self.codebook.weight.t() + self.codebook.weight.pow(2).sum(1)) indices = dists.argmin(dim=1) # nearest code z_q = self.codebook(indices).view_as(z) # Straight-through estimator: copy gradients to the encoder z_q = z + (z_q - z).detach() # Codebook loss (moves codes toward encoder outputs) codebook_loss = F.mse_loss(z_q.detach(), z) # Commitment loss (keeps encoder outputs near chosen codes) commit_loss = F.mse_loss(z_q, z.detach()) loss = codebook_loss + self.beta * commit_loss return z_q, loss, indices Once quantisation is in place, the full VQ-VAE wraps an encoder, the quantiser, and a decoder into one trainable model: class VQVAE(nn.Module): def __init__(self, input_dim=256, hidden=128, num_codes=64, code_dim=32): super().__init__() self.encoder = nn.Sequential( nn.Linear(input_dim, hidden), nn.ReLU(), nn.Linear(hidden, code_dim)) self.quantizer = VectorQuantizer(num_codes, code_dim) self.decoder = nn.Sequential( nn.Linear(code_dim, hidden), nn.ReLU(), nn.Linear(hidden, input_dim), nn.Sigmoid()) def forward(self, x): z_e = self.encoder(x) # continuous encoder output z_q, vq_loss, indices = self.quantizer(z_e) x_hat = self.decoder(z_q) # decode quantised codes recon_loss = F.binary_cross_entropy(x_hat, x, reduction="mean") return x_hat, recon_loss + vq_loss, indices # Training step for a batch of molecular fingerprints model = VQVAE() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) x_hat, loss, codes = model(fingerprint_batch) loss.backward() optimizer.step() print("Latent codes per sample:", codes[:8]) VQ-VAE replaces the continuous Gaussian latent space of a standard VAE with a learnable codebook of discrete vectors, matching the discrete nature of SMILES strings, protein sequences, and categorical descriptors. Quantisation is non-differentiable, so training relies on the straight-through estimator plus separate codebook and commitment losses. Codebook collapse (unused codes) is the main training failure mode; EMA updates and code restarts are standard remedies. Discrete latent codes pair naturally with autoregressive priors, enabling powerful two-stage generative pipelines for de novo molecular design. Because each sample maps to a short sequence of integer codes, VQ-VAE latents are compact, interpretable, and well suited to combinatorial search over chemical space. Practice exercise: Modify the VQVAE class above so that it encodes a batch of one-hot encoded amino acid sequences (length 50, 20 residue types) instead of molecular fingerprints. Train it for a few epochs on a small synthetic dataset of random sequences, then inspect the indices tensor: How many of the 64 codebook entries are actually used across your training batches? Propose one concrete strategy to reduce codebook collapse if you observe that fewer than half the codes are active, and explain why it should work.