Vector Quantized Variational Autoencoder

The problem, from first principles

VQ-VAE learns discrete latent codes by quantizing continuous latent representations. An encoder maps input to a continuous space, which is quantized to the nearest code in a learned codebook. A decoder reconstructs from the quantized code.

Loss combines reconstruction, codebook, and commitment terms: $$\mathcal{L}=\mathrm{\mid }\mathrm{\mid }x-\widehat{x}\mathrm{\mid }{\mathrm{\mid }}^2+\mathrm{\mid }\mathrm{\mid }sg[z_e]-e\mathrm{\mid }{\mathrm{\mid }}^2+\beta \mathrm{\mid }\mathrm{\mid }{z}_e-sg[e]\mathrm{\mid }{\mathrm{\mid }}^2\backslash mathcal\{L\}\; =\; ||x\; -\; \backslash hat\{x\}||^2\; +\; ||sg[z\_e]\; -\; e||^2\; +\; \backslash beta||z\_e\; -\; sg[e]||^2$$

Where sg denotes stop-gradient, z_e is the encoder output, and e is the quantized code.

From theory to code

Implement:

vq_vae_loss(x_recon, x_true, z_e, z_q, beta=0.25)

Computes reconstruction, codebook, and commitment losses.

Constraints

• x_recon, x_true shape: (N, D)
• z_e (encoder), z_q (quantized) shape: (N, D)
• beta: commitment loss weight (default 0.25).
• Return scalar loss.