Spectral Kernels

spectrans.kernels.spectral

Spectral kernel functions for attention mechanisms.

This module implements kernel functions based on spectral decomposition and eigenfunction expansions. These kernels provide alternatives to RFF-based approximations for spectral attention mechanisms.

The implementations include polynomial spectral kernels, eigenvalue-based decompositions, and learnable spectral filters that can be optimized during training.

Classes:

Name	Description
SpectralKernel	Base class for spectral kernel functions.
PolynomialSpectralKernel	Polynomial kernel with spectral decomposition.
TruncatedSVDKernel	Kernel approximation via truncated SVD.
LearnableSpectralKernel	Spectral kernel with learnable eigenvalues.
FourierKernel	Kernel defined in Fourier domain.

Examples:

Using polynomial spectral kernel:

>>> import torch
>>> from spectrans.kernels.spectral import PolynomialSpectralKernel
>>> kernel = PolynomialSpectralKernel(rank=32, degree=3)
>>> Q, K = torch.randn(2, 100, 64), torch.randn(2, 100, 64)
>>> attention_weights = kernel.compute_attention(Q, K)
>>> assert attention_weights.shape == (2, 100, 100)

Learnable spectral kernel:

>>> from spectrans.kernels.spectral import LearnableSpectralKernel
>>> kernel = LearnableSpectralKernel(input_dim=64, rank=16)
>>> features = kernel.extract_features(Q)
>>> assert features.shape == (2, 100, 16)

Notes

Spectral kernels leverage eigendecomposition for kernel computation through the representation \(K(\mathbf{X}, \mathbf{Y}) = \Phi(\mathbf{X}) \mathbf{\Lambda} \Phi(\mathbf{Y})^T\) where \(\Phi\) are eigenfunctions and \(\mathbf{\Lambda}\) are eigenvalues.

This decomposition enables low-rank approximations via truncation of the eigenspectrum and learnable spectral filters through trainable eigenvalues. The rank parameter determines the number of eigenmodes retained in the approximation.

References

Yunyang Chen, Yingfeng Luo, and Liping Zhang. 2021. Scatterbrain: Unifying sparse and low-rank attention approximation. In Advances in Neural Information Processing Systems 34 (NeurIPS 2021).

Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. 2020. Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768.

See Also

spectrans.kernels.base : Base kernel interfaces. spectrans.kernels.rff : Random Fourier Features.

Classes

SpectralKernel

SpectralKernel(rank: int, normalize: bool = True)

Bases: KernelFunction

Base class for spectral kernel functions.

Spectral kernels use eigendecomposition or spectral analysis for efficient kernel computation.

Parameters:

Name	Type	Description	Default
rank	int	Rank of spectral approximation.	required
normalize	bool	Whether to normalize kernel values.	True

Attributes:

Name	Type	Description
rank	int	Approximation rank.
normalize	bool	Normalization flag.

Methods:

Name	Description
spectral_decomposition	Compute spectral decomposition of input.

Source code in spectrans/kernels/spectral.py

def __init__(self, rank: int, normalize: bool = True):
    self.rank = rank
    self.normalize = normalize

Functions

spectral_decomposition

spectral_decomposition(x: Tensor) -> tuple[Tensor, Tensor]

Compute spectral decomposition of input.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor of shape (..., n, d).	required

Returns:

Name	Type	Description
eigenvectors	Tensor	Eigenvectors of shape (..., n, rank).
eigenvalues	Tensor	Eigenvalues of shape (..., rank).

Source code in spectrans/kernels/spectral.py

def spectral_decomposition(self, x: Tensor) -> tuple[Tensor, Tensor]:
    """Compute spectral decomposition of input.

    Parameters
    ----------
    x : Tensor
        Input tensor of shape (..., n, d).

    Returns
    -------
    eigenvectors : Tensor
        Eigenvectors of shape (..., n, rank).
    eigenvalues : Tensor
        Eigenvalues of shape (..., rank).
    """
    # Compute Gram matrix
    gram = torch.matmul(x, x.transpose(-2, -1))

    # Eigendecomposition
    eigenvalues, eigenvectors = torch.linalg.eigh(gram)

    # Keep top-k eigenvalues/vectors
    eigenvalues = eigenvalues[..., -self.rank :]
    eigenvectors = eigenvectors[..., -self.rank :]

    if self.normalize:
        # Normalize by trace
        trace = eigenvalues.sum(dim=-1, keepdim=True)
        eigenvalues = eigenvalues / (trace + 1e-8)

    return eigenvectors, eigenvalues

PolynomialSpectralKernel

PolynomialSpectralKernel(rank: int, degree: int = 2, coef0: float = 1.0, alpha: float = 1.0, normalize: bool = True)

Bases: SpectralKernel

Polynomial kernel with spectral decomposition.

Computes \((\mathbf{X}\mathbf{Y}^T + c)^d\) using eigendecomposition.

Parameters:

Name	Type	Description	Default
rank	int	Rank of spectral approximation.	required
degree	int	Polynomial degree.	2
coef0	float	Constant coefficient.	1.0
alpha	float	Scaling factor.	1.0
normalize	bool	Whether to normalize.	True

Attributes:

Name	Type	Description
degree	int	Polynomial degree.
coef0	float	Constant term.
alpha	float	Scale factor.

Methods:

Name	Description
compute	Compute polynomial spectral kernel.
compute_attention	Compute attention weights using spectral decomposition.

Source code in spectrans/kernels/spectral.py

def __init__(
    self,
    rank: int,
    degree: int = 2,
    coef0: float = 1.0,
    alpha: float = 1.0,
    normalize: bool = True,
):
    super().__init__(rank, normalize)
    self.degree = degree
    self.coef0 = coef0
    self.alpha = alpha

Functions

compute

compute(x: Tensor, y: Tensor) -> Tensor

Compute polynomial spectral kernel.

Parameters:

Name	Type	Description	Default
x	Tensor	First input of shape (..., n, d).	required
y	Tensor	Second input of shape (..., m, d).	required

Returns:

Type	Description
Tensor	Kernel matrix of shape (..., n, m).

Source code in spectrans/kernels/spectral.py

def compute(self, x: Tensor, y: Tensor) -> Tensor:
    """Compute polynomial spectral kernel.

    Parameters
    ----------
    x : Tensor
        First input of shape (..., n, d).
    y : Tensor
        Second input of shape (..., m, d).

    Returns
    -------
    Tensor
        Kernel matrix of shape (..., n, m).
    """
    # Standard polynomial kernel
    inner = torch.matmul(x, y.transpose(-2, -1))
    kernel = (self.alpha * inner + self.coef0) ** self.degree

    if self.normalize:
        # Normalize by geometric mean of norms
        x_norm = torch.norm(x, dim=-1, keepdim=True)
        y_norm = torch.norm(y, dim=-1, keepdim=True)
        norm_matrix = torch.matmul(x_norm, y_norm.transpose(-2, -1))
        kernel = kernel / (norm_matrix + 1e-8)

    return kernel

compute_attention

compute_attention(q: Tensor, k: Tensor) -> Tensor

Compute attention weights using spectral decomposition.

Parameters:

Name	Type	Description	Default
q	Tensor	Queries of shape (..., n, d).	required
k	Tensor	Keys of shape (..., m, d).	required

Returns:

Type	Description
Tensor	Attention weights of shape (..., n, m).

Source code in spectrans/kernels/spectral.py

def compute_attention(self, q: Tensor, k: Tensor) -> Tensor:
    """Compute attention weights using spectral decomposition.

    Parameters
    ----------
    q : Tensor
        Queries of shape (..., n, d).
    k : Tensor
        Keys of shape (..., m, d).

    Returns
    -------
    Tensor
        Attention weights of shape (..., n, m).
    """
    # Low-rank approximation via SVD
    # Q = U_q S_q V_q^T, K = U_k S_k V_k^T

    # Compute QK^T approximately
    q_reduced = self._reduce_rank(q)  # (..., n, r)
    k_reduced = self._reduce_rank(k)  # (..., m, r)

    # Polynomial kernel in reduced space
    inner = torch.matmul(q_reduced, k_reduced.transpose(-2, -1))
    attention = (self.alpha * inner + self.coef0) ** self.degree

    if self.normalize:
        attention = F.softmax(attention, dim=-1)

    return attention

TruncatedSVDKernel

TruncatedSVDKernel(rank: int, normalize: bool = True, use_randomized: bool = False)

Bases: SpectralKernel

Kernel approximation via truncated SVD.

Uses SVD to compute low-rank approximation of kernel matrix.

Parameters:

Name	Type	Description	Default
rank	int	Truncation rank.	required
normalize	bool	Whether to normalize.	True
use_randomized	bool	Use randomized SVD for large matrices.	False

Attributes:

Name	Type	Description
use_randomized	bool	Whether to use randomized algorithms.

Methods:

Name	Description
compute	Compute kernel via truncated SVD.

Source code in spectrans/kernels/spectral.py

def __init__(
    self,
    rank: int,
    normalize: bool = True,
    use_randomized: bool = False,
):
    super().__init__(rank, normalize)
    self.use_randomized = use_randomized

Functions

compute

compute(x: Tensor, y: Tensor) -> Tensor

Compute kernel via truncated SVD.

Parameters:

Name	Type	Description	Default
x	Tensor	First input of shape (..., n, d).	required
y	Tensor	Second input of shape (..., m, d).	required

Returns:

Type	Description
Tensor	Approximate kernel matrix of shape (..., n, m).

Source code in spectrans/kernels/spectral.py

def compute(self, x: Tensor, y: Tensor) -> Tensor:
    """Compute kernel via truncated SVD.

    Parameters
    ----------
    x : Tensor
        First input of shape (..., n, d).
    y : Tensor
        Second input of shape (..., m, d).

    Returns
    -------
    Tensor
        Approximate kernel matrix of shape (..., n, m).
    """
    # Compute full kernel matrix
    kernel_full = torch.matmul(x, y.transpose(-2, -1))

    if self.use_randomized:
        # Randomized SVD (faster for large matrices)
        kernel_approx = self._randomized_svd_approximation(kernel_full)
    else:
        # Standard SVD
        U, S, Vt = torch.linalg.svd(kernel_full, full_matrices=False)

        # Truncate to rank
        U_r = U[..., : self.rank]
        S_r = S[..., : self.rank]
        Vt_r = Vt[..., : self.rank, :]

        # Reconstruct
        kernel_approx = torch.matmul(U_r * S_r.unsqueeze(-2), Vt_r)

    if self.normalize:
        # Normalize rows
        row_norms = kernel_approx.norm(dim=-1, keepdim=True)
        kernel_approx = kernel_approx / (row_norms + 1e-8)

    return kernel_approx

LearnableSpectralKernel

LearnableSpectralKernel(input_dim: int, rank: int, init_scale: float = 1.0, trainable_eigenvectors: bool = True, normalize: bool = True)

Bases: Module, SpectralKernel

Spectral kernel with learnable eigenvalues and eigenfunctions.

Parameters:

Name	Type	Description	Default
input_dim	int	Input dimension.	required
rank	int	Number of spectral components.	required
init_scale	float	Initialization scale.	1.0
trainable_eigenvectors	bool	Whether eigenvectors are trainable.	True
normalize	bool	Whether to normalize.	True

Attributes:

Name	Type	Description
eigenvectors	Parameter	Learnable eigenvectors of shape (input_dim, rank).
eigenvalues	Parameter	Learnable eigenvalues of shape (rank,).

Methods:

Name	Description
compute	Compute learnable spectral kernel.
extract_features	Extract spectral features.
forward	Forward pass for nn.Module compatibility.
orthogonalize_eigenvectors	Orthogonalize eigenvectors via Gram-Schmidt.

Source code in spectrans/kernels/spectral.py

def __init__(
    self,
    input_dim: int,
    rank: int,
    init_scale: float = 1.0,
    trainable_eigenvectors: bool = True,
    normalize: bool = True,
):
    nn.Module.__init__(self)
    SpectralKernel.__init__(self, rank, normalize)

    self.input_dim = input_dim
    self.trainable_eigenvectors = trainable_eigenvectors

    # Initialize eigenvectors (orthogonal)
    eigenvectors = torch.randn(input_dim, rank) * init_scale
    eigenvectors, _ = torch.linalg.qr(eigenvectors)

    if trainable_eigenvectors:
        self.eigenvectors = nn.Parameter(eigenvectors)
    else:
        self.register_buffer("eigenvectors", eigenvectors)

    # Initialize eigenvalues (positive, decreasing)
    eigenvalues = torch.linspace(1.0, 0.1, rank) * init_scale
    self.eigenvalues = nn.Parameter(eigenvalues)

Functions

compute

compute(x: Tensor, y: Tensor) -> Tensor

Compute learnable spectral kernel.

Parameters:

Name	Type	Description	Default
x	Tensor	First input of shape (..., n, d).	required
y	Tensor	Second input of shape (..., m, d).	required

Returns:

Type	Description
Tensor	Kernel matrix of shape (..., n, m).

Source code in spectrans/kernels/spectral.py

def compute(self, x: Tensor, y: Tensor) -> Tensor:
    """Compute learnable spectral kernel.

    Parameters
    ----------
    x : Tensor
        First input of shape (..., n, d).
    y : Tensor
        Second input of shape (..., m, d).

    Returns
    -------
    Tensor
        Kernel matrix of shape (..., n, m).
    """
    # Project to eigenspace
    x_proj = torch.matmul(x, self.eigenvectors)  # (..., n, r)
    y_proj = torch.matmul(y, self.eigenvectors)  # (..., m, r)

    # Apply eigenvalue weighting
    x_weighted = x_proj * torch.sqrt(torch.abs(self.eigenvalues) + 1e-8)
    y_weighted = y_proj * torch.sqrt(torch.abs(self.eigenvalues) + 1e-8)

    # Compute kernel
    kernel = torch.matmul(x_weighted, y_weighted.transpose(-2, -1))

    if self.normalize:
        # Row normalization
        kernel = F.normalize(kernel, p=2, dim=-1)

    return kernel

extract_features

extract_features(x: Tensor) -> Tensor

Extract spectral features.

Parameters:

Name	Type	Description	Default
x	Tensor	Input of shape (..., n, d).	required

Returns:

Type	Description
Tensor	Spectral features of shape (..., n, r).

Source code in spectrans/kernels/spectral.py

def extract_features(self, x: Tensor) -> Tensor:
    """Extract spectral features.

    Parameters
    ----------
    x : Tensor
        Input of shape (..., n, d).

    Returns
    -------
    Tensor
        Spectral features of shape (..., n, r).
    """
    # Project to eigenspace
    features = torch.matmul(x, self.eigenvectors)

    # Weight by eigenvalues
    features = features * torch.sqrt(torch.abs(self.eigenvalues) + 1e-8)

    return features

forward

forward(x: Tensor, y: Tensor | None = None) -> Tensor

Forward pass for nn.Module compatibility.

Parameters:

Name	Type	Description	Default
x	Tensor	First input of shape (..., n, d).	required
y	Tensor | None	Second input. If None, returns features.	None

Returns:

Type	Description
Tensor	Kernel matrix or features.

Source code in spectrans/kernels/spectral.py

def forward(self, x: Tensor, y: Tensor | None = None) -> Tensor:
    """Forward pass for nn.Module compatibility.

    Parameters
    ----------
    x : Tensor
        First input of shape (..., n, d).
    y : Tensor | None, default=None
        Second input. If None, returns features.

    Returns
    -------
    Tensor
        Kernel matrix or features.
    """
    if y is None:
        return self.extract_features(x)
    else:
        return self.compute(x, y)

orthogonalize_eigenvectors

orthogonalize_eigenvectors() -> None

Orthogonalize eigenvectors via Gram-Schmidt.

Source code in spectrans/kernels/spectral.py

def orthogonalize_eigenvectors(self) -> None:
    """Orthogonalize eigenvectors via Gram-Schmidt."""
    if self.trainable_eigenvectors:
        with torch.no_grad():
            Q, _ = torch.linalg.qr(self.eigenvectors)
            self.eigenvectors.data = Q

FourierKernel

FourierKernel(rank: int, input_dim: int, learnable_filter: bool = True, filter_type: Literal['gaussian', 'butterworth', 'ideal'] = 'gaussian', cutoff_freq: float = 0.5)

Bases: Module, SpectralKernel

Kernel defined in Fourier domain.

Defines kernel through spectral filters in frequency space.

Parameters:

Name	Type	Description	Default
rank	int	Number of Fourier modes.	required
input_dim	int	Input dimension.	required
learnable_filter	bool	Whether filter is learnable.	True
filter_type	Literal['gaussian', 'butterworth', 'ideal']	Type of spectral filter.	"gaussian"
cutoff_freq	float	Normalized cutoff frequency.	0.5

Attributes:

Name	Type	Description
filter	Parameter or Tensor	Spectral filter of shape (rank,).

Methods:

Name	Description
compute	Compute Fourier kernel.

Source code in spectrans/kernels/spectral.py

def __init__(
    self,
    rank: int,
    input_dim: int,
    learnable_filter: bool = True,
    filter_type: Literal["gaussian", "butterworth", "ideal"] = "gaussian",
    cutoff_freq: float = 0.5,
):
    # Use super() to initialize nn.Module (first in MRO)
    super().__init__()
    # Manually set attributes that SpectralKernel.__init__ would set
    self.rank = rank
    self.normalize = True

    self.input_dim = input_dim
    self.filter_type = filter_type
    self.cutoff_freq = cutoff_freq

    # Initialize spectral filter
    filter_vals = self._init_filter()

    if learnable_filter:
        self.filter = nn.Parameter(filter_vals)
    else:
        self.register_buffer("filter", filter_vals)

Functions

compute

compute(x: Tensor, y: Tensor) -> Tensor

Compute Fourier kernel.

Parameters:

Name	Type	Description	Default
x	Tensor	First input of shape (..., n, d).	required
y	Tensor	Second input of shape (..., m, d).	required

Returns:

Type	Description
Tensor	Kernel matrix of shape (..., n, m).

Source code in spectrans/kernels/spectral.py

def compute(self, x: Tensor, y: Tensor) -> Tensor:
    """Compute Fourier kernel.

    Parameters
    ----------
    x : Tensor
        First input of shape (..., n, d).
    y : Tensor
        Second input of shape (..., m, d).

    Returns
    -------
    Tensor
        Kernel matrix of shape (..., n, m).
    """
    # Compute FFT of inputs
    x_freq = safe_rfft(x, dim=-1)
    y_freq = safe_rfft(y, dim=-1)

    # Truncate to rank modes
    x_freq = x_freq[..., : self.rank]
    y_freq = y_freq[..., : self.rank]

    # Apply spectral filter
    x_filtered = x_freq * self.filter
    y_filtered = y_freq * self.filter

    # Compute kernel in frequency domain
    # K(x,y) = Real(IFFT(X_filtered * conj(Y_filtered)))
    kernel_freq = x_filtered.unsqueeze(-2) * y_filtered.unsqueeze(-3).conj()

    # Average over frequency dimension
    kernel: Tensor = kernel_freq.real.mean(dim=-1)

    return kernel

Functions