class

VQVAEForImageGeneration

extendsImageGenerationModel

VQVAEForImageGeneration(config: VQVAEConfig)

source edit

VQ-VAE with the full training objective and a codebook sampler.

Wraps VQVAEModel with the three-term objective of van den Oord et al., 2017 — reconstruction, codebook, and $\beta$ -weighted commitment — and a convenience sampler.

forward(x) returns a VQVAEOutput carrying the reconstruction, the quantised latent, the integer code field, the codebook perplexity, and all four loss tensors.

Under recon_loss="bce" the decoder emits logits, which is what the Bernoulli likelihood needs; out.sample and generate().samples are squashed through a sigmoid so every reconstruction handed back lives in the same [0, 1] space as the input, while the loss is still computed from the raw logits.

Parameters

configVQVAEConfig

Hyperparameters. config.commitment_cost is the

\beta

coefficient; config.recon_loss selects the likelihood.

Attributes

vqvaeVQVAEModel

Underlying trunk providing encode / decode and the codebook.

Notes

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

Training objective:

L = \log p(x \mid z_q(x)) + \big\| \mathrm{sg}[z_e(x)] - e \big\|_2^2 + \beta \big\| z_e(x) - \mathrm{sg}[e] \big\|_2^2 .

On generate. The prior over the discrete latents is uniform during training — that is the paper's choice, and it is why the KL term is the constant $\log K$ and drops out of the objective. generate samples from exactly that uniform prior, so it is faithful to what this model was trained against, but it is not how the paper produces its figures: there, a PixelCNN is fit over the latent grid afterwards and sampled autoregressively. Expect incoherent images from the uniform sampler. Fitting a prior is a separate model over encode_indices output, deliberately outside this family.

Examples

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

Used by 2

Constructors

dunder

init

→None

__init__(config: VQVAEConfig)

source edit

Instance methods

forward

→VQVAEOutput

forward(x: Tensor)

source edit

generate

→GenerationOutput

generate(n_samples: int = 1, device: str | None = None)

source edit

Sample n_samples images from the uniform codebook prior.

Draws an index field uniformly over $[0, K)$ and decodes it.

Parameters

n_samplesint= 1

Number of images to draw.

device(str or None, optional, keyword - only)= None

Device to sample on. None resolves to the device the model's parameters already live on.

Returns

GenerationOutput

samples of shape (n_samples, out_channels, H, W), squashed through a sigmoid when recon_loss="bce" so the result matches the space forward reports.

Notes

This is the training-time prior — uniform is what makes the KL term the constant $\log K$ and drops it from the objective. It is not the paper's generative prior, which is a PixelCNN fit over the latent grid afterwards, so these samples are expected to be incoherent. Fitting such a prior is a separate model over VQVAEModel.encode_indices output.

Examples

>>> import lucid
>>> from lucid.models.generative.vqvae import VQVAEConfig
>>> from lucid.models.generative.vqvae import VQVAEForImageGeneration
>>> cfg = VQVAEConfig(sample_size=32, hidden_channels=16,
...                   residual_hidden_channels=16, embedding_dim=8,
...                   num_embeddings=32)
>>> VQVAEForImageGeneration(cfg).eval().generate(2).samples.shape
(2, 3, 32, 32)

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

>>> import lucid >>> from lucid.models.generative.vqvae import VQVAEConfig >>> from lucid.models.generative.vqvae import VQVAEForImageGeneration >>> cfg = VQVAEConfig(sample_size=32, hidden_channels=16, ... residual_hidden_channels=16, embedding_dim=8, ... num_embeddings=32) >>> VQVAEForImageGeneration(cfg).eval().generate(2).samples.shape (2, 3, 32, 32)