Skip to content

Add correct_multidimensional_mixture flag to SpectralMixtureKernel #2753

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
anjawa wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Add correct_multidimensional_mixture flag to SpectralMixtureKernel #2753

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a26905c
correct_multidimensional_mixture flag to SMK kernel
anjawa May 18, 2026
0f1b37b
docs updated
anjawa May 18, 2026
7734bc1
test both versions + compare equality for 1D
anjawa May 18, 2026
f3112a9
Merge branch 'main' into correct-dim-smk
anjawa Jun 9, 2026
88e8e13
Merge branch 'main' into correct-dim-smk
anjawa Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 63 additions & 18 deletions gpytorch/kernels/spectral_mixture_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import logging
import math
import warnings

import torch

Expand Down Expand Up @@ -37,6 +38,11 @@ class SpectralMixtureKernel(Kernel):
a few input dimensions. The ints corresponds to the indices of the dimensions. (Default: `None`.)
:type active_dims: float, optional

:param correct_multidimensional_mixture: If True, use the corrected multidimensional
formula from `Correction to Spectral Mixture (SM) Kernel Derivation for Multidimensional Inputs`_.
If False (default), use the product-of-1D-kernels formula as shown in the original paper
(deprecated for ard_num_dims > 1).
:type correct_multidimensional_mixture: bool, optional
:param mixture_scales_prior: A prior to set on the mixture_scales parameter
:type mixture_scales_prior: ~gpytorch.priors.Prior, optional
:param mixture_scales_constraint: A constraint to set on the mixture_scales parameter
Expand Down Expand Up @@ -69,6 +75,8 @@ class SpectralMixtureKernel(Kernel):

.. _Gaussian Process Kernels for Pattern Discovery and Extrapolation:
https://arxiv.org/pdf/1302.4245.pdf
.. _Correction to Spectral Mixture (SM) Kernel Derivation for Multidimensional Inputs:
https://www.cs.cmu.edu/~andrewgw/typo.pdf
"""

is_stationary = True # kernel is stationary even though it does not have a lengthscale
Expand All @@ -78,6 +86,7 @@ def __init__(
num_mixtures: int | None = None,
ard_num_dims: int | None = 1,
batch_shape: torch.Size | None = torch.Size([]),
correct_multidimensional_mixture: bool = False,
mixture_scales_prior: Prior | None = None,
mixture_scales_constraint: Interval | None = None,
mixture_means_prior: Prior | None = None,
Expand All @@ -94,6 +103,14 @@ def __init__(
# This kernel does not use the default lengthscale
super().__init__(ard_num_dims=ard_num_dims, batch_shape=batch_shape, **kwargs)
self.num_mixtures = num_mixtures
self.correct_multidimensional_mixture = correct_multidimensional_mixture

if not correct_multidimensional_mixture and ard_num_dims > 1:
warnings.warn(
"SpectralMixtureKernel with ard_num_dims > 1 computes the product kernel. "
"To use the correct formula, pass correct_multidimensional_mixture=True.",
DeprecationWarning,
)

if mixture_scales_constraint is None:
mixture_scales_constraint = Positive()
Expand Down Expand Up @@ -328,27 +345,55 @@ def forward(
x1_cos = x1_ * self.mixture_means
x2_cos = x2_ * self.mixture_means

# Create grids
x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, diag=diag, **params)
x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, diag=diag, **params)
if self.correct_multidimensional_mixture:
# Compute the exponential term
exp_term = self.covar_dist(
x1_exp, x2_exp, square_dist=True, diag=diag, last_dim_is_batch=last_dim_is_batch
).mul_(-2 * math.pi**2)

# Compute the exponential and cosine terms
exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi**2)
cos_term = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
res = exp_term.exp_() * cos_term.cos_()
# Compute the cosine term
x1_cos_, x2_cos_ = self._create_input_grid(
x1_cos, x2_cos, diag=diag, last_dim_is_batch=last_dim_is_batch, **params
)
cos_term = (x1_cos_ - x2_cos_).mul(2 * math.pi).sum(-1)

# Sum over mixtures
mixture_weights = self.mixture_weights.view(*self.mixture_weights.shape, 1, 1)
if not diag:
mixture_weights = mixture_weights.unsqueeze(-2)
res = exp_term.exp_() * cos_term.cos_()

res = (res * mixture_weights).sum(-3 if diag else -4)
# Sum over mixtures
if last_dim_is_batch:
# last_dim_is_batch=True: res (... x k x d x n x m)
mixture_weights = self.mixture_weights.view(*self.mixture_weights.shape, 1, 1, 1)
if diag:
mixture_weights = mixture_weights.squeeze(-1)
return (res * mixture_weights).sum(-3 if diag else -4)
else:
mixture_weights = self.mixture_weights.view(*self.mixture_weights.shape, 1, 1)
if diag:
mixture_weights = mixture_weights.squeeze(-1)
return (res * mixture_weights).sum(-2 if diag else -3)

# Product over dimensions
if last_dim_is_batch:
# Put feature-dimension in front of data1/data2 dimensions
res = res.permute(*list(range(0, res.dim() - 3)), -1, -3, -2)
else:
res = res.prod(-1)
# Create grids
x1_exp_, x2_exp_ = self._create_input_grid(x1_exp, x2_exp, diag=diag, **params)
x1_cos_, x2_cos_ = self._create_input_grid(x1_cos, x2_cos, diag=diag, **params)

# Compute the exponential and cosine terms
exp_term = (x1_exp_ - x2_exp_).pow_(2).mul_(-2 * math.pi**2)
cos_term = (x1_cos_ - x2_cos_).mul_(2 * math.pi)
res = exp_term.exp_() * cos_term.cos_()

# Sum over mixtures
mixture_weights = self.mixture_weights.view(*self.mixture_weights.shape, 1, 1)
if not diag:
mixture_weights = mixture_weights.unsqueeze(-2)

res = (res * mixture_weights).sum(-3 if diag else -4)

# Product over dimensions
if last_dim_is_batch:
# Put feature-dimension in front of data1/data2 dimensions
res = res.permute(*list(range(0, res.dim() - 3)), -1, -3, -2)
else:
res = res.prod(-1)

return res
return res
196 changes: 120 additions & 76 deletions test/kernels/test_spectral_mixture_kernel.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,119 +20,163 @@ def create_data_single_batch(self):
def create_data_double_batch(self):
return torch.randn(3, 2, 50, 2)

def create_data_1d_no_batch(self):
return torch.randn(50, 1)

def create_data_1d_single_batch(self):
return torch.randn(2, 50, 1)

def test_active_dims_list(self):
x = self.create_data_no_batch()
kernel = self.create_kernel(num_dims=4, active_dims=[0, 2, 4, 6])
y = torch.randn_like(x[..., 0])
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)

covar_mat = kernel(x).evaluate_kernel().to_dense()
kernel_basic = self.create_kernel(num_dims=4)
kernel_basic.raw_mixture_weights.data = kernel.raw_mixture_weights
kernel_basic.raw_mixture_means.data = kernel.raw_mixture_means
kernel_basic.raw_mixture_scales.data = kernel.raw_mixture_scales
covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().to_dense()

self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
for use_correct_formula in [False, True]:
kernel = self.create_kernel(
num_dims=4, active_dims=[0, 2, 4, 6], correct_multidimensional_mixture=use_correct_formula
)
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)

covar_mat = kernel(x).evaluate_kernel().to_dense()
kernel_basic = self.create_kernel(num_dims=4, correct_multidimensional_mixture=use_correct_formula)
kernel_basic.raw_mixture_weights.data = kernel.raw_mixture_weights
kernel_basic.raw_mixture_means.data = kernel.raw_mixture_means
kernel_basic.raw_mixture_scales.data = kernel.raw_mixture_scales
covar_mat_actual = kernel_basic(x[:, [0, 2, 4, 6]]).evaluate_kernel().to_dense()

self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)

def test_active_dims_range(self):
active_dims = list(range(3, 9))
x = self.create_data_no_batch()
kernel = self.create_kernel(num_dims=6, active_dims=active_dims)
y = torch.randn_like(x[..., 0])
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)

covar_mat = kernel(x).evaluate_kernel().to_dense()
kernel_basic = self.create_kernel(num_dims=6)
kernel_basic.raw_mixture_weights.data = kernel.raw_mixture_weights
kernel_basic.raw_mixture_means.data = kernel.raw_mixture_means
kernel_basic.raw_mixture_scales.data = kernel.raw_mixture_scales
covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().to_dense()

self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)
for use_correct_formula in [False, True]:
kernel = self.create_kernel(
num_dims=6, active_dims=active_dims, correct_multidimensional_mixture=use_correct_formula
)
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)

covar_mat = kernel(x).evaluate_kernel().to_dense()
kernel_basic = self.create_kernel(num_dims=6, correct_multidimensional_mixture=use_correct_formula)
kernel_basic.raw_mixture_weights.data = kernel.raw_mixture_weights
kernel_basic.raw_mixture_means.data = kernel.raw_mixture_means
kernel_basic.raw_mixture_scales.data = kernel.raw_mixture_scales
covar_mat_actual = kernel_basic(x[:, active_dims]).evaluate_kernel().to_dense()

self.assertLess(torch.norm(covar_mat - covar_mat_actual) / covar_mat_actual.norm(), 1e-4)

def test_no_batch_kernel_single_batch_x(self):
x = self.create_data_single_batch()
kernel = self.create_kernel(num_dims=x.size(-1))
y = torch.randn_like(x[..., 0])
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)
batch_covar_mat = kernel(x).evaluate_kernel().to_dense()
for use_correct_formula in [False, True]:
kernel = self.create_kernel(num_dims=x.size(-1), correct_multidimensional_mixture=use_correct_formula)
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)
batch_covar_mat = kernel(x).evaluate_kernel().to_dense()

actual_mat_1 = kernel(x[0]).evaluate_kernel().to_dense()
actual_mat_2 = kernel(x[1]).evaluate_kernel().to_dense()
actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
actual_mat_1 = kernel(x[0]).evaluate_kernel().to_dense()
actual_mat_2 = kernel(x[1]).evaluate_kernel().to_dense()
actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])

self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)

# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)

def test_single_batch_kernel_single_batch_x(self):
x = self.create_data_single_batch()
kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([]))
y = torch.randn_like(x[..., 0])
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)
batch_covar_mat = kernel(x).evaluate_kernel().to_dense()
for use_correct_formula in [False, True]:
kernel = self.create_kernel(
num_dims=x.size(-1), batch_shape=torch.Size([]), correct_multidimensional_mixture=use_correct_formula
)
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)
batch_covar_mat = kernel(x).evaluate_kernel().to_dense()

actual_mat_1 = kernel(x[0]).evaluate_kernel().to_dense()
actual_mat_2 = kernel(x[1]).evaluate_kernel().to_dense()
actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])
actual_mat_1 = kernel(x[0]).evaluate_kernel().to_dense()
actual_mat_2 = kernel(x[1]).evaluate_kernel().to_dense()
actual_covar_mat = torch.cat([actual_mat_1.unsqueeze(0), actual_mat_2.unsqueeze(0)])

self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)
self.assertLess(torch.norm(batch_covar_mat - actual_covar_mat) / actual_covar_mat.norm(), 1e-4)

# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = actual_covar_mat.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)

def test_smoke_double_batch_kernel_double_batch_x(self):
x = self.create_data_double_batch()
kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([3, 2]))
y = torch.randn_like(x[..., 0])
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)

batch_covar_mat = kernel(x).evaluate_kernel().to_dense()
kernel_diag = kernel(x, diag=True)
for use_correct_formula in [False, True]:
kernel = self.create_kernel(
num_dims=x.size(-1),
batch_shape=torch.Size([3, 2]),
correct_multidimensional_mixture=use_correct_formula,
)
kernel.initialize_from_data(x, y)
kernel.initialize_from_data_empspect(x, y)

batch_covar_mat = kernel(x).evaluate_kernel().to_dense()
kernel_diag = kernel(x, diag=True)
return batch_covar_mat, kernel_diag

def test_kernel_getitem_single_batch(self):
x = self.create_data_single_batch()
kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([2]))
for use_correct_formula in [False, True]:
kernel = self.create_kernel(
num_dims=x.size(-1), batch_shape=torch.Size([2]), correct_multidimensional_mixture=use_correct_formula
)

res1 = kernel(x).to_dense()[0] # Result of first kernel on first batch of data
res1 = kernel(x).to_dense()[0] # Result of first kernel on first batch of data

new_kernel = kernel[0]
res2 = new_kernel(x[0]).to_dense() # Should also be result of first kernel on first batch of data.
new_kernel = kernel[0]
res2 = new_kernel(x[0]).to_dense() # Should also be result of first kernel on first batch of data.

self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)
self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)

# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = res1.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = res1.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)

def test_kernel_getitem_double_batch(self):
x = self.create_data_double_batch()
kernel = self.create_kernel(num_dims=x.size(-1), batch_shape=torch.Size([3, 2]))

res1 = kernel(x).to_dense()[0, 1] # Result of first kernel on first batch of data

new_kernel = kernel[0, 1]
res2 = new_kernel(x[0, 1]).to_dense() # Should also be result of first kernel on first batch of data.

self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)

# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = res1.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)
for use_correct_formula in [False, True]:
kernel = self.create_kernel(
num_dims=x.size(-1),
batch_shape=torch.Size([3, 2]),
correct_multidimensional_mixture=use_correct_formula,
)

res1 = kernel(x).to_dense()[0, 1] # Result of first kernel on first batch of data

new_kernel = kernel[0, 1]
res2 = new_kernel(x[0, 1]).to_dense() # Should also be result of first kernel on first batch of data.

self.assertLess(torch.norm(res1 - res2) / res1.norm(), 1e-4)

# Test diagonal
kernel_diag = kernel(x, diag=True)
actual_diag = res1.diagonal(dim1=-1, dim2=-2)
self.assertLess(torch.norm(kernel_diag - actual_diag) / actual_diag.norm(), 1e-4)

def test_correct_multidimensional_mixture_1d_equivalence(self):
# For 1D inputs, correct=True and correct=False must give identical results
kernel = self.create_kernel(num_dims=1, correct_multidimensional_mixture=False)
kernel_correct = self.create_kernel(num_dims=1, correct_multidimensional_mixture=True)

x = self.create_data_1d_no_batch()
res = kernel(x).to_dense()
res_correct = kernel_correct(x).to_dense()
self.assertLess(torch.norm(res - res_correct) / res.norm(), 1e-4)

x_batch = self.create_data_1d_single_batch()
res_batch = kernel(x_batch).to_dense()
res_correct_batch = kernel_correct(x_batch).to_dense()
self.assertLess(torch.norm(res_batch - res_correct_batch) / res_batch.norm(), 1e-4)


if __name__ == "__main__":
Expand Down
Loading