Autoencoder Models

AE and VAE share build, encode and decode; VAE.encode returns the posterior mean by default. Build your model by passing your own encoder and decoder to these classes: the base classes they share are internal and not meant to be subclassed.

Vanilla Autoencoder

class pyautoencoder.vanilla.AE(encoder: Module, decoder: Module)

Vanilla Autoencoder composed of a user-defined encoder and decoder.

The model implements:

• _encode(x) – maps inputs x to latent codes z.

• _decode(z) – maps latent codes z to reconstructions x_hat.

• forward(x) – full training forward pass returning both z and x_hat.

The encoder and decoder are arbitrary torch.nn.Module instances that define the mapping between data space and latent space.

__init__(encoder: Module, decoder: Module)

Construct an Autoencoder from an encoder and decoder module.

Parameters:

• encoder (nn.Module) – Module implementing the mapping x → z.

• decoder (nn.Module) – Module implementing the mapping z → x_hat. Its output must have the shape of the input x, as compute_loss() requires.

build(input_sample: Tensor) → None

Materialize lazy layers with a representative input.

Performs a no-grad forward pass so that any LazyModuleMixin layers infer their shapes. Call this once before training, e.g. before creating the optimizer or applying a custom weight initialization.

The pass runs in eval mode, so it does not update BatchNorm running statistics and does not apply dropout. The train/eval mode of every submodule is restored afterwards.

Parameters:

input_sample (torch.Tensor) – A representative input batch (e.g., a single batch from the training set). Only its shape matters.

compute_loss(x: Tensor, output: AEOutput, likelihood: str | LikelihoodType = LikelihoodType.GAUSSIAN) → LossResult

Compute Autoencoder reconstruction loss.

The scalar loss is the batch-mean reconstruction negative log-likelihood (NLL). The method also computes diagnostics to monitor model behavior.

Parameters:

• x (torch.Tensor) – Ground-truth inputs of shape [B, ...].

• output (AEOutput) –

  Output from the AE forward pass. Expected fields include:

  • x_hat (torch.Tensor): Reconstructions, shape [B, ...].

  • z (torch.Tensor): Latent representation (unused by this method).

• likelihood (str | LikelihoodType, optional) – Likelihood model for computing the reconstruction term ('gaussian' or 'bernoulli'). Defaults to Gaussian.

Returns:

Result containing:

• objective – Scalar batch-mean reconstruction NLL (in nats).

• diagnostics – Dictionary of detached scalar tensors with:

  • "log_likelihood": Negative of the objective (batch-mean log-likelihood).

Return type:

LossResult

Raises:

ValueError – If output.x_hat and x do not have the same shape, or if likelihood is not a supported likelihood model.

Notes

Reductions follow:

1. Elementwise log-likelihood.

2. Sum over feature dimensions.

3. Mean over the batch.

Ensure that inputs match the chosen likelihood:

• Gaussian: continuous data (typically standardized).

• Bernoulli: targets in \([0, 1]\), predictions given as logits.

decode(z: Tensor, *args: Any, **kwargs: Any) → ModelOutput

Decode latent codes without tracking gradients.

This is a thin wrapper around _decode() executed under torch.no_grad(), making it suitable for evaluation-time decoding. It does not change the train/eval mode: call model.eval() first if the model has layers such as dropout or batch normalization.

Parameters:

• z (torch.Tensor) – Latent representation batch to decode.

• *args – Additional positional arguments passed to _decode().

• **kwargs – Additional keyword arguments passed to _decode().

Returns:

The output of _decode(), with at least the reconstruction x_hat: AEDecodeOutput for AE, VAEDecodeOutput for VAE.

Return type:

ModelOutput

encode(x: Tensor, *args: Any, **kwargs: Any) → ModelOutput

Encode inputs without tracking gradients.

This is a thin wrapper around _encode() executed under torch.no_grad(), making it suitable for evaluation-time encoding. It does not change the train/eval mode: call model.eval() first if the model has layers such as dropout or batch normalization.

Parameters:

• x (torch.Tensor) – Input batch to encode.

• *args – Additional positional arguments passed to _encode().

• **kwargs – Additional keyword arguments passed to _encode().

Returns:

The output of _encode(), with at least the latent code z: AEEncodeOutput for AE.

Return type:

ModelOutput

forward(x: Tensor) → AEOutput

Full training forward pass with gradients.

Parameters:

x (torch.Tensor) – Input batch of shape [B, ...].

Returns:

Dataclass containing both the reconstruction x_hat and the latent code z.

Return type:

AEOutput

Data Structures

class pyautoencoder.vanilla.AEEncodeOutput(z: Tensor)

Output of the Autoencoder encoder stage.

z

Latent code of shape [B, ...] produced by AE._encode() or AE.encode().

Type:

torch.Tensor

class pyautoencoder.vanilla.AEDecodeOutput(x_hat: Tensor)

Output of the Autoencoder decoder stage.

x_hat

Reconstruction (or logits) of shape [B, ...] produced by AE._decode() or AE.decode().

Type:

torch.Tensor

class pyautoencoder.vanilla.AEOutput(x_hat: Tensor, z: Tensor)

Output of the full Autoencoder forward pass.

x_hat

Reconstruction (or logits) of shape [B, ...] produced by AE.forward().

Type:

torch.Tensor

z

Latent code of shape [B, ...] produced by AE.forward().

Type:

torch.Tensor

Variational Autoencoder

class pyautoencoder.variational.VAE(encoder: Module, decoder: Module, latent_dim: int)

Variational Autoencoder following Kingma & Welling (2013).

The model consists of:

• an encoder mapping x → f(x) (feature representation),

• a fully factorized Gaussian head producing (z, mu, log_var),

• a decoder mapping latent samples z → x_hat.

Sampling is controlled explicitly by the sample argument, not by train/eval mode (which only affects layers such as dropout and batch normalization):

• forward() / _encode() sample z from \(q(z \mid x)\) by default (sample=True), as needed for training and for estimating the ELBO, also on validation data.

• encode() returns the posterior mean z = mu by default (sample=False), as usually needed at inference time.

Example

out = vae(x, S=5)                     # training: 5 samples per input
loss = vae.compute_loss(x, out)

vae.eval()
with torch.no_grad():                 # validation ELBO
    val = vae.compute_loss(x_val, vae(x_val))

enc = vae.encode(x)                   # inference: z = mu, [B, 1, D_z]
x_rec = vae.decode(enc.z).x_hat       # deterministic reconstruction

Shape convention

Latent codes and reconstructions always carry an explicit Monte Carlo sample dimension S right after the batch dimension, also when S = 1:

• z: [B, S, D_z]

• x_hat: [B, S, ...]

• mu, log_var: [B, D_z] (one posterior per input, no S)

Consequently, decode() expects z of shape [B, S, D_z]. To decode a batch of single codes [N, D_z] (e.g. samples from the prior \(p(z)\)), use z.unsqueeze(1) or call model.decoder directly.

__init__(encoder: Module, decoder: Module, latent_dim: int)

Construct a Variational Autoencoder from an encoder, decoder, and latent size.

Parameters:

• encoder (nn.Module) – Maps input x to a feature vector f(x) with shape [B, F].

• decoder (nn.Module) – Maps latent codes of shape [N, D_z] to reconstructions (or logits) x_hat. Its output for one code must have the shape of one input, x.shape[1:], as compute_loss() requires.

• latent_dim (int) – Dimensionality D_z of the latent space.

Notes

A FullyFactorizedGaussian sampling layer is created internally as self.sampling_layer; it is not a constructor parameter.

build(input_sample: Tensor) → None

Materialize lazy layers with a representative input.

Performs a no-grad forward pass so that any LazyModuleMixin layers infer their shapes. Call this once before training, e.g. before creating the optimizer or applying a custom weight initialization.

The pass runs in eval mode, so it does not update BatchNorm running statistics and does not apply dropout. The train/eval mode of every submodule is restored afterwards.

Parameters:

input_sample (torch.Tensor) – A representative input batch (e.g., a single batch from the training set). Only its shape matters.

compute_loss(x: Tensor, output: VAEOutput, beta: float = 1, likelihood: str | LikelihoodType = LikelihoodType.GAUSSIAN) → LossResult

Compute the Evidence Lower Bound (ELBO) for a (beta-)Variational Autoencoder.

This method implements the beta-VAE objective:

\[\mathcal{L}(x; \beta) = \mathbb{E}_{q(z \mid x)}[\log p(x \mid z)] \;-\; \beta \, \mathrm{KL}(q(z \mid x) \,\|\, p(z)).\]

The reconstruction term \(\log p(x \mid z)\) is computed using log_likelihood(), which supports both Gaussian and Bernoulli likelihoods.

Monte Carlo estimation

x_hat in output contains S Monte Carlo samples per input (S = 1 included), and the expectation \(\mathbb{E}_{q(z \mid x)}\) is approximated by:

\[\mathbb{E}_{q(z \mid x)}[\log p(x \mid z)] \approx \frac{1}{S} \sum_{s=1}^{S} \log p(x \mid z^{(s)}).\]

Broadcasting

• x_hat must have shape [B, S, *x.shape[1:]]: the dimension after the batch is always the sample dimension, also when S = 1, and it is never inferred from the shapes.

• x is broadcast over the sample dimension only, without copying data, so sample s of input i is always compared with x[i].

Parameters:

• x (torch.Tensor) – Ground-truth inputs of shape [B, ...].

• output (VAEOutput) –

  Output from the VAE forward pass. Expected fields include:

  • x_hat (torch.Tensor): Reconstructed samples, shape [B, S, ...], where ... is the shape of one input, x.shape[1:].

  • mu (torch.Tensor): Mean of \(q(z \mid x)\), shape [B, D_z] with D_z = latent_dim.

  • log_var (torch.Tensor): Log-variance of \(q(z \mid x)\), shape [B, D_z].

• beta (float, optional) – Weighting factor for the KL term (beta-VAE). beta = 1 yields the standard VAE. Defaults to 1.

• likelihood (str | LikelihoodType, optional) – Likelihood model for the reconstruction term ('gaussian' or 'bernoulli'). Defaults to Gaussian.

Returns:

Result containing:

• objective – Negative batch mean of \(\mathcal{L}(x; \beta)\) (the negative ELBO when beta = 1), scalar.

• diagnostics – Dictionary of detached scalar tensors with:

  • "elbo": Batch mean of \(\mathcal{L}(x; \beta)\) (the ELBO when beta = 1).

  • "log_likelihood": Mean reconstruction term \(\mathbb{E}_{q}[\log p(x \mid z)]\).

  • "kl_divergence": Mean \(\mathrm{KL}(q \,\|\, p)\) over the batch (not multiplied by beta).

Return type:

LossResult

Raises:

ValueError – If x_hat is not of shape [B, S, *x.shape[1:]] with S >= 1, if mu and log_var are not both of shape [B, latent_dim], or if likelihood is not a supported likelihood model.

Notes

• All returned diagnostics are batch means.

• objective is not detached: gradients flow through x_hat, mu and log_var to both the decoder and the encoder.

• "elbo" is a Monte Carlo estimate of the ELBO only if output was computed with sample=True (the default of forward(), also in eval mode). With sample=False the reconstruction term is evaluated at z = mu, which is not an ELBO estimate.

• With the Gaussian likelihood, "elbo" and "log_likelihood" omit the normalization constant of \(p(x \mid z)\) (see log_likelihood()).

decode(z: Tensor, *args: Any, **kwargs: Any) → ModelOutput

Decode latent codes without tracking gradients.

This is a thin wrapper around _decode() executed under torch.no_grad(), making it suitable for evaluation-time decoding. It does not change the train/eval mode: call model.eval() first if the model has layers such as dropout or batch normalization.

Parameters:

• z (torch.Tensor) – Latent representation batch to decode.

• *args – Additional positional arguments passed to _decode().

• **kwargs – Additional keyword arguments passed to _decode().

Returns:

The output of _decode(), with at least the reconstruction x_hat: AEDecodeOutput for AE, VAEDecodeOutput for VAE.

Return type:

ModelOutput

encode(x: Tensor, S: int = 1, sample: bool = False) → VAEEncodeOutput

Encode inputs without tracking gradients (inference API).

Same as _encode() under torch.no_grad(), but with an inference-friendly default: z is the posterior mean (sample=False). Pass sample=True to draw S samples.

It does not change the train/eval mode: call model.eval() first if the encoder has layers such as dropout or batch normalization, otherwise the posterior mean is not deterministic.

Parameters:

• x (torch.Tensor) – Input batch of shape [B, ...].

• S (int, optional) – Number of latent samples per input; must be 1 unless sample=True. Defaults to 1.

• sample (bool, optional) – Whether to sample z from \(q(z \mid x)\). Defaults to False (z = mu).

Returns:

Contains z of shape [B, S, D_z] ([B, 1, D_z] by default), and mu and log_var of shape [B, D_z].

Return type:

VAEEncodeOutput

Raises:

ValueError – If S < 1, if S != 1 with sample=False, or if the encoder output is not of shape [B, F].

forward(x: Tensor, S: int = 1, sample: bool = True) → VAEOutput

Full VAE pass: encode, draw S latent codes (or take the posterior mean), decode.

Parameters:

• x (torch.Tensor) – Input batch of shape [B, ...].

• S (int, optional) – Number of latent samples for Monte Carlo estimates. Must be >= 1, and 1 when sample=False. Defaults to 1.

• sample (bool, optional) – Whether to sample z from \(q(z \mid x)\) (True) or to use the posterior mean z = mu (False). Defaults to True, in both train and eval mode.

Returns:

Contains reconstructions x_hat of shape [B, S, ...], latent codes z of shape [B, S, D_z], and the posterior parameters mu and log_var of shape [B, D_z].

Return type:

VAEOutput

Raises:

ValueError – If S < 1, if S != 1 with sample=False, or if the encoder output is not of shape [B, F].

Notes

If S > 1, loss computation can broadcast x to shape [B, S, ...] without materializing copies. For Bernoulli likelihoods, the decoder must output logits.

Data Structures

class pyautoencoder.variational.VAEEncodeOutput(z: Tensor, mu: Tensor, log_var: Tensor)

Output of the VAE encoder stage.

z

Latent codes of shape [B, S, D_z], produced by VAE._encode() or VAE.encode(): S samples from \(q(z \mid x)\), or the posterior mean (S = 1) with sample=False, which is the default of VAE.encode().

Type:

torch.Tensor

mu

Mean of the approximate posterior \(q(z \mid x)\), shape [B, D_z].

Type:

torch.Tensor

log_var

Log-variance of \(q(z \mid x)\), shape [B, D_z].

Type:

torch.Tensor

class pyautoencoder.variational.VAEDecodeOutput(x_hat: Tensor)

Output of the VAE decoder stage.

x_hat

Reconstructions or logits of shape [B, S, ...], produced by VAE._decode() or VAE.decode().

Type:

torch.Tensor

class pyautoencoder.variational.VAEOutput(x_hat: Tensor, z: Tensor, mu: Tensor, log_var: Tensor)

Output of a full VAE forward pass.

x_hat

Reconstructions or logits of shape [B, S, ...], produced by VAE.forward().

Type:

torch.Tensor

z

Latent codes of shape [B, S, D_z], produced by VAE.forward(): S samples from \(q(z \mid x)\), or the posterior mean (S = 1) with sample=False.

Type:

torch.Tensor

mu

Mean of \(q(z \mid x)\), shape [B, D_z].

Type:

torch.Tensor

log_var

Log-variance of \(q(z \mid x)\), shape [B, D_z].

Type:

torch.Tensor

Adaptive Group Variational Autoencoder

class pyautoencoder.variational.AdaGVAE(vae: VAE)

Bases: Module

Adaptive Group Variational Autoencoder (Ada-GVAE), from Locatello et al. (2020).

Wraps a VAE and adds adaptive posterior grouping for feature disentanglement. All VAE parameters are tracked through this wrapper.

forward() expects a pair of inputs (x1, x2) and returns an AdaGVAEOutput for both. The adaptive grouping only changes training: in eval mode no grouping is applied and AdaGVAE behaves as the wrapped VAE, passing each input independently through VAE.forward(). For single-image inference after training, use model.vae.encode and model.vae.decode directly.

__init__(vae: VAE)

Wrap a VAE with adaptive posterior grouping.

Parameters:

vae (VAE) – A configured VAE instance whose encoder, decoder, and sampling layer are reused for the paired training objective.

build(input_sample: Tensor) → None

Materialize lazy layers by delegating to the wrapped VAE’s build step.

Parameters:

input_sample (torch.Tensor) – A representative batch of single inputs (not a pair), passed to VAE.build(). Only its shape matters.

compute_loss(x: tuple[Tensor, Tensor], output: AdaGVAEOutput, beta: float = 1, likelihood: str | LikelihoodType = LikelihoodType.GAUSSIAN) → LossResult

Compute the combined ELBO for a pair of inputs.

In training mode the two terms use the adapted posteriors \(q(\hat{z} \mid x_i)\) produced by forward(); in eval mode they use the plain VAE posteriors.

\[\mathcal{L}(x_1, x_2; \beta) = \left[ \mathbb{E}_{q(\hat{z} \mid x_1)}[\log p(x_1 \mid \hat{z})] \;-\; \beta \, \mathrm{KL}(q(\hat{z} \mid x_1) \,\|\, p(\hat{z})) \right] + \left[ \mathbb{E}_{q(\hat{z} \mid x_2)}[\log p(x_2 \mid \hat{z})] \;-\; \beta \, \mathrm{KL}(q(\hat{z} \mid x_2) \,\|\, p(\hat{z})) \right].\]

Parameters:

• x (tuple[torch.Tensor, torch.Tensor]) – The (x1, x2) pair of ground-truth inputs, each of shape [B, ...].

• output (AdaGVAEOutput) – Output from forward().

• beta (float, optional) – KL weighting factor. beta = 1 yields the standard objective. Defaults to 1.

• likelihood (str | LikelihoodType, optional) – Likelihood model for the reconstruction term ('gaussian' or 'bernoulli'). Defaults to Gaussian.

Returns:

Result containing:

• objective – Sum of the objective values of VAE.compute_loss() for x1 and x2 (scalar).

• diagnostics – Dictionary of detached scalar tensors with:

  • "elbo": Sum of the two "elbo" diagnostics of VAE.compute_loss() (weighted by beta).

  • "log_likelihood_x1": Mean reconstruction term for x1.

  • "log_likelihood_x2": Mean reconstruction term for x2.

  • "kl_divergence_x1": Mean KL divergence for x1.

  • "kl_divergence_x2": Mean KL divergence for x2.

Return type:

LossResult

Raises:

ValueError – If either output of the pair fails the shape checks of VAE.compute_loss(), or if likelihood is not a supported likelihood model.

forward(x: tuple[Tensor, Tensor], S: int = 1) → AdaGVAEOutput

AdaGVAE forward pass on a pair of inputs.

• train() – adaptive grouping of the two posteriors, then sampling and decoding. Dimensions whose KL divergence \(\mathrm{KL}(q_1 \,\|\, q_2)\) is below the per-sample threshold \(\tau = \tfrac{1}{2}(\max_d \mathrm{KL}_d + \min_d \mathrm{KL}_d)\) share the average posterior (mean of the two means and of the two variances); the other dimensions keep their own posterior.

• eval() – no grouping: each input goes independently through VAE.forward(), which samples z from its own posterior.

For single-image inference after training use model.vae.encode and model.vae.decode.

Parameters:

• x (tuple[torch.Tensor, torch.Tensor]) – A (x1, x2) pair, each of shape [B, ...].

• S (int, optional) – Number of latent samples for Monte Carlo estimates. Must be >= 1. Defaults to 1.

Returns:

Pair outputs containing reconstructions and posterior parameters for both inputs (adapted posteriors in train mode).

Return type:

AdaGVAEOutput

Data Structures

class pyautoencoder.variational.AdaGVAEOutput(output1: VAEOutput, output2: VAEOutput)

Output of a full AdaGVAE paired forward pass.

output1

VAE output for the first input of the pair: built from the adapted posterior in training mode, plain VAE.forward() output in eval mode.

Type:

VAEOutput

output2

VAE output for the second input of the pair, as for output1.

Type:

VAEOutput