Kernel methods in PyTorch

Fast implementations of standard utilities for kernel machines

Installation

pip install -I git+https://github.com/parthe/torchkernels

Requires a PyTorch installation

Stable behavior

Currently this code has been tested with n=10,000 samples.
with Python 3.11 and PyTorch 2.4

Test installation with Laplacian kernel

importtorchfromtorchkernels.kernels.radialimportlaplacian, LaplacianKerneln, p, d=300, 200, 100# number of samples, centers, dimensionsDEV=torch.device("cuda") iftorch.cuda.is_available() elsetorch.device("cpu") X=torch.randn(n, d, device=DEV)
Z=torch.randn(p, d, device=DEV)
kernel_matrix1=laplacian(X, Z, length_scale=1.)
K=LaplacianKernel(length_scale=1.)
kernel_matrix2=K(X, Z)
torch.testing.assert_close(kernel_matrix1, kernel_matrix2, msg='Laplacian test failed')
print('Laplacian test complete!')

Example: Differentiating the kernel function

importtorchfromtorchkernels.kernels.radialimportlaplacianfromtorch.funcimportvmap, grad, jacrevn, p, d, c=300, 200, 100, 3# number of samples, centers, dimensions, outputsDEV=torch.device("cuda") iftorch.cuda.is_available() elsetorch.device("cpu") X=torch.randn(n, d, device=DEV)
Z=torch.randn(p, d, device=DEV)
a=torch.randn(p, device=DEV)
A=torch.randn(p, c, device=DEV)
f=lambdax: laplacian(x, Z, length_scale=1., in_place=False).squeeze() @ aprint(vmap(grad(f))(X).shape)
# torch.Size([300, 100])F=lambdax: laplacian(x, Z, length_scale=1., in_place=False).squeeze() @ Aprint(vmap(jacrev(F))(X).shape)
# torch.Size([300, 3, 100])

Random features

See an example of Logistic regression with random features of the Laplacian kernel.

Currently supported Kernels

• Laplacian, Gaussian, Dispersal (Exponential power kernel)
• Normalized dot-product kernel for arbitrary functions
• Neural Network Gaussian Process (NNGP) and Tangent Kernel (NTK) with ReLU activations

Other utilities

• extracting top eigenvectors of a kernel matrix
• Random feature maps for:
  • Gaussian kernel
  • Laplacian kernel
  • Matern kernel
  • Exponential-power kernel $K(x,z) = \exp(-|x-z|^\gamma)$
• Differentiable models