Develop multi-output GP model with linear mean for each task and with learnable hyperparameters

Questions v5
1

Hello,
I am wondering if it’s possible to develop multi-output GP process to predict two correlated tasks and the mean for each task is a linear function of the input x with learnable hyperparameters.

Thank you.

2

Hello, @DanhPhan recently explored multi-output GPs in a great talk:

and a notebook:

I hope you find them helpful!

3 Likes
3

The “full post” of Danh’s PyMC Web Series event is here. I know, there were multiple posts created in the run-up to the event for marketing purposes and it’s a bit confusing. But hopefully one or more contains information you find useful!

1 Like
4

After examining the above resources. I am still confused. from the paper and the lecture, it seems that the application of multi-output method is based on the assumption that the distribution mean is zero. However, I am wondering if it’s possible to extend its application so that each task has a mean with learnable parameters
Thank you

5

You can def add your own mean function or whatever you’d like there. Something like

mean = pm.gp.mean.Linear(...)
mogp = pm.gp.Marginal(mean_func=mean, cov_func=cov_icm)

(taking code from @daniel-saunders-phil’s link) might be what you need. It’s important to know that the mean function is just a super simple wrapper around any PyMC code, so you can write your own easily. __init__ takes in random variables, and __call__ takes in the same data X the GP is over.

1 Like
6

Thank you @bwengals for your reply. I defined the following mean function

class Linear(pm.gp.mean.Mean):
def init(self, coeffs, intercept):
super().init()
self.b = intercept
self.A = coeffs

def __call__(self, X):
    return pt.squeeze(pt.dot(X, self.A) + self.b)

and then tried to used the function to define the model as the following code

with pm.Model() as model:

# Priors
ell = pm.Gamma("ell", alpha=3, beta=0.2, shape=2)
eta = pm.Gamma("eta", alpha=3, beta=0.2, shape=2)
kernels = [pm.gp.cov.ExpQuad, pm.gp.cov.Matern52]
sigma = pm.HalfNormal("sigma", sigma=3)
beta_m = pm.Gamma("beta_m", alpha=3, beta=0.2, shape=2)
b_m = pm.Gamma("b_m", alpha=3, beta=0.2, shape=2)

# Define a list of covariance functions
cov_list = [
    eta[idx] ** 2 * kernel(input_dim=2, ls=ell[idx], active_dims=[0])
    for idx, kernel in enumerate(kernels)
]

# Get the LCM kernel
cov_lcm = get_lcm(input_dim=2, active_dims=[1], num_outputs=n_outputs, kernels=cov_list)

# define linear mean
mean = Linear(coeffs=beta_m, intercept=b_m)

# Define a Multi-output GP
mogp = pm.gp.Marginal(mean_func=mean, cov_func=cov_lcm)
y_ = mogp.marginal_likelihood("f", x, y, sigma=sigma)

But I got the following error:

ValueError Traceback (most recent call last)
File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:441, in Elemwise.get_output_info(self, dim_shuffle, inputs)
439 try:
440 out_shapes = [
→ 441 [
442 get_most_specialized_shape(shape)
443 for shape in zip([inp.type.shape for inp in inputs])
444 ]
445 ] * shadow.nout
446 except ValueError:

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:442, in (.0)
439 try:
440 out_shapes = [
441 [
→ 442 get_most_specialized_shape(shape)
443 for shape in zip(*[inp.type.shape for inp in inputs])
444 ]
445 ] * shadow.nout
446 except ValueError:

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:434, in Elemwise.get_output_info..get_most_specialized_shape(shapes)
433 if len(shapes) > 1:
→ 434 raise ValueError
435 return tuple(shapes)[0]

ValueError:

During handling of the above exception, another exception occurred:

ValueError Traceback (most recent call last)
Cell In[64], line 25
23 # Define a Multi-output GP
24 mogp = pm.gp.Marginal(mean_func=mean, cov_func=cov_lcm)
—> 25 y_ = mogp.marginal_likelihood(“f”, x, y, sigma=sigma)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pymc\gp\gp.py:473, in Marginal.marginal_likelihood(self, name, X, y, sigma, noise, jitter, is_observed, **kwargs)
470 sigma = _handle_sigma_noise_parameters(sigma=sigma, noise=noise)
472 noise_func = sigma if isinstance(sigma, Covariance) else pm.gp.cov.WhiteNoise(sigma)
→ 473 mu, cov = self._build_marginal_likelihood(X=X, noise_func=noise_func, jitter=jitter)
474 self.X = X
475 self.y = y

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pymc\gp\gp.py:429, in Marginal._build_marginal_likelihood(self, X, noise_func, jitter)
428 def _build_marginal_likelihood(self, X, noise_func, jitter):
→ 429 mu = self.mean_func(X)
430 Kxx = self.cov_func(X)
431 Knx = noise_func(X)

Cell In[63], line 9, in Linear.call(self, X)
8 def call(self, X):
----> 9 return pt.squeeze(pt.dot(X, self.A) + self.b)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\var.py:106, in _tensor_py_operators.add(self, other)
104 def add(self, other):
105 try:
→ 106 return at.math.add(self, other)
107 # We should catch the minimum number of exception here.
108 # Otherwise this will convert error when PyTensor flags
109 # compute_test_value is used
110 # Evidently, we need to catch NotImplementedError
111 # TypeError from as_tensor_variable are caught in Elemwise.make_node
112 # Otherwise TensorVariable * SparseVariable won’t work!
113 except (NotImplementedError, TypeError):
114 # We must return NotImplemented and not an
115 # NotImplementedError or raise an NotImplementedError.
116 # That way python will give a good error message like this
117 # TypeError: unsupported operand type(s) for +: 118 # 'TensorVariable' and 'TensorVariable'

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\graph\op.py:295, in Op.call(self, *inputs, **kwargs)
253 r""“Construct an Apply node using :meth:Op.make_node and return its outputs.
254
255 This method is just a wrapper around :meth:Op.make_node.
(…)
292
293 “””
294 return_list = kwargs.pop(“return_list”, False)
→ 295 node = self.make_node(*inputs, **kwargs)
297 if config.compute_test_value != “off”:
298 compute_test_value(node)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:485, in Elemwise.make_node(self, *inputs)
479 “”"
480 If the inputs have different number of dimensions, their shape
481 is left-completed to the greatest number of dimensions with 1s
482 using DimShuffle.
483 “”"
484 inputs = [as_tensor_variable(i) for i in inputs]
→ 485 out_dtypes, out_shapes, inputs = self.get_output_info(DimShuffle, *inputs)
486 outputs = [
487 TensorType(dtype=dtype, shape=shape)()
488 for dtype, shape in zip(out_dtypes, out_shapes)
489 ]
490 return Apply(self, inputs, outputs)

File ~\anaconda3\envs\pymc_env\Lib\site-packages\pytensor\tensor\elemwise.py:447, in Elemwise.get_output_info(self, dim_shuffle, inputs)
440 out_shapes = [
441 [
442 get_most_specialized_shape(shape)
443 for shape in zip([inp.type.shape for inp in inputs])
444 ]
445 ] * shadow.nout
446 except ValueError:
→ 447 raise ValueError(
448 f"Incompatible Elemwise input shapes {[inp.type.shape for inp in inputs]}"
449 )
451 # inplace_pattern maps output idx → input idx
452 inplace_pattern = self.inplace_pattern

ValueError: Incompatible Elemwise input shapes [(22,), (2,)]

I appreciate any advice to get rid of this error.

Thank you

7

Something’s off with the shapes where it adds pt.dot(X, self.A) and self.b. It’s hard for me to debug exactly without some example data, or knowing what the shapes are you’re going for.

The thing to be aware of here though is that the multiple outputs have been stacked into y as if they were a single output. This might or might not change how you do the mean function.

Otherwise though, your approach will work just fine once the shape issue is sorted out.