Metrics Kernel

kernel

FUNCTION	DESCRIPTION
kernel_alignment	Compute centered kernel alignment (CKA) between two kernel matrices.
hsic	Compute Hilbert-Schmidt Independence Criterion (HSIC) between two kernel matrices.
kernel_frobenius_error	Compute Frobenius norm error between predicted and true kernel matrices.
effective_rank	Compute effective rank of kernel matrix.

Functions

kernel_alignment(K: torch.Tensor, M: torch.Tensor, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute centered kernel alignment (CKA) between two kernel matrices.

CKA measures similarity between two kernel matrices in a way that is invariant to orthogonal transformations and isotropic scaling [1]. CKA is defined in [0, 1]. Value of 1 indicates perfect alignment.

The metric is computed as: \(\mathrm{CKA}(K, M) = \left< K^c, M^c\right>/(\|K^c\| \|M^c\|)\), where \(c\) denotes the centered kernel matrix.

PARAMETER	DESCRIPTION
K	First kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
M	Second kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
eps	Small value to avoid division by zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Centered kernel alignment.

Examples:

>>> K = torch.randn(50, 50)
>>> K = K @ K.T  # Make positive semidefinite
>>> M = K + 0.1 * torch.randn(50, 50)
>>> kernel_alignment(K, M)
0.95...

Perfect alignment:

>>> kernel_alignment(K, K)
1.0

Source code in spectre/metrics/kernel.py

def kernel_alignment(
    K: torch.Tensor, M: torch.Tensor, eps: float = torch.finfo(torch.float32).eps
) -> torch.Tensor:
    """
    Compute centered kernel alignment (CKA) between two kernel matrices.

    CKA measures similarity between two kernel matrices in a way that is
    invariant to orthogonal transformations and isotropic scaling
    [@kornblith2019similarity]. CKA is defined in [0, 1]. Value of 1 indicates perfect
    alignment.

    The metric is computed as:
    $\\mathrm{CKA}(K, M) = \\left< K^c, M^c\\right>/(\\|K^c\\| \\|M^c\\|)$,
    where $c$ denotes the centered kernel matrix.

    Parameters
    ----------
    K : torch.Tensor
        First kernel matrix, shape (n_samples, n_samples).

    M : torch.Tensor
        Second kernel matrix, shape (n_samples, n_samples).

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small value to avoid division by zero.

    Returns
    -------
    torch.Tensor
        Centered kernel alignment.

    Examples
    --------
    >>> K = torch.randn(50, 50)
    >>> K = K @ K.T  # Make positive semidefinite
    >>> M = K + 0.1 * torch.randn(50, 50)
    >>> kernel_alignment(K, M)
    0.95...

    Perfect alignment:

    >>> kernel_alignment(K, K)
    1.0
    """
    check_2d(K)
    check_2d(M)
    check_square_shape(K)
    check_same_shape(K, M)
    check_same_device(K, M)

    # Center kernel matrices
    # K_c = K - row_mean - col_mean + grand_mean
    row_mean_K = K.mean(dim=1, keepdim=True)
    col_mean_K = K.mean(dim=0, keepdim=True)
    grand_mean_K = K.mean()
    K_c = K - row_mean_K - col_mean_K + grand_mean_K

    row_mean_M = M.mean(dim=1, keepdim=True)
    col_mean_M = M.mean(dim=0, keepdim=True)
    grand_mean_M = M.mean()
    M_c = M - row_mean_M - col_mean_M + grand_mean_M

    # Compute Frobenius inner product
    numerator = torch.sum(K_c * M_c)

    # Compute Frobenius norms
    norm_K_sq = torch.sum(K_c * K_c)
    norm_M_sq = torch.sum(M_c * M_c)

    if norm_K_sq < eps or norm_M_sq < eps:
        raise ValueError(
            "One or both centered kernel matrices have near-zero Frobenius norm."
        )

    cka = numerator / torch.sqrt(norm_K_sq * norm_M_sq)

    return cka

hsic(K: torch.Tensor, M: torch.Tensor, unbiased: bool = True, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute Hilbert-Schmidt Independence Criterion (HSIC) between two kernel matrices.

HSIC measures the dependence between two sets of variables using kernel methods. The unbiased estimator removes diagonal bias terms for improved statistical properties with finite samples. Higher values indicate stronger dependence. Unlike CKA, HSIC is unbounded and not scale-invariant.

The biased HSIC estimator is: \(\mathrm{HSIC}_b(K_1, K_2) = \frac{1}{n^2} \mathrm{trace}(K_1 H K_2 H)\) where \(H\) is the centering matrix. The unbiased estimator removes diagonal terms.

PARAMETER	DESCRIPTION
K	First kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
M	Second kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
unbiased	If True, use the unbiased HSIC estimator. If False, use biased version. TYPE: bool, optional, by default True DEFAULT: True
eps	Small value to avoid division by zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	HSIC value.

Examples:

>>> K = torch.randn(50, 50)
>>> K = K @ K.T
>>> M = torch.randn(50, 50)
>>> M = M @ M.T
>>> hsic_unbiased(K, M)
125.3...

Biased estimator:

>>> hsic_unbiased(K, M, unbiased=False)
130.1...

HSIC is not scale-invariant (unlike CKA):

>>> hsic_unbiased(K, 10 * K)  # 100x larger than hsic_unbiased(K, K)

Source code in spectre/metrics/kernel.py

def hsic(
    K: torch.Tensor,
    M: torch.Tensor,
    unbiased: bool = True,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute Hilbert-Schmidt Independence Criterion (HSIC) between two kernel matrices.

    HSIC measures the dependence between two sets of variables using kernel methods.
    The unbiased estimator removes diagonal bias terms for improved statistical
    properties with finite samples. Higher values indicate stronger dependence.
    Unlike CKA, HSIC is unbounded and not scale-invariant.

    The biased HSIC estimator is:
    $\\mathrm{HSIC}_b(K_1, K_2) = \\frac{1}{n^2} \\mathrm{trace}(K_1 H K_2 H)$
    where $H$ is the centering matrix. The unbiased estimator removes diagonal terms.

    Parameters
    ----------
    K : torch.Tensor
        First kernel matrix, shape (n_samples, n_samples).

    M : torch.Tensor
        Second kernel matrix, shape (n_samples, n_samples).

    unbiased : bool, optional, by default True
        If True, use the unbiased HSIC estimator. If False, use biased version.

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small value to avoid division by zero.

    Returns
    -------
    torch.Tensor
        HSIC value.

    Examples
    --------
    >>> K = torch.randn(50, 50)
    >>> K = K @ K.T
    >>> M = torch.randn(50, 50)
    >>> M = M @ M.T
    >>> hsic_unbiased(K, M)
    125.3...

    Biased estimator:

    >>> hsic_unbiased(K, M, unbiased=False)
    130.1...

    HSIC is not scale-invariant (unlike CKA):

    >>> hsic_unbiased(K, 10 * K)  # 100x larger than hsic_unbiased(K, K)
    """
    check_2d(K)
    check_2d(M)
    check_square_shape(K)
    check_same_shape(K, M)
    check_same_device(K, M)

    n_samples = K.shape[0]

    if n_samples < 4 and unbiased:
        raise ValueError("Unbiased HSIC requires at least 4 samples.")

    if unbiased:
        # Remove diagonal terms
        diag_mask = ~torch.eye(n_samples, dtype=torch.bool, device=K.device)
        K_zeroed = K * diag_mask
        M_zeroed = M * diag_mask

        trace_term = torch.sum(K_zeroed * M_zeroed)
        cross_term = torch.sum(K_zeroed @ M_zeroed)

        hsic_val = (
            trace_term
            + K_zeroed.sum() * M_zeroed.sum() / ((n_samples - 1) * (n_samples - 2))
            - 2 * cross_term / (n_samples - 2)
        ) / (n_samples * (n_samples - 3))

    else:
        # Biased estimator
        row_mean_K = K.mean(dim=1, keepdim=True)
        col_mean_K = K.mean(dim=0, keepdim=True)
        grand_mean_K = K.mean()
        K_c = K - row_mean_K - col_mean_K + grand_mean_K

        row_mean_M = M.mean(dim=1, keepdim=True)
        col_mean_M = M.mean(dim=0, keepdim=True)
        grand_mean_M = M.mean()
        M_c = M - row_mean_M - col_mean_M + grand_mean_M

        # trace(K_c @ M_c) = sum(K_c * M_c) when both are symmetric
        hsic_val = torch.sum(K_c * M_c) / (n_samples * n_samples)

    return hsic_val

kernel_frobenius_error(K: torch.Tensor, M: torch.Tensor, normalize: bool = True, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute Frobenius norm error between predicted and true kernel matrices.

Measures direct element-wise reconstruction quality of kernel matrices.

PARAMETER	DESCRIPTION
K	Predicted kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
M	True/reference kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
normalize	If True, normalize by Frobenius norm for relative error. TYPE: bool, optional, by default True DEFAULT: True
eps	Small value to avoid division by zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
float	Frobenius reconstruction error. Lower values indicate better reconstruction.

Examples:

>>> M = torch.randn(50, 50)
>>> M = M @ M.T
>>> K = M + 0.1 * torch.randn(50, 50)
>>> kernel_frobenius_error(K, M, normalize=True)
0.05...

Absolute error:

>>> kernel_frobenius_error(K, M, normalize=False)
2.5...

Source code in spectre/metrics/kernel.py

def kernel_frobenius_error(
    K: torch.Tensor,
    M: torch.Tensor,
    normalize: bool = True,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute Frobenius norm error between predicted and true kernel matrices.

    Measures direct element-wise reconstruction quality of kernel matrices.

    Parameters
    ----------
    K : torch.Tensor
        Predicted kernel matrix, shape (n_samples, n_samples).

    M : torch.Tensor
        True/reference kernel matrix, shape (n_samples, n_samples).

    normalize : bool, optional, by default True
        If True, normalize by Frobenius norm for relative error.

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small value to avoid division by zero.

    Returns
    -------
    float
        Frobenius reconstruction error. Lower values indicate better reconstruction.

    Examples
    --------
    >>> M = torch.randn(50, 50)
    >>> M = M @ M.T
    >>> K = M + 0.1 * torch.randn(50, 50)
    >>> kernel_frobenius_error(K, M, normalize=True)
    0.05...

    Absolute error:

    >>> kernel_frobenius_error(K, M, normalize=False)
    2.5...
    """
    check_2d(K)
    check_2d(M)
    check_square_shape(K)
    check_same_shape(K, M)
    check_same_device(K, M)

    # Compute Frobenius norm of difference
    error = torch.norm(K - M, p="fro")

    if normalize:
        norm_true = torch.norm(M, p="fro")
        if norm_true < eps:
            raise ValueError(
                "Cannot normalize: true kernel has near-zero Frobenius norm."
            )
        error = error / norm_true

    return error

effective_rank(eigenvalues: torch.Tensor | None = None, K: torch.Tensor | None = None, method: str = 'threshold', threshold: float = 0.99, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute effective rank of kernel matrix.

The effective rank measures the intrinsic dimensionality of the data represented by the kernel matrix. Two methods are supported:

• "threshold": Number of eigenvalues capturing specified variance fraction
• "entropy": Exponential of Shannon entropy of eigenvalue distribution

PARAMETER	DESCRIPTION
eigenvalues	Eigenvalues of the kernel matrix, shape (n_samples,). Should be sorted in descending order. If None, eigenvalues will be computed from the kernel matrix "K". TYPE: torch.Tensor | None = None DEFAULT: None
K	Kernel matrix, shape (n_samples, n_samples). Should be positive semidefinite. Used to compute eigenvalues if not provided. TYPE: torch.Tensor | None = None DEFAULT: None
method	Method to compute effective rank. Options: "threshold", "entropy". TYPE: str, optional, by default "threshold" DEFAULT: 'threshold'
threshold	Fraction of variance to capture in (0, 1). Only used when method="threshold". TYPE: float, optional, by default 0.99 DEFAULT: 0.99
eps	Small value to avoid division by zero and log of zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Effective rank. Integer for "threshold" method, continuous for "entropy" method.

Examples:

>>> # Low-rank kernel (rank ~5)
>>> X = torch.randn(100, 5)
>>> K = X @ X.T
>>> effective_rank(K, method="threshold", threshold=0.99)
5

>>> # Entropy-based measure
>>> effective_rank(K, method="entropy")
4.8...

>>> # Full-rank kernel
>>> K_full = torch.randn(100, 100)
>>> K_full = K_full @ K_full.T
>>> effective_rank(K_full, method="threshold", threshold=0.99)
85...

Source code in spectre/metrics/kernel.py

def effective_rank(
    eigenvalues: torch.Tensor | None = None,
    K: torch.Tensor | None = None,
    method: str = "threshold",
    threshold: float = 0.99,
    eps: float = torch.finfo(torch.float32).eps,
) -> torch.Tensor:
    """
    Compute effective rank of kernel matrix.

    The effective rank measures the intrinsic dimensionality of the data
    represented by the kernel matrix. Two methods are supported:

    - "threshold": Number of eigenvalues capturing specified variance fraction
    - "entropy": Exponential of Shannon entropy of eigenvalue distribution

    Parameters
    ----------
    eigenvalues : torch.Tensor | None = None
        Eigenvalues of the kernel matrix, shape (n_samples,). Should be sorted
        in descending order. If None, eigenvalues will be computed from
        the kernel matrix "K".

    K : torch.Tensor | None = None
        Kernel matrix, shape (n_samples, n_samples). Should be positive semidefinite.
        Used to compute eigenvalues if not provided.

    method : str, optional, by default "threshold"
        Method to compute effective rank. Options: "threshold", "entropy".

    threshold : float, optional, by default 0.99
        Fraction of variance to capture in (0, 1). Only used when `method="threshold"`.

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small value to avoid division by zero and log of zero.

    Returns
    -------
    torch.Tensor
        Effective rank. Integer for "threshold" method, continuous for "entropy" method.

    Examples
    --------
    >>> # Low-rank kernel (rank ~5)
    >>> X = torch.randn(100, 5)
    >>> K = X @ X.T
    >>> effective_rank(K, method="threshold", threshold=0.99)
    5

    >>> # Entropy-based measure
    >>> effective_rank(K, method="entropy")
    4.8...

    >>> # Full-rank kernel
    >>> K_full = torch.randn(100, 100)
    >>> K_full = K_full @ K_full.T
    >>> effective_rank(K_full, method="threshold", threshold=0.99)
    85...
    """
    if method not in ["threshold", "entropy"]:
        raise ValueError(f"method must be 'threshold' or 'entropy', got '{method}'.")

    if method == "threshold" and not (0 < threshold < 1):
        raise ValueError(f"Expected threshold in (0, 1), got {threshold}.")

    if eigenvalues is None and K is not None:
        check_2d(K)
        check_square_shape(K)
        n_samples = K.shape[0]
        try:
            eigenvalues = eigvals(K, n_eigen=0, symmetric=True, sort_descending=True)
        except (RuntimeError, ValueError):
            raise ValueError("Failed to compute eigenvalues of kernel matrix.")
    elif eigenvalues is not None:
        check_1d(eigenvalues)
        n_samples = eigenvalues.shape[0]
    else:
        raise ValueError("Either K or eigenvalues must be provided.")

    # Handle negative eigenvalues (numerical errors for PSD matrices)
    eigenvalues = torch.clamp(eigenvalues, min=0.0)
    if eigenvalues.sum() < eps:
        raise ValueError("Kernel matrix has near-zero total variance.")

    if method == "threshold":
        cumulative_variance = torch.cumsum(eigenvalues, dim=0) / eigenvalues.sum()

        # searchsorted returns 0-based index, add 1 for rank (1-indexed)
        idx = torch.searchsorted(cumulative_variance, threshold, right=False).item()
        effective_rank_value = min(idx + 1, n_samples)

        return torch.tensor(effective_rank_value, dtype=torch.int64)

    elif method == "entropy":
        eigenvalue_prob = eigenvalues / eigenvalues.sum()
        # Filter out near-zero probabilities to avoid log(0)
        eigenvalue_prob = eigenvalue_prob[eigenvalue_prob > eps]

        # Compute Shannon entropy. Effective rank is exp(entropy)
        entropy = -torch.sum(eigenvalue_prob * torch.log(eigenvalue_prob))
        effective_rank_value = torch.exp(entropy)

        return effective_rank_value

effective_dimensions(K: torch.Tensor, M: torch.Tensor, eps: float = torch.finfo(torch.float32).eps) -> torch.Tensor

Compute the Effective Number of Shared Dimensions (ENSD) between two kernel matrices.

ENSD is a scalar quantity that generalizes the participation ratio (PR) to a pair of datasets [2]. For kernel matrices K and M it is:

\(ENSD(K, M) = Tr(K) Tr(M) Tr(K M) / (Tr(K^2) Tr(M^2))\)

PARAMETER	DESCRIPTION
K	First kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
M	Second kernel matrix, shape (n_samples, n_samples). TYPE: Tensor
eps	Small value to guard against division by zero. TYPE: float, optional, by default torch.finfo(torch.float32).eps DEFAULT: eps

RETURNS	DESCRIPTION
Tensor	Scalar ENSD value (effective number of shared dimensions).

Examples:

>>> X = torch.randn(100, 50)
>>> Y = X[:, :25]  # Y lives in a 25-dim subspace of X
>>> K = X @ X.T
>>> M = Y @ Y.T
>>> effective_number_shared_dimensions(K, M)
tensor(25.0, ...)

Source code in spectre/metrics/kernel.py

def effective_dimensions(
    K: torch.Tensor, M: torch.Tensor, eps: float = torch.finfo(torch.float32).eps
) -> torch.Tensor:
    """
    Compute the Effective Number of Shared Dimensions (ENSD) between two kernel
    matrices.

    ENSD is a scalar quantity that generalizes the participation ratio (PR)
    to a pair of datasets [@giaffar2024effective]. For kernel matrices
    K and M it is:

    $ENSD(K, M) = Tr(K) Tr(M) Tr(K M) / (Tr(K^2) Tr(M^2))$

    Parameters
    ----------
    K : torch.Tensor
        First kernel matrix, shape (n_samples, n_samples).

    M : torch.Tensor
        Second kernel matrix, shape (n_samples, n_samples).

    eps : float, optional, by default torch.finfo(torch.float32).eps
        Small value to guard against division by zero.

    Returns
    -------
    torch.Tensor
        Scalar ENSD value (effective number of shared dimensions).

    Examples
    --------
    >>> X = torch.randn(100, 50)
    >>> Y = X[:, :25]  # Y lives in a 25-dim subspace of X
    >>> K = X @ X.T
    >>> M = Y @ Y.T
    >>> effective_number_shared_dimensions(K, M)
    tensor(25.0, ...)
    """
    check_2d(K)
    check_2d(M)
    check_square_shape(K)
    check_same_shape(K, M)
    check_same_device(K, M)

    # Center kernel matrices (same as in kernel_alignment)
    row_mean_K = K.mean(dim=1, keepdim=True)
    col_mean_K = K.mean(dim=0, keepdim=True)
    grand_mean_K = K.mean()
    K_c = K - row_mean_K - col_mean_K + grand_mean_K

    row_mean_M = M.mean(dim=1, keepdim=True)
    col_mean_M = M.mean(dim=0, keepdim=True)
    grand_mean_M = M.mean()
    M_c = M - row_mean_M - col_mean_M + grand_mean_M

    # Trace terms
    tr_K = torch.trace(K_c)
    tr_M = torch.trace(M_c)

    # Frobenius norm squared terms: Tr(K^2) = ||K||_F^2 for symmetric K
    tr_KK = torch.sum(K_c * K_c)
    tr_MM = torch.sum(M_c * M_c)

    if tr_KK < eps or tr_MM < eps:
        raise ValueError(
            "One or both centered kernel matrices have near-zero Frobenius norm. ENSD is not well-defined."
        )

    # Tr(K M) = Frobenius inner product for symmetric matrices
    tr_KM = torch.sum(K_c * M_c)

    ensd = tr_K * tr_M * tr_KM / (tr_KK * tr_MM)

    return ensd