Skip to content

RFF

Compute Random Fourier Feature (RFF) kernel approximations.

RFF 

RFF(base_kernel, num_basis_fns=50, frequencies=None, compute_engine=BasisFunctionComputation(), key=jr.PRNGKey(0))

Bases: AbstractKernel

Computes an approximation of the kernel using Random Fourier Features.

All stationary kernels are equivalent to the Fourier transform of a probability distribution. We call the corresponding distribution the spectral density. Using a finite number of basis functions, we can compute the spectral density using a Monte-Carlo approximation. This is done by sampling from the spectral density and computing the Fourier transform of the samples. The kernel is then approximated by the inner product of the Fourier transform of the samples with the Fourier transform of the data.

The key reference for this implementation is the following papers: - 'Random Features for Large-Scale Kernel Machines' by Rahimi and Recht (2008). - 'On the Error of Random Fourier Features' by Sutherland and Schneider (2015).

Parameters:

  • base_kernel (StationaryKernel) –

    The base kernel to be approximated.

  • num_basis_fns (int, default: 50 ) –

    The number of basis functions to use in the approximation.

  • frequencies (Float[Array, 'M D'] | None, default: None ) –

    The frequencies to use in the approximation. If None, the frequencies are sampled from the spectral density of the base kernel.

  • compute_engine (BasisFunctionComputation, default: BasisFunctionComputation() ) –

    The computation engine to use for the basis function computation.

  • key (KeyArray, default: PRNGKey(0) ) –

    The random key to use for sampling the frequencies.

cross_covariance 

cross_covariance(x, y)

Compute the cross-covariance matrix of the kernel.

Parameters:

  • x (Num[Array, 'N D']) –

    the first input matrix of shape (N, D).

  • y (Num[Array, 'M D']) –

    the second input matrix of shape (M, D).

Returns:

  • Float[Array, 'N M'] –

    The cross-covariance matrix of the kernel of shape (N, M).

gram 

gram(x)

Compute the gram matrix of the kernel.

Parameters:

  • x (Num[Array, 'N D']) –

    the input matrix of shape (N, D).

Returns:

  • LinearOperator –

    The gram matrix of the kernel of shape (N, N).

diagonal 

diagonal(x)

Compute the diagonal of the gram matrix of the kernel.

Parameters:

  • x (Num[Array, 'N D']) –

    the input matrix of shape (N, D).

Returns:

  • Float[Array, ' N'] –

    The diagonal of the gram matrix of the kernel of shape (N,).

slice_input 

slice_input(x)

Slice out the relevant columns of the input matrix.

Select the relevant columns of the supplied matrix to be used within the kernel's evaluation.

Parameters:

  • x (Float[Array, '... D']) –

    the matrix or vector that is to be sliced.

Returns:

  • Float[Array, '... Q'] –

    The sliced form of the input matrix.

__add__ 

__add__(other)

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns:

  • AbstractKernel ( AbstractKernel ) –

    A new kernel that is the sum of the two kernels.

__mul__ 

__mul__(other)

Multiply two kernels together.

Parameters:

  • other (AbstractKernel) –

    The kernel to be multiplied with the current kernel.

Returns:

  • AbstractKernel ( AbstractKernel ) –

    A new kernel that is the product of the two kernels.

__call__ 

__call__(x, y)

Superfluous for RFFs.

compute_features 

compute_features(x)

Compute the features for the inputs.

Parameters:

  • x (Float[Array, 'N D']) –

    A \(N \times D\) array of inputs.

Returns:

  • Float[Array, 'N L'] –

    Float[Array, "N L"]: A \(N \times L\) array of features where \(L = 2M\).