Kernel Adaptive Gaussian

adaptive_gaussian

CLASS	DESCRIPTION
AdaptiveGaussianKernel	Adaptive Gaussian kernel with neural network-based bandwidth prediction.

FUNCTION	DESCRIPTION
adaptive_gaussian_kernel	Compute adaptive Gaussian kernel from distance matrix (functional interface).

Classes

AdaptiveGaussianKernel(bandwidth_network: torch.nn.Module, prediction_mode: str = 'sample', normalization: KernelNormalizer | str | dict | None = None, dtype: torch.dtype = torch.float32)

Bases: Kernel

Adaptive Gaussian kernel with neural network-based bandwidth prediction.

Implements a data-adaptive Gaussian kernel where the bandwidth is predicted by a neural network from the distance matrix. This enables the kernel to learn optimal, sample-specific bandwidths during training based on the pairwise distance structure.

The kernel form is: \(K(d_{kl}) = \exp(-d_{kl}^2 / \sigma_{kl}^2)\), but now \(\sigma_{kl}\) is predicted by a neural network from the distance matrix rather than being fixed or estimated via heuristics.

PARAMETER	DESCRIPTION
bandwidth_network	Neural network that predicts bandwidth from distance matrix. Expected input/output shapes depend on prediction_mode: "scalar": Input (n_samples, n_samples) -> Output (1,) or scalar. Single bandwidth for entire matrix "sample": Input (n_samples, n_samples) -> Output (n_samples,) or (n_samples, 1). One bandwidth per sample (creates matrix via outer product) TYPE: Module
prediction_mode	Mode for bandwidth prediction. "scalar": Single bandwidth for entire kernel matrix "sample": One bandwidth per sample (row-wise, uses outer product) TYPE: str, optional, by default "sample" DEFAULT: 'sample'
normalization	Kernel normalization. TYPE: KernelNormalizer | str | dict | None, optional, by default None DEFAULT: None
dtype	Data type for tensors. TYPE: torch.dtype, optional, by default torch.float32 DEFAULT: float32

ATTRIBUTE	DESCRIPTION
bandwidth_network	The neural network used for bandwidth prediction. TYPE: Module
prediction_mode	The bandwidth prediction mode being used. TYPE: str
bw	Predicted squared bandwidth matrix with shape matching the input distance matrix. Set after each forward pass. TYPE: Tensor

Examples:

Basic usage with sample-wise bandwidth prediction:

>>> import torch
>>> from spectre.kernel import AdaptiveGaussianKernel
>>>
>>> # Define a bandwidth predictor network that takes distance matrix
>>> # and outputs per-sample bandwidth
>>> class SampleBandwidthNet(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.net = torch.nn.Sequential(
...             torch.nn.Linear(1, 10),
...             torch.nn.ReLU(),
...             torch.nn.Linear(10, 1),
...             torch.nn.Softplus(),  # Ensure positive output
...         )
...
...     def forward(self, D):
...         # Use mean distance per sample as feature
...         features = D.mean(dim=1, keepdim=True)
...         return self.net(features).squeeze(-1)
>>>
>>> bandwidth_net = SampleBandwidthNet()
>>>
>>> # Create adaptive kernel
>>> kernel = AdaptiveGaussianKernel(
...     bandwidth_network=bandwidth_net,
...     prediction_mode="sample",
... )
>>>
>>> # Compute distances (from any source)
>>> X = torch.randn(100, 10)
>>> D = torch.cdist(X, X, p=2)
>>>
>>> # Apply kernel (only needs D)
>>> K = kernel(D)
>>> K.shape
torch.Size([100, 100])

Scalar bandwidth mode (single bandwidth for entire matrix):

>>> class ScalarBandwidthNet(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.net = torch.nn.Sequential(
...             torch.nn.Linear(1, 10),
...             torch.nn.ReLU(),
...             torch.nn.Linear(10, 1),
...             torch.nn.Softplus(),  # Ensure positive output
...         )
...
...     def forward(self, D):
...         # Use median distance as feature
...         median_dist = D.median().unsqueeze(0)
...         return self.net(median_dist)
>>>
>>> kernel = AdaptiveGaussianKernel(
...     bandwidth_network=ScalarBandwidthNet(),
...     prediction_mode="scalar",
... )
>>> K = kernel(D)

With normalization:

>>> from spectre.kernel import KernelNormalizer
>>> normalization = KernelNormalizer(norm_type="markov_symmetric")
>>> kernel = AdaptiveGaussianKernel(
...     bandwidth_network=bandwidth_net,
...     normalization=normalization,
... )
>>> K_normalized = kernel(D)

Integration with SpectralMap for end-to-end learning:

>>> from spectre.parametric import SpectralMap
>>> from spectre.kernel import AdaptiveGaussianKernel
>>>
>>> # Encoder network
>>> encoder = torch.nn.Sequential(
...     torch.nn.Linear(10, 20),
...     torch.nn.ReLU(),
...     torch.nn.Linear(20, 3),
... )
>>>
>>> # Bandwidth predictor learns from latent space distances
>>> class LatentBandwidthNet(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.net = torch.nn.Sequential(
...             torch.nn.Linear(1, 10),
...             torch.nn.ReLU(),
...             torch.nn.Linear(10, 1),
...             torch.nn.Softplus(),  # Ensure positive output
...         )
...
...     def forward(self, D):
...         features = D.mean(dim=1, keepdim=True)
...         return self.net(features).squeeze(-1)
>>>
>>> kernel = AdaptiveGaussianKernel(bandwidth_network=LatentBandwidthNet())
>>>
>>> # SpectralMap computes distances in latent space and passes to kernel
>>> spectral_map = SpectralMap(model=encoder, kernel_fn=kernel, n_states=3)
>>>
>>> # Train - kernel learns optimal bandwidth from distance patterns
>>> spectral_map.fit(data, n_epochs=100)

Notes

• The bandwidth network should output positive values (use Softplus, ReLU, or exponential activation)
• The network can use any features derived from the distance matrix (e.g., row means, medians, quantiles, k-NN distances)

METHOD	DESCRIPTION
forward_step	Compute adaptive Gaussian kernel matrix.
get_kernel_params	Return current predicted bandwidth from last forward pass.

Source code in spectre/kernel/adaptive_gaussian.py

def __init__(
    self,
    bandwidth_network: torch.nn.Module,
    prediction_mode: str = "sample",
    normalization: KernelNormalizer | str | dict | None = None,
    dtype: torch.dtype = torch.float32,
) -> None:
    # AdaptiveGaussianKernel always has learnable parameters (the network)
    super().__init__(
        learnable=True,
        learnable_log=False,  # Network handles parameterization
        normalization=normalization,
        dtype=dtype,
    )

    if not isinstance(bandwidth_network, torch.nn.Module):
        raise TypeError(
            f"Parameter `bandwidth_network` must be a torch.nn.Module, "
            f"got {type(bandwidth_network).__name__}."
        )

    if prediction_mode not in ["scalar", "sample"]:
        raise ValueError(
            f"Parameter `prediction_mode` must be 'scalar' or 'sample', "
            f"got '{prediction_mode}'."
        )

    self.bandwidth_network = bandwidth_network
    self.prediction_mode = prediction_mode

Functions

forward_step(D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs) -> torch.Tensor

Compute adaptive Gaussian kernel matrix.

PARAMETER	DESCRIPTION
D	Pairwise distance matrix of shape (n_samples, n_samples). TYPE: Tensor
weights	Sample weights (not used in bandwidth prediction, but may be used by normalization). TYPE: torch.Tensor | None, optional, by default None DEFAULT: None
**kwargs	Additional keyword arguments passed to the kernel. TYPE: dict DEFAULT: {}

RETURNS	DESCRIPTION
Tensor	Adaptive Gaussian kernel matrix of shape (n_samples, n_samples).

Notes

• Predicts bandwidth using bandwidth_network from distance matrix D
• Computes kernel using predicted bandwidth
• Caches predicted squared bandwidth in self.bw
• Normalization is handled by parent Kernel.forward() if specified

Source code in spectre/kernel/adaptive_gaussian.py

def forward_step(
    self, D: torch.Tensor, weights: torch.Tensor | None = None, **kwargs
) -> torch.Tensor:
    """
    Compute adaptive Gaussian kernel matrix.

    Parameters
    ----------
    D : torch.Tensor
        Pairwise distance matrix of shape (n_samples, n_samples).

    weights : torch.Tensor | None, optional, by default None
        Sample weights (not used in bandwidth prediction, but may be used
        by normalization).

    **kwargs : dict
        Additional keyword arguments passed to the kernel.

    Returns
    -------
    torch.Tensor
        Adaptive Gaussian kernel matrix of shape (n_samples, n_samples).

    Notes
    -----
    - Predicts bandwidth using `bandwidth_network` from distance matrix `D`
    - Computes kernel using predicted bandwidth
    - Caches predicted squared bandwidth in `self.bw`
    - Normalization is handled by parent `Kernel.forward()` if specified
    """
    result = adaptive_gaussian_kernel(
        D,
        bandwidth_network=self.bandwidth_network,
        prediction_mode=self.prediction_mode,
    )

    self.bw = result.bandwidth

    return result.kernel

get_kernel_params() -> dict[str, torch.Tensor]

Return current predicted bandwidth from last forward pass.

Returns the actual bandwidth matrix used in kernel computation. This is the squared bandwidth matrix (σ²) from the most recent forward pass.

RETURNS	DESCRIPTION
dict[str, Tensor]	Dictionary with "bandwidth" key containing the predicted squared bandwidth matrix from the last forward pass. Returns empty dict if forward() has not been called yet.

Examples:

>>> import torch
>>> from spectre.kernel import AdaptiveGaussianKernel
>>>
>>> class SimpleBandwidthNet(torch.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.net = torch.nn.Sequential(
...             torch.nn.Linear(1, 1),
...             torch.nn.Softplus(),
...         )
...
...     def forward(self, D):
...         features = D.mean(dim=1, keepdim=True)
...         return self.net(features).squeeze(-1)
>>>
>>> kernel = AdaptiveGaussianKernel(bandwidth_network=SimpleBandwidthNet())
>>> D = torch.rand(50, 50)
>>> _ = kernel(D)
>>> params = kernel.get_kernel_params()
>>> params["bandwidth"].shape
torch.Size([50, 50])

Notes

The returned bandwidth is σ², not σ, representing the actual values used in the kernel computation: K = exp(-D²/σ²).

Source code in spectre/kernel/adaptive_gaussian.py

def get_kernel_params(self) -> dict[str, torch.Tensor]:
    """
    Return current predicted bandwidth from last forward pass.

    Returns the actual bandwidth matrix used in kernel computation. This is
    the squared bandwidth matrix (σ²) from the most recent forward pass.

    Returns
    -------
    dict[str, torch.Tensor]
        Dictionary with "bandwidth" key containing the predicted squared
        bandwidth matrix from the last forward pass. Returns empty dict if
        `forward()` has not been called yet.

    Examples
    --------
    >>> import torch
    >>> from spectre.kernel import AdaptiveGaussianKernel
    >>>
    >>> class SimpleBandwidthNet(torch.nn.Module):
    ...     def __init__(self):
    ...         super().__init__()
    ...         self.net = torch.nn.Sequential(
    ...             torch.nn.Linear(1, 1),
    ...             torch.nn.Softplus(),
    ...         )
    ...
    ...     def forward(self, D):
    ...         features = D.mean(dim=1, keepdim=True)
    ...         return self.net(features).squeeze(-1)
    >>>
    >>> kernel = AdaptiveGaussianKernel(bandwidth_network=SimpleBandwidthNet())
    >>> D = torch.rand(50, 50)
    >>> _ = kernel(D)
    >>> params = kernel.get_kernel_params()
    >>> params["bandwidth"].shape
    torch.Size([50, 50])

    Notes
    -----
    The returned bandwidth is σ², not σ, representing the actual values
    used in the kernel computation: K = exp(-D²/σ²).
    """
    if hasattr(self, "bw"):
        return {"bandwidth": self.bw}
    else:
        return {}

Functions

adaptive_gaussian_kernel(D: torch.Tensor, bandwidth_network: torch.nn.Module, prediction_mode: str = 'sample') -> GaussianKernelResult

Compute adaptive Gaussian kernel from distance matrix (functional interface).

Pure function that computes an unnormalized adaptive Gaussian kernel where bandwidth is predicted by a neural network from the distance matrix.

For normalized kernels, use the :class:AdaptiveGaussianKernel class with a normalization.

PARAMETER	DESCRIPTION
D	Pairwise distance matrix of shape (n_samples, n_samples). TYPE: Tensor
bandwidth_network	Neural network that predicts bandwidth from distance matrix. TYPE: Module
prediction_mode	Mode for bandwidth prediction. "scalar": Single bandwidth for entire kernel matrix "sample": One bandwidth per sample (row-wise) TYPE: str, optional, by default "sample" DEFAULT: 'sample'
dtype	Data type for tensors. TYPE: torch.dtype, optional, by default torch.float32

RETURNS	DESCRIPTION
GaussianKernelResult	NamedTuple containing: kernel: Unnormalized adaptive Gaussian kernel matrix bandwidth: Predicted squared bandwidth matrix used in computation

Source code in spectre/kernel/adaptive_gaussian.py

def adaptive_gaussian_kernel(
    D: torch.Tensor,
    bandwidth_network: torch.nn.Module,
    prediction_mode: str = "sample",
) -> GaussianKernelResult:
    """
    Compute adaptive Gaussian kernel from distance matrix (functional interface).

    Pure function that computes an unnormalized adaptive Gaussian kernel where
    bandwidth is predicted by a neural network from the distance matrix.

    For normalized kernels, use the :class:`AdaptiveGaussianKernel` class with
    a normalization.

    Parameters
    ----------
    D : torch.Tensor
        Pairwise distance matrix of shape (n_samples, n_samples).

    bandwidth_network : torch.nn.Module
        Neural network that predicts bandwidth from distance matrix.

    prediction_mode : str, optional, by default "sample"
        Mode for bandwidth prediction.

        - "scalar": Single bandwidth for entire kernel matrix
        - "sample": One bandwidth per sample (row-wise)

    dtype : torch.dtype, optional, by default torch.float32
        Data type for tensors.

    Returns
    -------
    GaussianKernelResult
        NamedTuple containing:

        - kernel: Unnormalized adaptive Gaussian kernel matrix
        - bandwidth: Predicted squared bandwidth matrix used in computation
    """
    check_2d(D)

    # Predict bandwidth
    bw_pred = bandwidth_network(D)

    if prediction_mode == "scalar":
        # Ensure scalar shape
        if bw_pred.numel() != 1:
            raise ValueError(
                f"For 'scalar' mode, bandwidth network must output a scalar, "
                f"got shape {bw_pred.shape}."
            )
        bw_pred = bw_pred.reshape(())
        # Broadcast to matrix
        bw_matrix = (bw_pred**2).expand(D.shape[0], D.shape[1])

    elif prediction_mode == "sample":
        # Ensure shape is (n_samples,) or (n_samples, 1)
        if bw_pred.ndim == 2 and bw_pred.shape[1] == 1:
            bw_pred = bw_pred.squeeze(-1)
        if bw_pred.ndim != 1 or bw_pred.shape[0] != D.shape[0]:
            raise ValueError(
                f"For 'sample' mode, bandwidth network must output shape "
                f"(n_samples,) or (n_samples, 1), got {bw_pred.shape}. "
                f"Expected n_samples={D.shape[0]}."
            )
        # Broadcast to matrix via outer product
        bw_matrix = torch.outer(bw_pred, bw_pred)

    else:
        raise ValueError(
            f"Parameter `prediction_mode` must be 'scalar' or 'sample', "
            f"got '{prediction_mode}'."
        )

    # Compute Gaussian kernel
    kernel = torch.exp(-torch.pow(D, 2) / bw_matrix)

    return GaussianKernelResult(kernel=kernel, bandwidth=bw_matrix)