class

VQVAEModel

extendsPretrainedModel

VQVAEModel(config: VQVAEConfig)

source edit

Bare discrete-latent auto-encoder — encoder, codebook, decoder.

Implements the VQ-VAE of van den Oord, Vinyals, and Kavukcuoglu, 2017. An image is encoded to a spatial field of continuous vectors, each snapped to its nearest entry in a learned codebook of VQVAEConfig.num_embeddings codes, and decoded back. The quantisation step is non-differentiable, so forward routes the decoder's gradient onto the encoder with the straight-through estimator.

This class computes no loss. Use VQVAEForImageGeneration for the training objective, or read the codebook_loss / commitment_loss terms off the returned VQVAEOutput and combine them yourself.

The pair encode_indices / decode_indices is the tokeniser interface: images in, an integer field out, and back again. Downstream discrete models over the latent grid — the paper's own PixelCNN prior among them — consume exactly that.

Parameters

configVQVAEConfig

Hyperparameters controlling image shape, trunk width, codebook size, and the commitment coefficient. See VQVAEConfig.

Attributes

encodernn.Module

Strided-conv encoder emitting (B, embedding_dim, H', W').

quantizernn.VectorQuantizer

Core-API codebook layer holding (num_embeddings, embedding_dim) entries.

decodernn.Module

Transposed-conv decoder mapping the quantised field back to image space.

Notes

Reference: van den Oord, Vinyals, and Kavukcuoglu, "Neural Discrete Representation Learning", NeurIPS, 2017 (arXiv:1711.00937).

The straight-through estimator used here is

z_q^{\mathrm{st}} = z_e + \mathrm{sg}\!\big[z_q - z_e\big],

which equals $z_q$ in the forward pass and has the identity Jacobian with respect to $z_e$ in the backward pass.

Examples

>>> import lucid
>>> from lucid.models.generative.vqvae import VQVAEConfig, VQVAEModel
>>> cfg = VQVAEConfig(sample_size=32, hidden_channels=32,
...                   residual_hidden_channels=32, embedding_dim=16,
...                   num_embeddings=64)
>>> model = VQVAEModel(cfg).eval()
>>> x = lucid.randn((1, 3, 32, 32))
>>> out = model(x)
>>> out.sample.shape, out.indices.shape
((1, 3, 32, 32), (1, 8, 8))

Used by 2

Constructors

dunder

init

→None

__init__(config: VQVAEConfig)

source edit

Properties

prop

latent_grid_size

→tuple[int, int]

latent_grid_size: tuple[int, int]

source edit

Spatial (H', W') of the discrete latent field.

prop

num_embeddings

→int

num_embeddings: int

source edit

Codebook size $K$ .

Instance methods

decode

→Tensor

decode(z_q: Tensor)

source edit

Decode a quantised latent field back to image space.

decode_indices

→Tensor

decode_indices(indices: Tensor)

source edit

Detokenise an index field (B, H', W') back to an image.

encode

→Tensor

encode(x: Tensor)

source edit

Return the continuous pre-quantisation latent z_e(x).

encode_indices

→Tensor

encode_indices(x: Tensor)

source edit

Tokenise x into codebook indices shaped (B, H', W').

forward

→VQVAEOutput

forward(x: Tensor)

source edit

quantize

→nn.VectorQuantizerOutput

quantize(z_e: Tensor)

source edit

Quantise an (B, D, H', W') latent field.

lucid.nn.VectorQuantizer follows the framework's trailing-axis convention, so the channel axis is moved last for the lookup and moved back afterwards; the returned quantized is in image layout, while indices is (B, H', W').

>>> import lucid >>> from lucid.models.generative.vqvae import VQVAEConfig, VQVAEModel >>> cfg = VQVAEConfig(sample_size=32, hidden_channels=32, ... residual_hidden_channels=32, embedding_dim=16, ... num_embeddings=64) >>> model = VQVAEModel(cfg).eval() >>> x = lucid.randn((1, 3, 32, 32)) >>> out = model(x) >>> out.sample.shape, out.indices.shape ((1, 3, 32, 32), (1, 8, 8))