class

VAEModel

extendsPretrainedModel

VAEModel(config: VAEConfig)

source edit

Bare convolutional VAE — vanilla / $\beta$ -VAE / hierarchical.

Implements the variational auto-encoder of Kingma and Welling, 2014 in three operating modes selected implicitly by VAEConfig.latent_dim:

Vanilla / $\beta$ -VAE (int latent_dim): a single bottleneck z is projected from the flattened final encoder stage. encode(x) -> (mu, logvar) with shape (B, D) each; decode(z: (B, D)) -> x_hat.
Hierarchical / Ladder VAE (tuple latent_dim, length $L$ = len(down_block_channels)): one z_l per encoder stage; the decoder injects z_l back at the matching upsampling spatial scale. encode(x) -> (mus: list[Tensor], logvars: list[Tensor]) with mus[l] of shape (B, latent_dims[l]); decode(zs: list[Tensor]) -> x_hat.

Reparameterisation is applied internally in forward — $z = \mu + \sigma \odot \epsilon$ , $\epsilon \sim \mathcal{N}(0, \mathbf{I})$ — so gradients flow pathwise through the stochastic sample.

Parameters

configVAEConfig

Hyperparameters controlling sample size, channel widths (down_block_channels), latent topology (latent_dim), and loss-related options. See VAEConfig for the full list.

Attributes

encodernn.Module

Conv-stack encoder producing the posterior parameters

(\mu, \log \sigma^2)

. Output shape depends on the mode (see class summary).

decodernn.Module

Top-down ConvTranspose decoder mapping the latent (or list of per-level latents) back to image space.

Notes

Reference: Kingma and Welling, "Auto-Encoding Variational Bayes", ICLR, 2014 (arXiv:1312.6114); hierarchical extension from Sønderby, Raiko, Maaløe, Sønderby, and Winther, "Ladder Variational Autoencoders", NeurIPS, 2016 (arXiv:1602.02282).

The closed-form KL term used by the loss head is

\mathrm{KL}\!\big( \mathcal{N}(\mu, \sigma^2) \,\big\|\, \mathcal{N}(0, \mathbf{I}) \big) = \tfrac{1}{2} \sum_i \!\big( \mu_i^2 + \sigma_i^2 - 1 - \log \sigma_i^2 \big).

Examples

>>> import lucid
>>> from lucid.models.generative.vae import VAEConfig, VAEModel
>>> cfg = VAEConfig(sample_size=32, in_channels=3, out_channels=3,
...                 latent_dim=16, down_block_channels=(16, 32, 64))
>>> model = VAEModel(cfg).eval()
>>> x = lucid.randn((1, 3, 32, 32))
>>> out = model(x)
>>> out.sample.shape, out.latent.shape   # reconstruction + latent
((1, 3, 32, 32), (1, 16))

Used by 2

Constructors

dunder

init

→None

__init__(config: VAEConfig)

source edit

Properties

prop

is_hierarchical

→bool

is_hierarchical: bool

source edit

prop

latent_dims

→tuple[int, ...]

latent_dims: tuple[int, ...]

source edit

Instance methods

decode

→Tensor

decode(z_or_zs: Tensor | list[Tensor])

source edit

Map a latent (or list of per-level latents) back to image space.

encode

→tuple[Tensor, Tensor] | tuple[list[Tensor], list[Tensor]]

encode(x: Tensor)

source edit

Return posterior parameters of q(z|x).

Vanilla mode → (mu, logvar). Hierarchical mode → (mus, logvars) with one entry per encoder stage.

forward

→VAEOutput

forward(x: Tensor)

source edit

>>> import lucid >>> from lucid.models.generative.vae import VAEConfig, VAEModel >>> cfg = VAEConfig(sample_size=32, in_channels=3, out_channels=3, ... latent_dim=16, down_block_channels=(16, 32, 64)) >>> model = VAEModel(cfg).eval() >>> x = lucid.randn((1, 3, 32, 32)) >>> out = model(x) >>> out.sample.shape, out.latent.shape # reconstruction + latent ((1, 3, 32, 32), (1, 16))