Architecture & Design#

Overview#

PyAutoencoder provides clean, well-documented implementations of the following fundamental autoencoder architectures:

  1. Autoencoder (AE) – Vanilla encoder-decoder pair

  2. Variational Autoencoder (VAE) – Probabilistic encoder with sampling

  3. Adaptive Group VAE (AdaGVAE) – Paired-input VAE for disentangled representations

All follow a consistent build/forward/compute_loss pattern and integrate naturally with PyTorch. AdaGVAE wraps a VAE and applies it to paired inputs.

Vanilla Autoencoder (AE)#

The AE is a deterministic encoder-decoder model:

Architecture

  • Encoder – Maps input \(x\) to latent code \(z = f(x)\)

  • Decoder – Maps latent \(z\) back to reconstruction \(\hat{x} = g(z)\), with the same shape as \(x\)

  • Forward pass – Returns both \(z\) and \(\hat{x}\)

Training

from pyautoencoder.vanilla import AE

model = AE(encoder=encoder_nn, decoder=decoder_nn)
model.build(sample_input)

# Training step
output = model(x_batch)          # AEOutput with z and x_hat
loss_result = model.compute_loss(x_batch, output, likelihood='bernoulli')
loss_result.objective.backward()

Inference

model.eval()                             # encode/decode keep the train/eval mode

# Extract latent codes (no gradients; usable as inputs of other models)
z = model.encode(x_batch).z              # AEEncodeOutput → latent tensor

# Reconstruct from latents
x_reconstructed = model.decode(z).x_hat  # AEDecodeOutput → reconstruction tensor

Output Structure

class AEOutput:
    x_hat: torch.Tensor       # Reconstruction [B, ...]
    z: torch.Tensor           # Latent code [B, ...] (the encoder output)

Variational Autoencoder (VAE)#

The VAE implements the VAE framework (Kingma & Welling, 2013) with probabilistic inference:

Architecture

  • Encoder – Maps input \(x\) to a flat feature vector of shape \([B, F]\) (end it with nn.Flatten() if needed)

  • Sampling layer – Maps the features to the mean \(\mu\) and log-variance \(\log\sigma^2\) of \(q(z|x)\), then draws \(z\) with the reparameterization trick

  • Decoder – Maps \(z\) to the parameters of \(p(x|z)\) (the reconstruction, or logits for the Bernoulli likelihood), with the shape of one input for each latent code

Training

from pyautoencoder.variational import VAE

model = VAE(encoder=encoder_nn, decoder=decoder_nn, latent_dim=64)
model.build(sample_input)

# Training step (sample multiple times for Monte Carlo estimates)
output = model(x_batch, S=5)  # 5 samples
loss_result = model.compute_loss(x_batch, output, beta=1.0, likelihood='gaussian')
loss_result.objective.backward()

Sampling vs Posterior Mean

Whether \(z\) is sampled is chosen with the sample argument, not with model.train() / model.eval() (which only affect layers such as dropout and batch normalization):

  • model(x, S=5) – samples \(S\) latent codes per input (sample=True is the default of forward, in train and eval mode alike). compute_loss on this output is a Monte Carlo estimate of the ELBO, also for validation data.

  • model.encode(x) – inference: no gradients and \(z = \mu\) by default (sample=False), with shape \([B, 1, D_z]\). It keeps the train/eval mode, so call model.eval() first. Use model.encode(x, S=10, sample=True) to draw samples instead.

  • model(x, sample=False) – \(z = \mu\) in the full forward pass. S must be 1 in this case.

# Validation ELBO (e.g. for early stopping)
model.eval()
with torch.no_grad():
    val_loss = model.compute_loss(x_val, model(x_val))

# Inference: posterior mean and deterministic reconstruction
enc = model.encode(x)
x_rec = model.decode(enc.z).x_hat

Output Structure

class VAEOutput:
    x_hat: torch.Tensor       # Reconstructions [B, S, ...]
    z: torch.Tensor           # Latent codes [B, S, D_z] (samples, or the mean with sample=False)
    mu: torch.Tensor          # Posterior mean [B, D_z]
    log_var: torch.Tensor     # Posterior log-variance [B, D_z]

Sampling New Data

Generate samples from the prior \(p(z) = \mathcal{N}(0, I)\):

model.eval()
with torch.no_grad():
    z_prior = torch.randn(n_samples, latent_dim)  # [N, D_z], no sample dimension
    x_samples = model.decoder(z_prior)            # logits for the Bernoulli likelihood: apply torch.sigmoid

Adaptive Group Variational Autoencoder (AdaGVAE)#

The AdaGVAE implements the Ada-GVAE framework (Locatello et al., 2020). It wraps a VAE and adds an adaptive posterior-alignment step during training to encourage disentanglement.

Architecture

  • Backbone – A fully configured VAE (encoder, sampling layer, decoder are reused)

  • Paired encoding – Both inputs \(x_1\) and \(x_2\) are encoded independently to obtain \(q_1(z|x_1)\) and \(q_2(z|x_2)\)

  • Adaptive alignment – Dimensions where the per-dimension KL divergence \(\mathrm{KL}(q_1 \| q_2)\) falls below a threshold \(\tau\) are shared (posterior averaged); the rest are kept independent

  • Decoder – The standard VAE decoder is applied to samples from the adapted posteriors

The threshold is computed per sample as the midpoint of the min and max per-dimension KL values:

\[\tau = \tfrac{1}{2}(\max_d \mathrm{KL}_d + \min_d \mathrm{KL}_d)\]

The alignment only affects training. In eval mode (model.eval()) no grouping is applied: each input goes through the wrapped VAE’s forward independently.

Training

from pyautoencoder.variational import VAE, AdaGVAE

vae = VAE(encoder=encoder_nn, decoder=decoder_nn, latent_dim=64)
model = AdaGVAE(vae=vae)
model.build(sample_input)  # a batch of single inputs, not a pair

# Training step — forward takes a pair
output = model((x1_batch, x2_batch), S=5)
loss_result = model.compute_loss((x1_batch, x2_batch), output, beta=4.0, likelihood='bernoulli')
loss_result.objective.backward()

Inference

After training, AdaGVAE reuses the underlying VAE for single-image inference:

# Encode / decode through the wrapped VAE as usual
model.eval()
z = model.vae.encode(x).z
x_reconstructed = model.vae.decode(z).x_hat

Output Structure

class AdaGVAEOutput:
    output1: VAEOutput    # Output for x1
    output2: VAEOutput    # Output for x2

Each VAEOutput contains x_hat, z, mu, and log_var for the corresponding posterior: the adapted one in train mode, the plain VAE posterior in eval mode.

Loss Functions#

The loss computation is integrated into each model via its compute_loss method. This approach keeps loss logic close to the model implementation and ensures consistency.

Reconstruction Loss (AE)#

For standard autoencoders, use AE.compute_loss to compute reconstruction loss:

output = model(x)
loss_result = model.compute_loss(x, output, likelihood='gaussian')
loss_result.objective.backward()

Supported likelihoods (per element; the loss sums over features and averages over the batch):

  • Gaussian – Continuous data (unit variance, no normalization constant)

    \[\text{NLL} = \frac{1}{2}(x-\hat{x})^2\]

  • Bernoulli – Binary data or values in \([0, 1]\) (x_hat are logits)

    \[\text{NLL} = \text{BCE}_{\text{logits}}(x, \hat{x})\]

ELBO Loss (VAE)#

For variational autoencoders, use VAE.compute_loss to compute the negative (beta-weighted) ELBO:

output = model(x, S=5)
loss_result = model.compute_loss(x, output, beta=1.0, likelihood='gaussian')
loss_result.objective.backward()

The ELBO decomposes into reconstruction and regularization terms:

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

Beta-VAE Weighting

The \(\beta\) hyperparameter controls the KL weight:

  • \(\beta = 1.0\) – Standard VAE (matches true ELBO)

  • \(\beta > 1.0\) – Stronger regularization (more disentangled latents)

  • \(\beta < 1.0\) – Weaker regularization (better reconstruction)

The elbo diagnostic is \(\mathcal{L}(x; \beta)\), so it equals the ELBO only for \(\beta = 1\); kl_divergence is not multiplied by \(\beta\).

Paired ELBO Loss (AdaGVAE)#

For AdaGVAE, use AdaGVAE.compute_loss to compute the combined objective over the pair:

output = model((x1, x2), S=5)
loss_result = model.compute_loss((x1, x2), output, beta=4.0, likelihood='bernoulli')
loss_result.objective.backward()

The objective is the negative of the sum of two beta-VAE ELBOs, evaluated on the adapted posteriors \(q(\hat{z}|x_1)\) and \(q(\hat{z}|x_2)\) in train mode (on the plain VAE posteriors in eval mode):

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

The diagnostics dictionary exposes per-input reconstruction and KL terms (log_likelihood_x1, log_likelihood_x2, kl_divergence_x1, kl_divergence_x2) so both views can be monitored independently.

Loss Result Structure

AE.compute_loss, VAE.compute_loss, and AdaGVAE.compute_loss all return a LossResult with:

  • objective – scalar loss to optimize (backward-differentiable)

  • diagnostics – dictionary of scalar metrics, as detached 0-dim tensors (call .item() when logging; this avoids a GPU sync at every step)

Example:

loss_result = model.compute_loss(x, output)
loss_result.objective.backward()              # Optimize this
for name, val in loss_result.diagnostics.items():
    log(name, val.item())                     # Monitor these

Key Design Principles#

User-Provided Networks

You provide the encoder and decoder as arbitrary PyTorch modules:

encoder = my_custom_encoder_network()  # Any nn.Module
decoder = my_custom_decoder_network()  # Any nn.Module
model = AE(encoder, decoder)

This keeps the library flexible and composable with existing PyTorch code.

Explicit Initialization

The build method initializes size-dependent parameters:

model.build(sample_input)  # Required once before training

This catches shape mismatches early and makes model behavior transparent.

Clear Loss Integration

Loss functions are methods on the models themselves, providing a clean API:

output = model(x_batch)
loss_result = model.compute_loss(x_batch, output)
loss_result.objective.backward()              # Optimize this
for name, val in loss_result.diagnostics.items():
    log(name, val.item())                     # Monitor these

Consistent Interfaces

All the architectures follow the same pattern:

model.build(sample)                          # Initialize
output = model(x_batch)                      # Forward pass
loss_result = model.compute_loss(x, output)  # Compute loss
loss_result.objective.backward()             # Backprop