Vector Quantization (VQ-VAE) kernel

[superposition]

Description

Vector quantization for discrete latent spaces in world models.

Operation

# Quantize continuous vectors to codebook entries
# z: (B, T, D) continuous latents
# codebook: (K, D) learnable embeddings
quantized, indices, commitment_loss = vector_quantize(z, codebook)

Implementation

• Efficient L2 distance to all codebook entries
• Straight-through estimator for gradients
• EMA codebook update option
• Codebook utilization tracking

Tests

• Output matches nearest codebook entry
• Straight-through gradient correct (grad flows to encoder)
• Commitment loss computed correctly
• EMA update modifies codebook
• Handles codebook collapse (utilization metric)
• gradcheck on encoder path

Used In

VQ-VAE, VQGAN, Genie, DALL-E, world models