Parametric Variational Autoencoder

variational_autoencoder

CLASS	DESCRIPTION
VariationalAutoencoder	Variational Autoencoder (VAE) with Gaussian latent space.

Classes

VariationalAutoencoder(*, model: dict[str, ModelType], z_dim: int = 2, loss_fn: Literal['elbo', 'elbo_log_cosh'] = 'elbo', loss_kwargs: dict | None = None, optimizer_fn: type[torch.optim.Optimizer] | str | None = None, optimizer_kwargs: dict | None = None, lr_scheduler_fn: Any | str | None = None, lr_scheduler_kwargs: dict | None = None, device: Literal['cuda', 'cpu'] = 'cuda')

Bases: Parametric

Variational Autoencoder (VAE) with Gaussian latent space.

This class implements a Variational Autoencoder that learns to encode input data into a probabilistic latent space and decode it back to the original space. The model uses the reparameterization trick to enable gradient-based optimization of the evidence lower bound (ELBO).

The VAE architecture consists of:

1. Encoder network that maps inputs to latent distribution parameters
2. Latent space parameterization (mean and log-variance)
3. Stochastic sampling using the reparameterization trick
4. Decoder network that reconstructs inputs from latent samples

The model optimizes the ELBO, which combines reconstruction loss with KL divergence regularization to encourage the latent space to match a standard Gaussian prior.

The negative ELBO is given by

\(-\mathbb{E}_{q(z|x)}[\log p(x|z)] + D_{KL}(q(z|x) || p(z))\),

where \(q(z|x)\) is the approximate posterior, \(p(x|z)\) is the likelihood, and \(p(z)\) is the prior (standard normal).

Another option is to use a log-cosh reconstruction loss instead of Gaussian likelihood, which can be more robust to outliers.

The log-cosh loss is defined as

\(-\mathbb{E}_{q(z|x)}[\log \cosh(p(x|z))] + D_{KL}(q(z|x) || p(z))\).

PARAMETER	DESCRIPTION
model	Dictionary containing 'encoder' and 'decoder' model specifications. 'encoder': Maps input to latent distribution parameters (must have output activation) 'decoder': Maps latent samples back to reconstructed input space Example: model={"encoder": encoder_net, "decoder": decoder_net} Note: Encoder must have an output activation function. TYPE: dict[str, ModelType]
z_dim	Dimensionality of the latent space. TYPE: int, by default 2 DEFAULT: 2
loss_fn	Loss function type. "elbo": Standard ELBO with Gaussian reconstruction likelihood "elbo_log_cosh": ELBO with log-cosh reconstruction loss TYPE: Literal["elbo", "elbo_log_cosh"], by default "elbo" DEFAULT: 'elbo'
loss_kwargs	Additional keyword arguments to pass to the loss function constructor. See ELBOGaussianLoss and ELBOLogCoshLoss for available options. TYPE: dict, optional, by default None DEFAULT: None
optimizer_fn	Optimizer specification for training. None: Uses Adam optimizer (default) str: Optimizer name from torch.optim (e.g., "AdamW", "SGD") type[torch.optim.Optimizer]: Optimizer class (supports third-party) TYPE: type[torch.optim.Optimizer] | str | None, optional, by default None DEFAULT: None
optimizer_kwargs	Keyword arguments for optimizer instantiation (e.g., {"lr": 1e-3, "weight_decay": 1e-4}). TYPE: dict | None, optional, by default None DEFAULT: None
lr_scheduler_fn	Learning rate scheduler specification. None: No scheduler (default) str: Scheduler name from torch.optim.lr_scheduler (e.g., "ReduceLROnPlateau") type: Scheduler class TYPE: Any | str | None, optional, by default None DEFAULT: None
lr_scheduler_kwargs	Keyword arguments for scheduler instantiation. TYPE: dict | None, optional, by default None DEFAULT: None
device	Computation device. Automatically falls back to "cpu" if CUDA unavailable. TYPE: Literal["cuda", "cpu"], by default "cuda" DEFAULT: 'cuda'

Examples:

>>> import torch
>>> from spectre.core import Model
>>> from spectre.parametric import VariationalAutoencoder
>>>
>>> # Create encoder and decoder networks
>>> encoder = Model(in_features=784, out_features=128, has_output_fn=True)
>>> decoder = Model(in_features=2, out_features=784)
>>>
>>> # Initialize VAE with model dictionary
>>> vae = VariationalAutoencoder(
...     model={"encoder": encoder, "decoder": decoder},
...     z_dim=2,
...     loss_fn="elbo",
... )
>>>
>>> # Training data
>>> X = torch.randn(100, 784)
>>> weights = torch.ones(100)
>>> batch = (X, weights)
>>>
>>> # Training step
>>> loss = vae.training_step(batch, 0)
>>>
>>> # Generate samples (after training)
>>> vae.is_trained = True
>>> z_sample = torch.randn(10, 2)  # Sample from prior
>>> reconstructions = vae.decoder(z_sample)

Notes

• The encoder model must have an output activation function
• Latent space is regularized to match standard Gaussian prior
• The score_step method provides VAE-specific scoring (negative ELBO)

See Also

Autoencoder: Standard deterministic autoencoder

METHOD	DESCRIPTION
forward_step	Forward pass through the encoder network.
loss	Compute VAE loss (ELBO or log-cosh).
score_step	Compute VAE reconstruction score.
training_step	Training step for PyTorch Lightning with VAE loss.
encode	Encode input data to latent distribution parameters.
decode	Decode latent codes back to input space.
sample	Generate samples from the prior latent distribution.
reconstruct	Reconstruct input data through the VAE.

Source code in spectre/parametric/variational_autoencoder.py

def __init__(
    self,
    *,
    model: dict[str, ModelType],
    z_dim: int = 2,
    loss_fn: Literal["elbo", "elbo_log_cosh"] = "elbo",
    loss_kwargs: dict | None = None,
    optimizer_fn: type[torch.optim.Optimizer] | str | None = None,
    optimizer_kwargs: dict | None = None,
    lr_scheduler_fn: Any | str | None = None,
    lr_scheduler_kwargs: dict | None = None,
    device: Literal["cuda", "cpu"] = "cuda",
) -> None:
    if not isinstance(model, dict):
        raise TypeError(
            "model must be a dictionary containing 'encoder' and 'decoder' keys."
        )
    if "encoder" not in model or "decoder" not in model:
        raise ValueError(
            "model dict must contain both 'encoder' and 'decoder' keys."
        )

    super().__init__(
        model=model,
        device=device,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

    self.decoder = self.model_dict["decoder"]

    if not self.model.has_output_fn:
        raise ValueError("For VAE, encoder model must have an output function.")

    if not isinstance(z_dim, int):
        raise TypeError(f"z_dim must be an integer, got {type(z_dim)}")
    check_in_interval(z_dim, "(0, inf)")
    self.z_dim = z_dim
    self.z_loc_nn = torch.nn.Linear(
        in_features=self.model.out_features, out_features=z_dim
    )
    self.z_scale_nn = torch.nn.Linear(
        in_features=self.model.out_features, out_features=z_dim
    )

    self.z_loc_nn = self.z_loc_nn.to(device=device)
    self.z_scale_nn = self.z_scale_nn.to(device=device)

    self.decoder = self.decoder.to(device=device)

    if loss_kwargs is None:
        loss_kwargs = {}
    if not isinstance(loss_kwargs, dict):
        raise TypeError(
            "Additional loss parameters `loss_kwargs` must be a dictionary."
        )

    if loss_fn == "elbo":
        self.loss_fn = ELBOGaussianLoss(**loss_kwargs)
    elif loss_fn == "elbo_log_cosh":
        self.loss_fn = ELBOLogCoshLoss(**loss_kwargs)
    else:
        raise ValueError(
            f"Expected `loss_fn` to be 'elbo' or 'elbo_log_cosh', got '{loss_fn}'."
        )

Functions

forward_step(X: torch.Tensor) -> torch.Tensor

Forward pass through the encoder network.

PARAMETER	DESCRIPTION
X	Input data. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, encoder_out_features)	Encoded representations for latent parameter computation.

Source code in spectre/parametric/variational_autoencoder.py

def forward_step(self, X: torch.Tensor) -> torch.Tensor:
    """
    Forward pass through the encoder network.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data.

    Returns
    -------
    torch.Tensor of shape (n_samples, encoder_out_features)
        Encoded representations for latent parameter computation.
    """
    encoded = self.model(X)
    return encoded

loss(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute VAE loss (ELBO or log-cosh).

Computes the variational autoencoder loss combining reconstruction loss with KL divergence regularization of the latent space.

PARAMETER	DESCRIPTION
X	Input data to encode and reconstruct. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Ignored for VAE (unsupervised learning). TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar VAE loss (ELBO or log-cosh variant).

Source code in spectre/parametric/variational_autoencoder.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute VAE loss (ELBO or log-cosh).

    Computes the variational autoencoder loss combining reconstruction
    loss with KL divergence regularization of the latent space.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to encode and reconstruct.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted loss computation.

    target : torch.Tensor or None, by default None
        Ignored for VAE (unsupervised learning).

    Returns
    -------
    torch.Tensor
        Scalar VAE loss (ELBO or log-cosh variant).
    """
    z_batch = self.forward_step(X)
    z_loc = self.z_loc_nn(z_batch)
    z_logvar = self.z_scale_nn(z_batch)
    z_scale = torch.exp(0.5 * z_logvar)
    z_sample = torch.distributions.Normal(z_loc, z_scale).rsample()
    z_target = self.decoder(z_sample)

    context = {
        "X": X,
        "Y": z_target,
        "mean": z_loc,
        "logvar": z_logvar,
    }
    if weights is not None:
        context["weights"] = weights

    loss = self.loss_fn(context)

    return loss

score_step(X: torch.Tensor, weights: torch.Tensor | None = None, target: torch.Tensor | None = None) -> torch.Tensor

Compute VAE reconstruction score.

For VAE, the natural scoring metric is the negative ELBO, which measures how well the model can reconstruct the input data while maintaining a regularized latent space.

PARAMETER	DESCRIPTION
X	Input data to score. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted scoring. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Ignored for VAE (unsupervised learning). TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	VAE score (negative ELBO - lower is better).

Source code in spectre/parametric/variational_autoencoder.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute VAE reconstruction score.

    For VAE, the natural scoring metric is the negative ELBO, which
    measures how well the model can reconstruct the input data while
    maintaining a regularized latent space.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to score.

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Sample weights for weighted scoring.

    target : torch.Tensor or None, by default None
        Ignored for VAE (unsupervised learning).

    Returns
    -------
    torch.Tensor
        VAE score (negative ELBO - lower is better).
    """
    score = self.loss(X, target=None, weights=weights)
    return score

training_step(batch: DataBatch, batch_idx: int) -> torch.Tensor

Training step for PyTorch Lightning with VAE loss.

Computes the VAE loss (ELBO or log-cosh) and logs it for monitoring during training and validation.

PARAMETER	DESCRIPTION
batch	Input batch containing (X, weights) where: X : torch.Tensor of shape (n_samples, in_features) Input data weights : torch.Tensor of shape (n_samples,) Sample weights TYPE: Batch
batch_idx	Batch index. TYPE: int

RETURNS	DESCRIPTION
Tensor	Scalar training loss (VAE ELBO or log-cosh loss).

Source code in spectre/parametric/variational_autoencoder.py

def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    Training step for PyTorch Lightning with VAE loss.

    Computes the VAE loss (ELBO or log-cosh) and logs it for monitoring
    during training and validation.

    Parameters
    ----------
    batch : Batch
        Input batch containing (X, weights) where:

        - X : torch.Tensor of shape (n_samples, in_features)
            Input data
        - weights : torch.Tensor of shape (n_samples,)
            Sample weights

    batch_idx : int
        Batch index.

    Returns
    -------
    torch.Tensor
        Scalar training loss (VAE ELBO or log-cosh loss).
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss)

    return loss

encode(X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

Encode input data to latent distribution parameters.

PARAMETER	DESCRIPTION
X	Input data to encode. TYPE: torch.Tensor of shape (n_samples, in_features)

RETURNS	DESCRIPTION
tuple[Tensor, Tensor]	Tuple of (z_loc, z_scale) where: z_loc: Mean of latent distribution, shape (n_samples, z_dim) z_scale: Scale of latent distribution, shape (n_samples, z_dim)

Source code in spectre/parametric/variational_autoencoder.py

def encode(self, X: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Encode input data to latent distribution parameters.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to encode.

    Returns
    -------
    tuple[torch.Tensor, torch.Tensor]
        Tuple of (`z_loc`, `z_scale`) where:

        - `z_loc`: Mean of latent distribution, shape (n_samples, z_dim)
        - `z_scale`: Scale of latent distribution, shape (n_samples, z_dim)
    """
    encoded = self.forward_step(X)
    z_loc = self.z_loc_nn(encoded)
    z_scale = torch.exp(0.5 * self.z_scale_nn(encoded))

    return z_loc, z_scale

decode(Z: torch.Tensor) -> torch.Tensor

Decode latent codes back to input space.

PARAMETER	DESCRIPTION
Z	Latent codes to decode. TYPE: torch.Tensor of shape (n_samples, z_dim)

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, decoder_out_features)	Reconstructed data.

Source code in spectre/parametric/variational_autoencoder.py

def decode(self, Z: torch.Tensor) -> torch.Tensor:
    """
    Decode latent codes back to input space.

    Parameters
    ----------
    Z : torch.Tensor of shape (n_samples, z_dim)
        Latent codes to decode.

    Returns
    -------
    torch.Tensor of shape (n_samples, decoder_out_features)
        Reconstructed data.
    """
    return self.decoder(Z)

sample(n_samples: int = 1) -> torch.Tensor

Generate samples from the prior latent distribution.

PARAMETER	DESCRIPTION
n_samples	Number of samples to generate. TYPE: int, by default 1 DEFAULT: 1

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, decoder_out_features)	Generated samples.

Source code in spectre/parametric/variational_autoencoder.py

def sample(self, n_samples: int = 1) -> torch.Tensor:
    """
    Generate samples from the prior latent distribution.

    Parameters
    ----------
    n_samples : int, by default 1
        Number of samples to generate.

    Returns
    -------
    torch.Tensor of shape (n_samples, decoder_out_features)
        Generated samples.
    """
    device = next(self.decoder.parameters()).device
    z_prior = torch.randn(n_samples, self.z_dim, device=device)

    return self.decode(z_prior)

reconstruct(X: torch.Tensor, use_mean: bool = False) -> torch.Tensor

Reconstruct input data through the VAE.

PARAMETER	DESCRIPTION
X	Input data to reconstruct. TYPE: torch.Tensor of shape (n_samples, in_features)
use_mean	If True, use the mean of the latent distribution instead of sampling. If False, use stochastic sampling (reparameterization trick). TYPE: bool, by default False DEFAULT: False

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, decoder_out_features)	Reconstructed data.

Source code in spectre/parametric/variational_autoencoder.py

def reconstruct(self, X: torch.Tensor, use_mean: bool = False) -> torch.Tensor:
    """
    Reconstruct input data through the VAE.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, in_features)
        Input data to reconstruct.

    use_mean : bool, by default False
        If True, use the mean of the latent distribution instead of sampling.
        If False, use stochastic sampling (reparameterization trick).

    Returns
    -------
    torch.Tensor of shape (n_samples, decoder_out_features)
        Reconstructed data.
    """
    z_loc, z_scale = self.encode(X)
    z_batch = (
        z_loc if use_mean else torch.distributions.Normal(z_loc, z_scale).rsample()
    )

    return self.decode(z_batch)

Functions