Vector-quantized VAE (VQ-VAE) spec

VQ-VAE replaces a continuous latent sample with a discrete codebook lookup. This file exposes the core mechanism in a theorem-friendly way:

1. an encoder produces a continuous latent z_e(x);
2. a code index selects a codebook vector z_q;
3. a decoder reconstructs from z_q;
4. the loss combines reconstruction, codebook, and commitment terms.

The nearest-neighbor assignment is deliberately an explicit Fin numCodes argument. That keeps the spec total and avoids hiding tie-breaking policy in the mathematical layer; runtime code can compute the index however it likes and then pass the verified index into this spec.

Reference:

• van den Oord, Vinyals, and Kavukcuoglu (2017), "Neural Discrete Representation Learning".

The loss definitions below describe scalar values. Training also needs a rule for which parameters receive each gradient. trainingGradients states the latent gradient estimator: reconstruction passes through the encoder, codebook loss updates the selected embedding, and commitment loss updates the encoder. The index remains fixed during that backward pass. This estimator is specified separately from the classical derivative of a nearest-neighbor lookup.

NN.Runtime.Autograd.Model.VqVae implements these boundaries with the runtime's detach operation and a straight-through decoder input. Callers supply the encoder, decoder, and code assignment; there is no fixed-backbone VQ-VAE builder. The theorems in NN.MLTheory.Generative.Latent.VQVAE establish loss decomposition, monotonicity in the commitment weight, and nearest-code minimization for the value definitions.

structure Generative.VQVAE.Encoder (α : Type) (obs latent : Spec.Shape) [TorchLean.Storage α] [Context α] : Type

Encoder producing the pre-quantized latent vector z_e(x).

• forward : TorchLean.Tensor α obs → TorchLean.Tensor α latent

  Continuous encoder output before codebook lookup.

structure Generative.VQVAE.Decoder (α : Type) (latent obs : Spec.Shape) [TorchLean.Storage α] [Context α] : Type

Decoder mapping a codebook vector back to observation space.

• forward : TorchLean.Tensor α latent → TorchLean.Tensor α obs

  Decode a quantized latent vector.

structure Generative.VQVAE.Model (α : Type) (obs latent : Spec.Shape) (numCodes : ℕ) [TorchLean.Storage α] [Context α] : Type

VQ-VAE model: encoder, codebook, and decoder.

• encoder : Encoder α obs latent

  Continuous encoder.

• codebook : Latent.Codebook α numCodes latent

  Finite codebook.

• decoder : Decoder α latent obs

  Decoder from quantized latent vectors.

def Generative.VQVAE.encode {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (x : TorchLean.Tensor α obs) : TorchLean.Tensor α latent

Pre-quantized latent z_e(x).

def Generative.VQVAE.quantized {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (idx : Fin numCodes) : TorchLean.Tensor α latent

Quantized latent z_q, using an explicit code index.

def Generative.VQVAE.forward {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (_x : TorchLean.Tensor α obs) (idx : Fin numCodes) : TorchLean.Tensor α obs

VQ-VAE reconstruction from an explicit code assignment.

def Generative.VQVAE.reconstructionLoss {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (x : TorchLean.Tensor α obs) (idx : Fin numCodes) : α

Mean squared reconstruction error between dec(z_q) and the observation.

def Generative.VQVAE.codebookLoss {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (x : TorchLean.Tensor α obs) (idx : Fin numCodes) : α

Mean squared distance from the selected embedding to the encoder output.

This definition records the loss value. The training program treats the encoder output as a fixed target for this term, so only the selected codebook embedding receives its gradient.

def Generative.VQVAE.commitmentLoss {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (x : TorchLean.Tensor α obs) (idx : Fin numCodes) : α

Mean squared distance from the encoder output to the selected embedding.

The training program treats the embedding as a fixed target for this term. Its gradient therefore updates the encoder alone, with weight β in the total objective.

def Generative.VQVAE.loss {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (beta : α) (x : TorchLean.Tensor α obs) (idx : Fin numCodes) : α

VQ-VAE objective: reconstruction + codebook + β commitment.

structure Generative.VQVAE.TrainingGradients (α : Type) (latent : Spec.Shape) [TorchLean.Storage α] : Type

Cotangents at the encoder output and the selected codebook embedding for one assignment.

• encoder : TorchLean.Tensor α latent

  Reconstruction cotangent plus the weighted commitment contribution.

• codebook : TorchLean.Tensor α latent

  Codebook-loss contribution; other embeddings receive zero for this assignment.

def Generative.VQVAE.trainingGradients {α : Type} [TorchLean.Storage α] [Context α] {latent : Spec.Shape} (encoded selected reconstructionCotangent : TorchLean.Tensor α latent) (beta : α) : TrainingGradients α latent

The VQ-VAE latent gradient estimator for a fixed code assignment.

Let g be the reconstruction cotangent at the decoder input and n the number of latent coordinates. The straight-through rule sends g to the encoder and zero to the embedding. Adding the two auxiliary terms gives encoder = g + β * (2 / n) * (encoded - selected) and codebook = (2 / n) * (selected - encoded).

mseDerivSpec supplies the same mean reduction as the value losses, including their empty-shape convention. These are prescribed training signals. A hard nearest-neighbor assignment does not have this classical derivative; the runtime constructs the estimator with explicit stop-gradient boundaries. Decoder parameters receive their usual reconstruction gradients outside this pair.

@[simp]

theorem Generative.VQVAE.quantized_eq_embedding {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (idx : Fin numCodes) : quantized model idx = model.codebook.embedding idx

Quantization by explicit index is exactly codebook lookup.

@[simp]

theorem Generative.VQVAE.loss_eq_reconstruction_add_codebook_add_commitment {α : Type} [TorchLean.Storage α] [Context α] {obs latent : Spec.Shape} {numCodes : ℕ} (model : Model α obs latent numCodes) (beta : α) (x : TorchLean.Tensor α obs) (idx : Fin numCodes) : loss model beta x idx = reconstructionLoss model x idx + codebookLoss model x idx + beta * commitmentLoss model x idx

The VQ-VAE objective decomposes into the three standard terms.