Parametric Spectral Autoencoder

spectral_autoencoder

CLASS	DESCRIPTION
SpectralAutoencoder	Spectral autoencoder combining eigenvalue optimization with reconstruction.

Classes

SpectralAutoencoder(*, model: dict[str, ModelType], kernel_fn: Kernel | str = 'gaussian', kernel_kwargs: dict | None = None, distance_fn: PairwiseDistance | str = 'euclidean', distance_kwargs: dict | None = None, loss_fn: Literal['eigenvalue', 'adaptive_eigenvalue'] = 'eigenvalue', loss_kwargs: dict | None = None, loss_weights: dict[str, float] | None = None, additional_loss: list | dict | None = None, additional_loss_weights: list[float] | dict[str, float] | 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, n_eigen: int = 10, symmetric_eigendecomposition: bool = True, device: Literal['cuda', 'cpu'] = 'cuda', store_loss_context: bool = False)

Bases: SpectralMap

Spectral autoencoder combining eigenvalue optimization with reconstruction.

Learns low-dimensional representations by jointly optimizing spectral properties of kernel matrices in the latent space and reconstruction quality. Inherits from SpectralMap and extends it with a decoder network for reconstruction, combining spectral gap maximization with the reconstruction objective of Autoencoder.

The model consists of an encoder network that maps input data to a latent space and a decoder network that reconstructs the original input. The total loss is a combination of spectral loss and reconstruction loss \(L_{\lambda} + L_{\mathrm{MSE}}\), where \(L_{\lambda}\) optimizes eigenvalue properties (e.g., spectral gap) and \(L_{\mathrm{MSE}}\) measures reconstruction error.

PARAMETER	DESCRIPTION
model	Dictionary containing 'encoder' and 'decoder' model specifications. Both networks are required for spectral autoencoder functionality: 'encoder': Maps input data to latent representation where spectral properties are optimized 'decoder': Maps latent representations back to original input space for reconstruction Example: model={"encoder": encoder_net, "decoder": decoder_net} TYPE: dict[str, ModelType]
kernel_fn	Kernel function specification for latent space. Kernel object: Pre-configured kernel instance str: Kernel type name (see available options below) Available string options: "gaussian": GaussianKernel "t": Student's TKernel TYPE: Kernel | str, optional, by default "gaussian" DEFAULT: 'gaussian'
kernel_kwargs	Keyword arguments for kernel instantiation when kernel_fn is a string. See kernels for available options. TYPE: dict | None, optional, by default None DEFAULT: None
distance_fn	Distance metric specification for latent space. PairwiseDistance object: Pre-configured distance instance str: Distance type name (see available options below) Available string options: "euclidean": Euclidean L2 distance "mahalanobis": Mahalanobis distance "covariance": Covariance-based distance TYPE: PairwiseDistance | str, optional, by default "euclidean" DEFAULT: 'euclidean'
distance_kwargs	Keyword arguments for distance instantiation when distance_fn is a string. TYPE: dict | None, optional, by default None DEFAULT: None
loss_fn	Loss function specification. Only EigenvalueLoss (loss_fn = "eigenvalue") and AdaptiveEigenvalueLoss (loss_fn = "adaptive_eigenvalue") are supported for spectral optimization. TYPE: str, by default "eigenvalue" DEFAULT: 'eigenvalue'
loss_kwargs	Additional keyword arguments for the spectral loss function. reduce: Literal["gap", "gap_exp", "sum", "sum_pow", "single", "none", None], optional, by default "gap" See losses for more information. sign: float, optional, by default -1.0 Loss sign multiplier for spectral loss. TYPE: dict | None, optional, by default None DEFAULT: None
loss_weights	Weights for built-in loss components. Keys must match names in DEFAULT_LOSS_WEIGHTS. Unspecified keys keep their default weight. For SpectralAutoencoder, the available keys are: "loss_eigenvalue": Weight for the spectral (eigenvalue) loss (default 1.0) "loss_reconstruction": Weight for the MSE reconstruction loss (default 1.0) Example: loss_weights={"loss_eigenvalue": 0.5, "loss_reconstruction": 2.0} TYPE: dict[str, float] | None, optional, by default None DEFAULT: None
additional_loss	Additional loss functions to combine with spectral and reconstruction losses. list[Loss]: Unnamed losses combined with primary losses dict[str, Loss]: Named losses for interpretable logging None: Use only spectral and MSE losses (default) The total loss becomes a weighted sum of spectral, MSE, and additional losses. Useful for adding regularization (e.g., OrthogonalLoss, DecorrelationLoss). TYPE: list[Loss] | dict[str, Loss] | None, optional, by default None DEFAULT: None
additional_loss_weights	Weights for additional loss functions. For list losses: provide list of floats (same length as additional_loss) For dict losses: provide dict with matching keys If None: all additional losses weighted at 1.0 TYPE: list[float] | dict[str, float] | None, 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
n_eigen	Number of eigenvalues to compute. TYPE: int, optional, by default 10 DEFAULT: 10
symmetric_eigendecomposition	Whether to use symmetric eigendecomposition for spectral loss. TYPE: bool, optional, by default False DEFAULT: True
device	Computation device. Falls back to "cpu" if CUDA unavailable. TYPE: Literal["cuda", "cpu"], optional, by default "cuda" DEFAULT: 'cuda'
store_loss_context	Whether to store the loss context dictionary for metric extraction in callbacks. When enabled, the loss_context attribute will contain the context dict passed to the loss function, enabling custom extractors to compute metrics from intermediate values like kernel matrices, embeddings, reconstructions, etc. TYPE: bool, optional, by default False DEFAULT: False

ATTRIBUTE	DESCRIPTION
DEFAULT_LOSS_WEIGHTS	Default weights for built-in loss components: {"loss_eigenvalue": 1.0, "loss_reconstruction": 1.0} TYPE: ClassVar[dict[str, float]]
decoder	Decoder network on the specified device. Also available as model_dict["decoder"] for consistency with multi-model interface. TYPE: Model
kernel_fn	Kernel function instance for latent space. TYPE: Kernel
distance_fn	Distance function instance for latent space. TYPE: PairwiseDistance
loss_fn	Composite loss function containing spectral, reconstruction, and any additional losses. TYPE: CompositeLoss
eigenvalues	Most recently computed eigenvalues from forward pass, shape (n_eigen,). Updated during loss() computation. TYPE: Tensor

Examples:

String-based configuration:

>>> import torch
>>> from spectre.parametric import SpectralAutoencoder
>>> from spectre.core import Model
>>>
>>> # Create encoder (input -> latent)
>>> encoder = Model(in_features=784, out_features=64, multipliers=[0.5])
>>> # Create decoder (latent -> output)
>>> decoder = Model(in_features=64, out_features=784, multipliers=[2.0])
>>>
>>> sae = SpectralAutoencoder(
...     model={"encoder": encoder, "decoder": decoder},
...     kernel_fn="gaussian",
...     kernel_kwargs={"bw_method": "median"},
...     distance_fn="euclidean",
...     loss_fn="eigenvalue",
...     loss_kwargs={"n_eigen": 2, "reduce": "gap"},
...     device="cpu",
... )
>>>
>>> X = torch.randn(100, 784)
>>> loss = sae.loss(X)  # Combined spectral + MSE loss

With custom loss weights:

>>> sae = SpectralAutoencoder(
...     model={"encoder": encoder, "decoder": decoder},
...     loss_kwargs={"n_eigen": 2, "reduce": "gap"},
...     loss_weights={"loss_eigenvalue": 0.5, "loss_reconstruction": 2.0},
...     device="cpu",
... )
>>>
>>> # Total 0.5 * spectral + 2.0 * MSE
>>> X = torch.randn(100, 784)
>>> loss = sae.loss(X)

With additional loss functions for regularization:

>>> from spectre.loss import OrthogonalLoss, DecorrelationLoss
>>>
>>> sae = SpectralAutoencoder(
...     model={"encoder": encoder, "decoder": decoder},
...     kernel_fn="gaussian",
...     loss_fn="eigenvalue",
...     loss_kwargs={"n_eigen": 2, "reduce": "gap"},
...     loss_weights={"loss_eigenvalue": 0.5},
...     additional_loss={
...         "orthogonal": OrthogonalLoss(),
...         "decorrelation": DecorrelationLoss(method="correlation"),
...     },
...     additional_loss_weights={
...         "orthogonal": 0.1,
...         "decorrelation": 0.05,
...     },
...     device="cpu",
... )
>>>
>>> # Total 0.5 * spectral + 1.0 * MSE + 0.1 * orthogonal + 0.05 * decorrelation
>>> X = torch.randn(100, 784)
>>> loss = sae.loss(X)  # Composite loss value

METHOD	DESCRIPTION
loss	Compute combined spectral and reconstruction losses.

Source code in spectre/parametric/spectral_autoencoder.py

def __init__(
    self,
    *,
    model: dict[str, ModelType],
    kernel_fn: Kernel | str = "gaussian",
    kernel_kwargs: dict | None = None,
    distance_fn: PairwiseDistance | str = "euclidean",
    distance_kwargs: dict | None = None,
    loss_fn: Literal["eigenvalue", "adaptive_eigenvalue"] = "eigenvalue",
    loss_kwargs: dict | None = None,
    loss_weights: dict[str, float] | None = None,
    additional_loss: list | dict | None = None,
    additional_loss_weights: list[float] | dict[str, float] | 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,
    n_eigen: int = 10,
    symmetric_eigendecomposition: bool = True,
    device: Literal["cuda", "cpu"] = "cuda",
    store_loss_context: bool = False,
) -> 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["encoder"],
        kernel_fn=kernel_fn,
        kernel_kwargs=kernel_kwargs,
        distance_fn=distance_fn,
        distance_kwargs=distance_kwargs,
        loss_fn=loss_fn,
        loss_kwargs=loss_kwargs,
        loss_weights=loss_weights,
        additional_loss=additional_loss,
        additional_loss_weights=additional_loss_weights,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
        n_eigen=n_eigen,
        symmetric_eigendecomposition=symmetric_eigendecomposition,
        device=device,
        store_loss_context=store_loss_context,
    )

    self.decoder = self.parse_model(model["decoder"], device)
    self.model_dict["decoder"] = self.decoder

    weight_reconstruction = self._loss_weights["loss_reconstruction"]
    self.loss_fn = self.loss_fn + CompositeLoss(
        losses={"loss_reconstruction": MSELoss()},
        weights={"loss_reconstruction": weight_reconstruction},
    )

Functions

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

Compute combined spectral and reconstruction losses.

Performs encoding and decoding operations, computes kernel matrix eigenvalues in the latent space for spectral loss, and measures reconstruction error between input and decoded output. When additional losses are provided, computes a composite loss combining spectral, MSE, and regularization terms.

PARAMETER	DESCRIPTION
X	Input data tensor to encode, evaluate spectrally, and reconstruct. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Optional sample weights for weighted kernel computation and weighted MSE loss. TYPE: torch.Tensor of shape (n_samples,) or None, by default None DEFAULT: None
target	Not used for loss computation (unsupervised method). TYPE: torch.Tensor or None, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar loss value. If additional losses are provided, returns weighted sum of all loss components (spectral + MSE + additional).

Source code in spectre/parametric/spectral_autoencoder.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute combined spectral and reconstruction losses.

    Performs encoding and decoding operations, computes kernel matrix
    eigenvalues in the latent space for spectral loss, and measures
    reconstruction error between input and decoded output. When additional
    losses are provided, computes a composite loss combining spectral, MSE,
    and regularization terms.

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

    weights : torch.Tensor of shape (n_samples,) or None, by default None
        Optional sample weights for weighted kernel computation and
        weighted MSE loss.

    target : torch.Tensor or None, by default None
        Not used for loss computation (unsupervised method).

    Returns
    -------
    torch.Tensor
        Scalar loss value. If additional losses are provided, returns
        weighted sum of all loss components (spectral + MSE + additional).
    """
    z_batch, z_kernel, eigenvalues, eigenvectors = self.compute_loss_context(
        X=X,
        weights=weights,
    )

    self.eigenvalues = eigenvalues
    self.eigenvectors = eigenvectors

    context = {
        "eigenvalues": eigenvalues,
        "eigenvectors": eigenvectors,
        "weights": weights,
        "K": z_kernel,
        "Z": z_batch,
        "Y": self.decoder(z_batch),
        "X": X,
    }

    if self.store_loss_context:
        self.loss_context = context

    loss = self.loss_fn(context)

    return loss