data

VQVAEConfig

extendsGenerativeModelConfig

VQVAEConfig(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: Literal['silu', 'swish', 'relu', 'gelu'] = 'relu', num_embeddings: int = 512, embedding_dim: int = 256, hidden_channels: int = 256, num_downsample_layers: int = 2, num_residual_layers: int = 2, residual_hidden_channels: int = 256, commitment_cost: float = 0.25, recon_loss: Literal['mse', 'bce'] = 'mse')

source edit

Frozen configuration for the VQ-VAE family.

Defaults reproduce the image experiments of van den Oord et al., 2017 (Section 4.1): two stride-2 encoder convolutions with a 4x4 window, two residual blocks, 256 hidden units throughout, a 512 x 256 codebook, and commitment_cost = 0.25.

Parameters

sample_sizeint or tuple of int= 32

Input spatial resolution. Inherited from GenerativeModelConfig. Must be divisible by 2 ** num_downsample_layers.

in_channelsint= 3

Input image channels.

out_channelsint= 3

Reconstruction channels; normally equal to in_channels.

act_fn(silu, swish, relu, gelu)= "silu"

Trunk activation. The paper uses ReLU throughout, which is why this family overrides the generative-domain default of "silu".

num_embeddingsint= 512

Codebook size

K

— the number of discrete latent codes.

embedding_dimint= 256

Code dimension

D

. The encoder's final projection emits this many channels so encoder outputs and codebook entries live in the same space.

hidden_channelsint= 256

Width of the encoder / decoder trunk.

num_downsample_layersint= 2

Number of stride-2 convolutions. The latent grid is sample_size / 2 ** num_downsample_layers on a side.

num_residual_layersint= 2

Number of residual blocks after the downsampling stack (and, mirrored, before the decoder's upsampling stack).

residual_hidden_channelsint= 256

Inner width of each residual block's 3x3 convolution.

commitment_costfloat= 0.25

The coefficient

\beta

weighting the commitment term.

recon_loss(mse, bce)= "mse"

Reconstruction likelihood — Gaussian ("mse") or Bernoulli ("bce", for data in [0, 1]).

Attributes

latent_grid_sizetuple of int

Spatial (H, W) of the discrete latent field, derived from sample_size and num_downsample_layers.

Notes

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

The paper reports a single image architecture rather than a table of sized variants, so this family exposes one nominal factory rather than _small / _base / _large siblings. Scale it by overriding the fields above at create_model time.

The paper's Appendix A.1 also describes an exponential-moving-average alternative to the codebook loss. That variant changes how the codebook is updated rather than what the network computes, so it belongs to a training loop rather than to the architecture, and is not modelled here.

Examples

>>> from lucid.models.generative.vqvae import VQVAEConfig
>>> cfg = VQVAEConfig(sample_size=32)
>>> cfg.latent_grid_size
(8, 8)
>>> cfg.num_embeddings, cfg.embedding_dim
(512, 256)

Used by 3

Constructors

dunder

init

→None

__init__(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: Literal['silu', 'swish', 'relu', 'gelu'] = 'relu', num_embeddings: int = 512, embedding_dim: int = 256, hidden_channels: int = 256, num_downsample_layers: int = 2, num_residual_layers: int = 2, residual_hidden_channels: int = 256, commitment_cost: float = 0.25, recon_loss: Literal['mse', 'bce'] = 'mse')

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.

VQVAEConfig(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: Literal['silu', 'swish', 'relu', 'gelu'] = 'relu', num_embeddings: int = 512, embedding_dim: int = 256, hidden_channels: int = 256, num_downsample_layers: int = 2, num_residual_layers: int = 2, residual_hidden_channels: int = 256, commitment_cost: float = 0.25, recon_loss: Literal['mse', 'bce'] = 'mse')

__init__(sample_size: int | tuple[int, int] = 32, in_channels: int = 3, out_channels: int = 3, act_fn: Literal['silu', 'swish', 'relu', 'gelu'] = 'relu', num_embeddings: int = 512, embedding_dim: int = 256, hidden_channels: int = 256, num_downsample_layers: int = 2, num_residual_layers: int = 2, residual_hidden_channels: int = 256, commitment_cost: float = 0.25, recon_loss: Literal['mse', 'bce'] = 'mse')