Loss Elbo Gmm

elbo_gmm

CLASS	DESCRIPTION
ELBOGMMLoss	ELBO loss for Variational Autoencoder with Gaussian Mixture Model prior and

FUNCTION	DESCRIPTION
elbo_gmm_loss	Compute ELBO loss for GMM-VAE with categorical latent variables (functional API).

Classes

ELBOGMMLoss(n_mc_samples: int = 1, reduce: Literal['sum', 'mean'] = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps, context_prefix: str | None = None, **kwargs)

Bases: Loss

ELBO loss for Variational Autoencoder with Gaussian Mixture Model prior and categorical latent variables.

This loss function combines reconstruction loss with KL divergence terms for both continuous and categorical latent variables, enabling clustering in the latent space via a Gaussian mixture prior.

The loss is defined as:

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

where:

• \(p(x|z)\) is the reconstruction likelihood (Gaussian)
• \(q(z|x,c)\) is the posterior over continuous latent given component
• \(p(z|c)\) is the GMM component-conditional prior
• \(q(c|x)\) is the posterior over component assignments
• \(p(c)\) is the categorical prior (mixture weights)

The KL divergence for continuous latent \(z\) is computed via Monte Carlo estimation as there is no closed form for \(D_{KL}(N(\mu, \sigma^2) \| GMM)\).

PARAMETER	DESCRIPTION
n_mc_samples	Number of Monte Carlo samples for KL divergence estimation. More samples provide better gradient estimates but increase computation. TYPE: int, by default 1 DEFAULT: 1
reduce	Reduction operation for loss aggregation. "sum": Sum over batch dimension "mean": Mean over batch dimension TYPE: Literal["sum", "mean"], by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, by default 1.0 DEFAULT: 1.0
eps	Regularization term for numerical stability in log computations. TYPE: float, by default torch.finfo(torch.float32).eps DEFAULT: eps
**kwargs	Additional keyword arguments. DEFAULT: {}

Examples:

>>> import torch
>>> from spectre.loss import ELBOGMMCategoricalLoss
>>>
>>> loss_fn = ELBOGMMCategoricalLoss(n_mc_samples=1, reduce="mean")
>>>
>>> # Build context dictionary with required keys
>>> context = {
...     "X": torch.randn(32, 784),  # Input data
...     "Y": torch.randn(32, 784),  # Reconstructed data
...     "mean": torch.randn(32, 2),  # Encoder mean
...     "logvar": torch.randn(32, 2),  # Encoder log-variance
...     "z_sample": torch.randn(32, 2),  # Latent sample
...     "categorical_logits": torch.randn(32, 3),  # Component logits
...     "gmm_means": torch.randn(3, 2),  # GMM component means
...     "gmm_covariances": torch.randn(3, 2),  # GMM covariances (diag)
...     "gmm_weights": torch.softmax(torch.randn(3), dim=0),  # Mixing weights
...     "covariance_type": "diag",
... }
>>>
>>> loss = loss_fn(context)
>>> loss.shape
torch.Size([])

Notes

• Supports all covariance types: "diag", "spherical", "full", "tied"
• KL divergence for categorical is computed in closed form
• KL divergence for continuous latent uses Monte Carlo estimation
• Numerical stability via log-sum-exp trick for GMM probabilities

See Also

ELBOGaussianLoss: Standard VAE ELBO loss

METHOD	DESCRIPTION
forward	Compute ELBO loss with GMM prior and categorical latent variables.

Source code in spectre/loss/elbo_gmm.py

def __init__(
    self,
    n_mc_samples: int = 1,
    reduce: Literal["sum", "mean"] = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
    context_prefix: str | None = None,
    **kwargs,
):
    super().__init__(
        reduce=reduce, sign=sign, context_prefix=context_prefix, **kwargs
    )

    if not isinstance(n_mc_samples, int):
        raise ValueError(f"`n_mc_samples` must be an integer, got {n_mc_samples}")
    check_in_interval(n_mc_samples, "[1, inf)")
    self.n_mc_samples = n_mc_samples

    if not isinstance(eps, float):
        raise ValueError(f"`eps` must be a float, got {eps}")
    check_in_interval(eps, "[0, 1)")
    self.eps = eps

Functions

forward(context: dict, **kwargs) -> torch.Tensor

Compute ELBO loss with GMM prior and categorical latent variables.

PARAMETER	DESCRIPTION
context	Dictionary containing required keys: X : torch.Tensor of shape (n_samples, in_features) - Input data Y : torch.Tensor of shape (n_samples, in_features) - Reconstructed data mean : torch.Tensor of shape (n_samples, z_dim) - Encoder mean logvar : torch.Tensor of shape (n_samples, z_dim) - Encoder log-variance z_sample : torch.Tensor of shape (n_samples, z_dim) - Latent samples categorical_logits : torch.Tensor of shape (n_samples, n_components) - Component logits from encoder gmm_means : torch.Tensor of shape (n_components, z_dim) - GMM component means gmm_covariances : torch.Tensor - GMM covariances (shape depends on type) gmm_weights : torch.Tensor of shape (n_components,) - Mixing weights covariance_type : str - Type of covariance ("diag", "spherical", "full", "tied") weights : torch.Tensor of shape (n_samples,), optional - Sample weights TYPE: dict
kwargs	Additional keyword arguments. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Scalar ELBO loss value.

Source code in spectre/loss/elbo_gmm.py

def forward(self, context: dict, **kwargs) -> torch.Tensor:
    """
    Compute ELBO loss with GMM prior and categorical latent variables.

    Parameters
    ----------
    context : dict
        Dictionary containing required keys:

        - X : torch.Tensor of shape (n_samples, in_features) - Input data
        - Y : torch.Tensor of shape (n_samples, in_features) - Reconstructed data
        - mean : torch.Tensor of shape (n_samples, z_dim) - Encoder mean
        - logvar : torch.Tensor of shape (n_samples, z_dim) - Encoder log-variance
        - z_sample : torch.Tensor of shape (n_samples, z_dim) - Latent samples
        - categorical_logits : torch.Tensor of shape (n_samples, n_components) -
          Component logits from encoder
        - gmm_means : torch.Tensor of shape (n_components, z_dim) -
          GMM component means
        - gmm_covariances : torch.Tensor - GMM covariances (shape depends on type)
        - gmm_weights : torch.Tensor of shape (n_components,) - Mixing weights
        - covariance_type : str - Type of covariance ("diag", "spherical",
          "full", "tied")
        - weights : torch.Tensor of shape (n_samples,), optional - Sample weights

    kwargs : dict
        Additional keyword arguments.

    Returns
    -------
    torch.Tensor
        Scalar ELBO loss value.
    """
    X = self._get_context(context, "X", None)
    Y = self._get_context(context, "Y", None)
    mean = self._get_context(context, "mean", None)
    logvar = self._get_context(context, "logvar", None)
    z_sample = self._get_context(context, "z_sample", None)
    categorical_logits = self._get_context(context, "categorical_logits", None)
    gmm_means = self._get_context(context, "gmm_means", None)
    gmm_covariances = self._get_context(context, "gmm_covariances", None)
    gmm_weights = self._get_context(context, "gmm_weights", None)
    covariance_type = self._get_context(context, "covariance_type", None)
    weights = self._get_context(context, "weights", None)

    # Validate required fields
    required_fields = [
        X,
        Y,
        mean,
        logvar,
        z_sample,
        gmm_means,
        gmm_covariances,
        gmm_weights,
        covariance_type,
    ]

    if any(x is None for x in required_fields):
        raise ValueError(
            "ELBOGMMCategoricalLoss requires X, Y, mean, logvar, z_sample, "
            "gmm_means, gmm_covariances, gmm_weights, and covariance_type in context."
        )

    # If no categorical_logits, use uniform distribution over components
    if categorical_logits is None:
        categorical_logits = torch.zeros(
            X.shape[0], gmm_weights.shape[0], device=X.device
        )

    if weights is None:
        weights = torch.ones(X.shape[0], device=X.device, dtype=X.dtype)
    weights = weights.squeeze()

    return _elbo_gmm_loss_impl(
        X=X,
        Y=Y,
        mean=mean,
        logvar=logvar,
        z_sample=z_sample,
        categorical_logits=categorical_logits,
        gmm_means=gmm_means,
        gmm_covariances=gmm_covariances,
        gmm_weights=gmm_weights,
        covariance_type=covariance_type,
        weights=weights,
        n_mc_samples=self.n_mc_samples,
        reduce=self.reduce,
        sign=self.sign,
        eps=self.eps,
    )

Functions

elbo_gmm_loss(X: torch.Tensor, Y: torch.Tensor, mean: torch.Tensor, logvar: torch.Tensor, z_sample: torch.Tensor, categorical_logits: torch.Tensor, gmm_means: torch.Tensor, gmm_covariances: torch.Tensor, gmm_weights: torch.Tensor, covariance_type: Literal['diag', 'spherical', 'full', 'tied'], weights: torch.Tensor | None = None, n_mc_samples: int = 1, reduce: str = 'mean', sign: float = 1.0, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute ELBO loss for GMM-VAE with categorical latent variables (functional API).

The loss combines reconstruction error with two KL divergence terms:

1. KL divergence for continuous latent: \(D_{KL}(q(z|x,c) \| p(z|c))\)
2. KL divergence for categorical: \(D_{KL}(q(c|x) \| p(c))\)

PARAMETER	DESCRIPTION
X	Input data of shape (n_samples, in_features). TYPE: Tensor
Y	Reconstructed data of shape (n_samples, in_features). TYPE: Tensor
mean	Encoder mean of shape (n_samples, z_dim). TYPE: Tensor
logvar	Encoder log-variance of shape (n_samples, z_dim). TYPE: Tensor
z_sample	Latent samples of shape (n_samples, z_dim). TYPE: Tensor
categorical_logits	Component logits of shape (n_samples, n_components). TYPE: Tensor
gmm_means	GMM component means of shape (n_components, z_dim). TYPE: Tensor
gmm_covariances	GMM covariances (shape depends on covariance_type). TYPE: Tensor
gmm_weights	GMM mixing weights of shape (n_components,). TYPE: Tensor
covariance_type	Type of covariance parameterization. TYPE: Literal['diag', 'spherical', 'full', 'tied']
weights	Sample weights of shape (n_samples,). TYPE: torch.Tensor | None, by default None DEFAULT: None
n_mc_samples	Number of Monte Carlo samples for KL estimation. TYPE: int, by default 1 DEFAULT: 1
reduce	Reduction operation: "sum" or "mean". TYPE: str, by default "mean" DEFAULT: 'mean'
sign	Loss sign multiplier. TYPE: float, by default 1.0 DEFAULT: 1.0
eps	Regularization for numerical stability. TYPE: float, by default 1e-7 DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Scalar ELBO loss value.

Source code in spectre/loss/elbo_gmm.py

def elbo_gmm_loss(
    X: torch.Tensor,
    Y: torch.Tensor,
    mean: torch.Tensor,
    logvar: torch.Tensor,
    z_sample: torch.Tensor,
    categorical_logits: torch.Tensor,
    gmm_means: torch.Tensor,
    gmm_covariances: torch.Tensor,
    gmm_weights: torch.Tensor,
    covariance_type: Literal["diag", "spherical", "full", "tied"],
    weights: torch.Tensor | None = None,
    n_mc_samples: int = 1,
    reduce: str = "mean",
    sign: float = 1.0,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute ELBO loss for GMM-VAE with categorical latent variables (functional API).

    The loss combines reconstruction error with two KL divergence terms:

    1. KL divergence for continuous latent: $D_{KL}(q(z|x,c) \\| p(z|c))$
    2. KL divergence for categorical: $D_{KL}(q(c|x) \\| p(c))$

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

    Y : torch.Tensor
        Reconstructed data of shape (n_samples, in_features).

    mean : torch.Tensor
        Encoder mean of shape (n_samples, z_dim).

    logvar : torch.Tensor
        Encoder log-variance of shape (n_samples, z_dim).

    z_sample : torch.Tensor
        Latent samples of shape (n_samples, z_dim).

    categorical_logits : torch.Tensor
        Component logits of shape (n_samples, n_components).

    gmm_means : torch.Tensor
        GMM component means of shape (n_components, z_dim).

    gmm_covariances : torch.Tensor
        GMM covariances (shape depends on covariance_type).

    gmm_weights : torch.Tensor
        GMM mixing weights of shape (n_components,).

    covariance_type : Literal["diag", "spherical", "full", "tied"]
        Type of covariance parameterization.

    weights : torch.Tensor | None, by default None
        Sample weights of shape (n_samples,).

    n_mc_samples : int, by default 1
        Number of Monte Carlo samples for KL estimation.

    reduce : str, by default "mean"
        Reduction operation: "sum" or "mean".

    sign : float, by default 1.0
        Loss sign multiplier.

    eps : float, by default 1e-7
        Regularization for numerical stability.

    Returns
    -------
    torch.Tensor
        Scalar ELBO loss value.
    """
    check_same_shape(X, Y)
    check_same_shape(mean, logvar)
    check_same_shape(mean, z_sample)

    if weights is not None:
        if weights.ndim > 1:
            weights = weights.squeeze()
        check_1d(weights)
        check_same_len(X, weights)
        check_sample_weights(weights)
    else:
        weights = torch.ones(X.shape[0], device=X.device, dtype=X.dtype)

    if reduce not in ["sum", "mean"]:
        raise ValueError("Unknown reduce operation. Options: 'sum', 'mean'.")

    if covariance_type not in ["diag", "spherical", "full", "tied"]:
        raise ValueError(
            f"Unknown covariance_type: {covariance_type}. "
            "Options: 'diag', 'spherical', 'full', 'tied'."
        )

    return _elbo_gmm_loss_impl(
        X,
        Y,
        mean,
        logvar,
        z_sample,
        categorical_logits,
        gmm_means,
        gmm_covariances,
        gmm_weights,
        covariance_type,
        weights,
        n_mc_samples,
        reduce,
        sign,
        eps,
    )