Skip to content

Fix prediction slowdown from eager kernel evaluation#2749

Open
kayween wants to merge 1 commit into
mainfrom
fix-prediction-slowdown
Open

Fix prediction slowdown from eager kernel evaluation#2749
kayween wants to merge 1 commit into
mainfrom
fix-prediction-slowdown

Conversation

@kayween
Copy link
Copy Markdown
Collaborator

@kayween kayween commented May 1, 2026

This fixes the test-time slowdown for batched inputs as reported in #2736

Root Cause

Note that .evaluate_kernel is called for test covariances in ExactGP._get_test_prior_mean_and_covariances. This would materialize the K(test, test) matrix explicitly (e.g., RBF and Matern kernels), even in cases when only predictive variances are needed.

test_test_covar = joint_covar[..., num_train:, num_train:].evaluate_kernel()
test_train_covar = joint_covar[..., num_train:, :num_train].evaluate_kernel()

The Fix

Remove .evaluate_kernel calls in ExactGP. The prediction strategies should call .evaluated_kernel when needed.

Keep kernel covariances lazy in `_get_test_prior_mean_and_covariances`. Let prediction strategies call `.evaluate_kernel()` only when needed.
@kayween
Copy link
Copy Markdown
Collaborator Author

kayween commented May 1, 2026

This PR moves .evaluate_kernel calls from ExactGP to prediction strategies. In particular, it changes _get_test_prior_mean_and_covariances, which LatentKroneckerGP overrides.

@SebastianAment Would this cause any friction on the BoTorch side?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant