Parametric Spectral Net

spectral_net

CLASS	DESCRIPTION
SpectralNet	SpectralNet for spectral clustering via deep neural networks.

Classes

SpectralNet(*, model: ModelType | 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: str = 'laplacian', 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, device: Literal['cuda', 'cpu'] = 'cuda', store_loss_context: bool = False)

Bases: Parametric

SpectralNet for spectral clustering via deep neural networks.

This is a simplified version of SpectraNet. See [1] for the original method.

Learns embeddings that approximate the eigenvectors of the graph Laplacian by optimizing a combination of orthogonality constraints and Rayleigh quotient minimization. This enables learning of spectral clustering without explicit eigendecomposition.

The method trains a neural network to map input data to low-dimensional embeddings \(Y\) such that:

1. \(\text{Tr}(Y^\top L Y)\) is minimized (Rayleigh quotient),
2. \(Y\) is regularized to be orthogonal,

where \(L\) is the graph Laplacian constructed from the input data.

PARAMETER	DESCRIPTION
model	Neural network architecture for learning spectral embeddings. TYPE: ModelType | dict[str, ModelType]
kernel_fn	Kernel function specification for constructing affinity matrix. Kernel object: Pre-configured kernel instance str: Kernel type name ("gaussian" or "t") TYPE: Kernel | str, optional, by default "gaussian" DEFAULT: 'gaussian'
kernel_kwargs	Keyword arguments for kernel instantiation when kernel_fn is a string. TYPE: dict | None, optional, by default None DEFAULT: None
distance_fn	Distance metric specification for computing pairwise distances. PairwiseDistance object: Pre-configured distance instance str: Distance type name ("euclidean", "mahalanobis", "covariance") 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_weights	Weights for built-in loss components. Keys must match names in DEFAULT_LOSS_WEIGHTS. Unspecified keys keep their default weight. For SpectralNet, the available keys are: "loss_eigenvalue": Weight for the Laplacian (Rayleigh quotient) loss (default 1.0) "loss_orthogonal": Weight for the orthogonality constraint loss (default 1.0) Example: loss_weights={"loss_eigenvalue": 1.0, "loss_orthogonal": 0.1} TYPE: dict[str, float] | None, 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. TYPE: Literal["cuda", "cpu"], optional, by default "cuda" DEFAULT: 'cuda'

ATTRIBUTE	DESCRIPTION
DEFAULT_LOSS_WEIGHTS	Default weights for built-in loss components: {"loss_eigenvalue": 1.0, "loss_orthogonal": 1.0} TYPE: ClassVar[dict[str, float]]
kernel_fn	Kernel function instance. TYPE: Kernel
distance_fn	Distance function instance. TYPE: PairwiseDistance
loss_fn	Composite loss function containing spectral (Laplacian) and orthogonality losses. TYPE: CompositeLoss

Examples:

String-based configuration:

>>> from spectre.parametric import SpectralNet
>>> from spectre.core import Model
>>> import torch
>>>
>>> model = Model(in_features=10, out_features=3, multipliers=[2.0])
>>>
>>> snet = SpectralNet(
...     model=model,
...     kernel_fn="gaussian",
...     kernel_kwargs={"bw_method": "median"},
...     distance_fn="euclidean",
...     loss_weights={"loss_eigenvalue": 1.0, "loss_orthogonal": 0.1},
... )
>>>
>>> X = torch.randn(100, 10)
>>> loss = snet.loss(X)

Object-based configuration:

>>> from spectre.kernel import GaussianKernel
>>> from spectre.pairwise_distance import PairwiseDistanceEuclidean
>>>
>>> snet = SpectralNet(
...     model=model,
...     kernel_fn=GaussianKernel(bw_method="median"),
...     distance_fn=PairwiseDistanceEuclidean(),
... )

METHOD	DESCRIPTION
forward_step	Compute forward pass through the network.
loss	Compute loss combining orthogonality and Laplacian terms.
training_step	PyTorch Lightning training step.
score_step	Compute loss for evaluation without training status checks.
predict_states	Predict cluster assignments from embeddings.

Source code in spectre/parametric/spectral_net.py

def __init__(
    self,
    *,
    model: ModelType | 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: str = "laplacian",
    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,
    device: Literal["cuda", "cpu"] = "cuda",
    store_loss_context: bool = False,
) -> None:
    super().__init__(
        model=model,
        device=device,
        store_loss_context=store_loss_context,
        optimizer_fn=optimizer_fn,
        optimizer_kwargs=optimizer_kwargs,
        lr_scheduler_fn=lr_scheduler_fn,
        lr_scheduler_kwargs=lr_scheduler_kwargs,
    )

    self.kernel_fn = initialize_kernel_fn(kernel_fn, kernel_kwargs or {})
    self.distance_fn = initialize_distance_fn(distance_fn, distance_kwargs or {})

    if loss_kwargs is None:
        loss_kwargs = {}

    if loss_fn == "laplacian":
        self.loss_fn = LaplacianLoss(**loss_kwargs)
    else:
        raise ValueError(
            f"Unknown loss function '{loss_fn}'. Currently only 'laplacian' is supported."
        )

    self._loss_weights = {**self.DEFAULT_LOSS_WEIGHTS}
    if loss_weights is not None:
        if not isinstance(loss_weights, dict):
            raise TypeError(
                f"loss_weights must be dict or None, got {type(loss_weights)}."
            )
        unknown = set(loss_weights.keys()) - set(self.DEFAULT_LOSS_WEIGHTS)
        if unknown:
            raise ValueError(
                f"Unknown loss weight keys: {unknown}. "
                f"Available: {set(self.DEFAULT_LOSS_WEIGHTS.keys())}"
            )
        self._loss_weights.update(loss_weights)

    self.loss_fn = CompositeLoss(
        losses={
            "loss_eigenvalue": self.loss_fn,
            "loss_orthogonal": OrthogonalLoss(),
        },
        weights={
            "loss_eigenvalue": self._loss_weights["loss_eigenvalue"],
            "loss_orthogonal": self._loss_weights["loss_orthogonal"],
        },
    )

    if additional_loss is not None:
        self.loss_fn = self.loss_fn + CompositeLoss(
            losses=additional_loss, weights=additional_loss_weights
        )

Functions

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

Compute forward pass through the network.

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

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples, n_states)	Embedding vectors (approximate eigenvectors).

Source code in spectre/parametric/spectral_net.py

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

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

    Returns
    -------
    torch.Tensor of shape (n_samples, n_states)
        Embedding vectors (approximate eigenvectors).
    """
    return self.model(X)

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

Compute loss combining orthogonality and Laplacian terms.

PARAMETER	DESCRIPTION
X	Input data. TYPE: torch.Tensor of shape (n_samples, in_features)
weights	Sample weights for weighted loss computation. TYPE: torch.Tensor (n_samples,) or None, optional, by default None DEFAULT: None
target	Not used (unsupervised method). TYPE: torch.Tensor or None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor or tuple[Tensor, Tensor]	Scalar loss value or tuple of scalar loss values.

Source code in spectre/parametric/spectral_net.py

def loss(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
    """
    Compute loss combining orthogonality and Laplacian terms.

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

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

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

    Returns
    -------
    torch.Tensor or tuple[torch.Tensor, torch.Tensor]
        Scalar loss value or tuple of scalar loss values.
    """
    Z = self.forward_step(X)

    z_pdist = self.distance_fn(Z)
    z_kernel = self.kernel_fn(z_pdist, weights=weights)

    context = {
        "K": z_kernel,
        "Z": Z,
        "weights": weights,
    }

    if self.store_loss_context:
        self.loss_context = context

    loss = self.loss_fn(context)

    return loss

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

PyTorch Lightning training step.

PARAMETER	DESCRIPTION
batch	Training batch containing data and optional weights. TYPE: DataBatch
batch_idx	Index of current batch within the epoch. TYPE: int

RETURNS	DESCRIPTION
Tensor	Loss value.

Source code in spectre/parametric/spectral_net.py

def training_step(
    self,
    batch: DataBatch,
    batch_idx: int,
) -> torch.Tensor:
    """
    PyTorch Lightning training step.

    Parameters
    ----------
    batch : DataBatch
        Training batch containing data and optional weights.

    batch_idx : int
        Index of current batch within the epoch.

    Returns
    -------
    torch.Tensor
        Loss value.
    """
    loss = self.loss(X=batch.data, weights=batch.weights, target=batch.target)

    self.log_training_metrics(loss, **self.loss_fn.loss_components)

    return loss

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

Compute loss for evaluation without training status checks.

PARAMETER	DESCRIPTION
X	Input data for evaluation. TYPE: torch.Tensor of shape (n_samples, n_features)
weights	by default None Sample weights. TYPE: torch.Tensor of shape (n_samples,) or None, optional, DEFAULT: None
target	Not used (unsupervised method). TYPE: torch.Tensor or None, optional, by default None DEFAULT: None

RETURNS	DESCRIPTION
Tensor	Scalar loss value.

Source code in spectre/parametric/spectral_net.py

def score_step(
    self,
    X: torch.Tensor,
    weights: torch.Tensor | None = None,
    target: torch.Tensor | None = None,
) -> torch.Tensor:
    """
    Compute loss for evaluation without training status checks.

    Parameters
    ----------
    X : torch.Tensor of shape (n_samples, n_features)
        Input data for evaluation.

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

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

    Returns
    -------
    torch.Tensor
        Scalar loss value.
    """
    return self.loss(X, weights=weights, target=target)

predict_states(X: torch.Tensor, method: str = 'argmax') -> torch.Tensor

Predict cluster assignments from embeddings.

PARAMETER	DESCRIPTION
X	Input data. TYPE: torch.Tensor of shape (n_samples, n_features)
method	Clustering method to apply to embeddings. "argmax": Take argmax of embedding dimensions TYPE: str, optional, by default "argmax" DEFAULT: 'argmax'

RETURNS	DESCRIPTION
torch.Tensor of shape (n_samples,)	Cluster assignments.

Source code in spectre/parametric/spectral_net.py

def predict_states(self, X: torch.Tensor, method: str = "argmax") -> torch.Tensor:
    """
    Predict cluster assignments from embeddings.

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

    method : str, optional, by default "argmax"
        Clustering method to apply to embeddings.

        - "argmax": Take argmax of embedding dimensions

    Returns
    -------
    torch.Tensor of shape (n_samples,)
        Cluster assignments.
    """
    self.eval()
    with torch.no_grad():
        Z = self.forward_step(X)

        if method == "argmax":
            return torch.argmax(Z, dim=1)
        else:
            raise ValueError(
                f"Unknown clustering method '{method}'. Only 'argmax' is supported."
            )

Functions