Variance function for the prior

[frazane]

Hi all! I recently started to use the library and I am very happy so far. Thanks a lot to all contributors 😄 I don't have a strong mathematical background so forgive me if I say something idiotic.

I have a problem where I want to use a neural network as mean_function in the Prior. This was fairly easy to implement: I subclassed the AbstractMeanFunction and wrapped a hk.Module the same way it is done in the Deep Kernel Learning example.
Now, my problem is that I would also expect the variance to be different for each input sample, but I see that by default the variance in the RBF kernel (which I'm using) is a single parameter. So, if I understand correctly, all values along the diagonal will take the same constant value.

Is there a way to have the variance depend on the input values of each sample, instead of being a single adjustable parameter?
How would I then combine the computed variance with the Kernel function?

Might be related to #113?

[daniel-dodd]

Hi @frazane,

Thank you for your kind words. :)

This sounds super cool!

Yeah by default, all dimensions use a single parameter. To override this, currently you can use active_dims and pass through the sequence for which you would like the dims to act on.

import jaxkern as jk

 rbf = jk.RBF(active_dims=[0,1,2,4])

print(rbf.init_params(123))
#{'lengthscale': Array([1., 1., 1., 1.], dtype=float32), 'variance': Array([1.], dtype=float32)}

# The lengthscale parameters correspond to dims 0, 1, 2, 4 of your inputs.

[frazane]

Hi @frazane,

Thank you for your kind words. :)

This sounds super cool!

Yeah by default, all dimensions use a single parameter. To override this, currently you can use active_dims and pass through the sequence for which you would like the dims to act on.

import jaxkern as jk



 rbf = jk.RBF(active_dims=[0,1,2,4])



print(rbf.init_params(123))

#{'lengthscale': Array([1., 1., 1., 1.], dtype=float32), 'variance': Array([1.], dtype=float32)}



# The lengthscale parameters correspond to dims 0, 1, 2, 4 of your inputs.

Hi Daniel, thanks for answering! However, I wasn't referring to the lengthscale parameter of the Kernel, but the variance parameter. I'd like to have multiple values (one for each point, calculated from the input values) for the variance, instead of a single one.

[daniel-dodd]

Hi @frazane,

Sorry that I interpreted your question incorrectly!

I missed this in your original question:

Might be related to #113?

Are we talking about hetroscedastic noise here:

$$ y_i \sim N(f(x_i), \sigma_i^2) = N(0, k(x_i, x_i) + \sigma_i^2) , \qquad f \sim \mathcal{GP}(0, k) $$

i.e., have a different noise variance $\sigma_i^2$ for each data point?

As the RBF kernel variance $\alpha^2$ does not mean noise variance, it corresponds to the kernel:
$$k(x, x') = \alpha^2 \exp\left(-\frac{\lVert x - x' \rVert_2^2}{2 \ell^2}\right).$$

[daniel-dodd]

If it is something else, please may you express the maths to me, and we can go from there. :)

[frazane]

Thanks for taking the time @daniel-dodd I really appreciate it. I don't think that the noise variance is what I'm looking for. I'll try to express my problem a bit more precisely.

Starting from the definition of the GP:

We write a GP $f(\cdot) \sim \mathcal{GP}(\mu(\cdot), k(\cdot, \cdot))$ with mean function $\mu: \mathcal{X} \rightarrow \mathbb{R}$ and $\boldsymbol{\theta}$-parameterised kernel $k: \mathcal{X} \times \mathcal{X}\rightarrow \mathbb{R}$.

Then the $(i, j)^{\text{th}}$ entry of the (Gram) $\mathbf{K}_{ff}$ matrix is given by:

$$ [\mathbf{K}{ff}]{i, j} = k(\mathbf{x}_i, \mathbf{x}_j) $$

and the diagonal elements on the resulting covariance matrix, if we use the definition of $k(x,x')$ that you provided, are entirely defined by the variance parameter of the kernel ($\alpha^2$ in your definition or $\sigma^2$ more commonly). In other words, for the prior, the variance for the marginal distribution of each point is the same.

In my problem, for the prior I have $\mu(\cdot)$ defined by a neural network. But my prior knowledge goes beyond estimating the mean: my input $\mathbf{x}$ actually informs about the uncertainty of the marginals as well. So I'd like that instead of having a single parameter $\sigma$ to define the variance of the marginal distributions of the prior, I have it estimated from $\mathbf{x}$. In practice, the diagonal elements of the resulting Multivariate Normal for the prior will be calculated from $\mathbf{x}$.