Implement Hadamard Gaussian noise #2481

TobyBoyne · 2024-02-23T23:29:15Z

Addresses #877, #2416. Implements noise that is homoskedastic in inputs, but task-dependent.

Existing multi-task noises (https://docs.gpytorch.ai/en/latest/likelihoods.html#multi-dimensional-likelihoods) only work where all tasks are observed for each input. In the 'Hadamard' setting, where each input corresponds to exactly one task, there is no existing implementation that supports a different noise for each task.

See the updated Hadamard Multitask GP Regression for usage.

Unit tests, documentation, type hinting will all also be provided soon!

TobyBoyne · 2024-02-23T23:33:54Z

Two areas that likely need improving:

The implementation of _shaped_noise_covar probably needs some work. I'm not sure if there is some linear_operator magic to make this neater?
The syntax can be a little questionable. For example, I don't like the fact that the task indexes need to be in a list in mll(output, full_train_y, [full_train_i]). It is also not obvious that the implementation requires the task indexes to be passed as the last argument to the model, since they are accessed by params[0][-1]. At least some type checking is required here?

Also, I should credit @jpfolch for help with this implementation!

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

Balandat · 2025-03-12T04:24:14Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+        num_of_tasks : number of tasks in the multi output GP
+        noise_prior : any prior you want to put on the noise
+        noise_constraint : constraint to put on the noise


Suggested change

num_of_tasks : number of tasks in the multi output GP

noise_prior : any prior you want to put on the noise

noise_constraint : constraint to put on the noise

num_of_tasks: Number of tasks in the multi-output GP.

noise_prior: Prior for the noise.

noise_constraint: Constraint on the noise value.

batch_shape is missing here

Balandat · 2025-03-12T04:25:44Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+        num_tasks,
+        noise_prior=None,
+        noise_constraint=None,
+        batch_shape=torch.Size(),


can you add type annotations here?

Balandat · 2025-03-12T04:35:05Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+    def __init__(
+        self,
+        num_tasks,
+        noise_prior=None,


So we can only define a single prior for all tasks here - I think that's ok but probably worth mentioning this limitation in the docstring.

Would it not be possible to have a multi-dimensional prior on your task noises? I had assumed that this will work by-default with the current implementation, but I will confirm this behaviour works as intended.

noise_prior=LogNormalPrior( loc=torch.tensor([0., 1.]), scale=torch.tensor([1., 1.]) )

hmm you're right, that should work with a multi-dimensional prior (but not with different types of priors on different tasks)

Balandat · 2025-03-12T04:36:57Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+
+    def _shaped_noise_covar(self, base_shape: torch.Size, *params: Any, **kwargs: Any):
+        # params contains task indexes
+        task_idxs = params[0][-1]


Language nit

Suggested change

task_idxs = params[0][-1]

task_idcs = params[0][-1]

As you said, some runtime typechecking to make sure these are indeed the task indices would probably be good here.

Also, if you can introduce some comments with the shapes of the intermediate tensors / matrices in this function that would be helpful.

Is there a "canonical" way to check that these are the task indices? I will check shape and dtype, but I wasn't sure if there was a way to guarantee that task_idcs contains what we want.

I don't think there is, unfortunately

Balandat · 2025-03-12T04:41:41Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+        *params: Any,
+        **kwargs: Any,
+    ) -> base_distributions.Normal:
+        noise = self._shaped_noise_covar(function_samples.shape, *params, **kwargs).diag()


It's a bit unfortunate that IIUC self._shaped_noise_covar always returns the full covariance but that you're calling diag() on it here - ideally we can have a way to just compute the diagonal here.

From my understanding, if self._shaped_noise_covar returns a LinearOperator, then calling .diagonal should only calculate the diagonal (as in _MultiTaskGaussianLikelihoodBase.forward()). I'm not quite sure if my _shaped_noise_covar function does return a correct LinearOperator that shows this behaviour. I'll look into MaskedLinearOperator as you mentioned below.

I'm not quite sure if my _shaped_noise_covar function does return a correct LinearOperator that shows this behaviour

I was worried since you're using .sum(dim=-3) in self._shaped_noise_covar, will that return a LinearOperator in your case?

I have checked, and yes - the function does return a LinearOperator, since summing a LinearOperator will return another LinearOperator here.

Balandat · 2025-03-12T04:42:30Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+
+    def marginal(self, function_dist: MultivariateNormal, *params: Any, **kwargs: Any) -> MultivariateNormal:
+        mean, covar = function_dist.mean, function_dist.lazy_covariance_matrix
+        noise_covar = self._shaped_noise_covar(mean.shape, *params, **kwargs).squeeze(0)


Why is the .squeeze(0) necessary here? Does it work for general batch shapes?

The squeeze here is due to the shape of MultitaskHomoskedasticNoise.forward. It seems to add an extra dimension than I would expect, as the output is (1, num_tasks, num_data, num_data). I have changed this to .squeeze(-4), and moved it to where it is actually relevant for clarity.

base_shape = torch.Size([40]) HomoskedasticNoise()(shape=base_shape).shape # (40, 40) MultitaskHomoskedasticNoise(num_tasks=2)(shape=base_shape).shape # (1, 2, 40, 40)

Balandat · 2025-03-12T04:48:19Z

gpytorch/likelihoods/hadamard_gaussian_likelihood.py

+        diag = torch.eq(all_tasks, task_idxs.mT)
+        mask = DiagLinearOperator(diag)
+        return (noise_base_covar_matrix @ mask).sum(dim=-3)


Could you use a MaskedLinearOperator here? Not sure if that would necessarily be better (seems like it wouldn't help with the diag comment below b/c of this).

I don't think it can be used here. In my code, diag, which is the diagonal of the mask matrix, has shape (num_tasks, num_data). However, it looks like MaskedLinearOperator would only allow for a mask of shape (num_data,).

This could feasibly be rewritten as a comprehension, summing over num_tasks different MaskedLinearOperator, but I don't think it would help with the readability of the code, which does already return a LinearOperator.

Balandat · 2025-03-12T04:49:57Z

examples/03_Multitask_Exact_GPs/Hadamard_Multitask_GP_Regression.ipynb

@@ -29,7 +29,22 @@
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
-   "outputs": [],
+   "outputs": [


ideally we could avoid emitting this unrelated & unnecessary warning

Balandat · 2025-03-12T04:52:33Z

examples/03_Multitask_Exact_GPs/Hadamard_Multitask_GP_Regression.ipynb

@@ -108,7 +123,7 @@
    "        \n",
    "        return gpytorch.distributions.MultivariateNormal(mean_x, covar)\n",
    "\n",
-    "likelihood = gpytorch.likelihoods.GaussianLikelihood()\n",
+    "likelihood = gpytorch.likelihoods.HadamardGaussianLikelihood(num_tasks=2)\n",


Rather than changing the tutorial to only use task-specific noises, it would be good to instead show both cases and discuss the pros and cons of both approaches in the tutorial. Ideally you could produce the output you shared in pytorch/botorch#2765 (comment)

We could also try to highlight a potential failure mode where in the very low data regime andwith uniform noise levels across tasks, estimating task-specific noises results in worse performance due to the less parsimonious model.

Co-authored-by: Max Balandat <[email protected]>

TobyBoyne · 2025-03-12T10:00:02Z

Thanks for the thorough review, @Balandat! I've left a couple of comments and will update the PR soon to reflect your suggestions.

TobyBoyne · 2025-03-14T15:23:34Z

Hi @Balandat, I've pushed some changes that should address some of the concerns raised:

Added tests to check behaviour in batched settings
Rearranged notebook to discuss pros and cons of task-specific noise in appendix
Added the requested type hints and shape comments

I'm not sure why the readthedocs check is failing - seems to be having trouble importing matplotlib just for my notebook?

Hope you have a lovely weekend! Please let me know if there are any other changes to be made.

Balandat

Thanks, code looks great. The notebook has no outputs included though? Can you re-run this and make sure that it includes the figures so that users can see the plots without having to run them themselves (and so they render on GitHub)?

I'm not sure why the readthedocs check is failing - seems to be having trouble importing matplotlib just for my notebook?

Not sure what this is - could be transient, let's just try again maybe it'll work on another push.

Create task-dependent noise; update example notebook

40452ba

TobyBoyne added 2 commits February 27, 2024 12:06

Cleaner shaped_noise_covar using linear operator

66b0cb6

Merge remote-tracking branch 'upstream/main' into hadamardlikelihood

3631a90

TobyBoyne mentioned this pull request Mar 11, 2025

[FEATURE REQUEST]: Task-specific noise in MultiTaskGPs pytorch/botorch#2765

Open

1 task

Balandat reviewed Mar 12, 2025

View reviewed changes

TobyBoyne and others added 3 commits March 12, 2025 09:56

Update gpytorch/likelihoods/hadamard_gaussian_likelihood.py

29d14f4

Co-authored-by: Max Balandat <[email protected]>

Update gpytorch/likelihoods/hadamard_gaussian_likelihood.py

a3f500d

Co-authored-by: Max Balandat <[email protected]>

Update gpytorch/likelihoods/hadamard_gaussian_likelihood.py

2cffa6b

Co-authored-by: Max Balandat <[email protected]>

TobyBoyne added 4 commits March 13, 2025 22:07

Type hints and shape comments

b1fe818

Update Multitask notebook to compare to uniform noise

64f7b2c

Add tests; more careful shape checking in batch settings

62b14d2

Move task-specific noise to end of notebook; add comparison

90623b0

Add comparison to GaussianLikelihood in tests

1459ada

TobyBoyne requested a review from Balandat March 18, 2025 19:38

Balandat reviewed Mar 21, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement Hadamard Gaussian noise #2481

Implement Hadamard Gaussian noise #2481

TobyBoyne commented Feb 23, 2024

TobyBoyne commented Feb 23, 2024

Balandat Mar 12, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

TobyBoyne Mar 12, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

TobyBoyne Mar 12, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

TobyBoyne Mar 12, 2025

Balandat Mar 12, 2025

TobyBoyne Mar 13, 2025

Balandat Mar 12, 2025

TobyBoyne Mar 13, 2025

Balandat Mar 12, 2025

TobyBoyne Mar 13, 2025

Balandat Mar 12, 2025

Balandat Mar 12, 2025

TobyBoyne commented Mar 12, 2025

TobyBoyne commented Mar 14, 2025 •

edited

Loading

Balandat left a comment

Implement Hadamard Gaussian noise #2481

Are you sure you want to change the base?

Implement Hadamard Gaussian noise #2481

Conversation

TobyBoyne commented Feb 23, 2024

TobyBoyne commented Feb 23, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

TobyBoyne commented Mar 12, 2025

TobyBoyne commented Mar 14, 2025 • edited Loading

Balandat left a comment

Choose a reason for hiding this comment

TobyBoyne commented Mar 14, 2025 •

edited

Loading