bug: Zero mean function doesn't necessarily return zero after optimising a GP posterior

[Thomas-Christie]

Bug Report

GPJax version:

0.6.7

Current behaviour:

Currently, the Zero mean function doesn't necessarily stay as zero once the posterior model has been optimised. This is because, under the hood, it is modelled using the Constant mean function, with the constant field set to 0. However, the constant field is trainable, and so upon inspecting a GP posterior which has had its hyperparameters optimised, the Zero mean function can in fact be non-zero.

Expected behaviour:

I believe that if a user explicitly chooses the Zero mean function then it should always return zero (so, in effect, the mean is not trainable). This is also consistent with existing GP modelling libraries such as GPFlow (https://github.com/GPflow/GPflow/blob/develop/gpflow/functions.py#L187).

Steps to reproduce:

Create a GP model with a Zero mean function prior, and generate an optimised posterior on some data with non-zero mean. Inspecting the posterior model's "prior" mean will reveal that the constant field of the Zero mean function will have changed from 0. This can be seen in the regression notebook.

Feel free to let me know if this behaviour is in fact intended!

[daniel-dodd]

Thanks for spotting this @Thomas-Christie.

Changing param_field's trainable=False, won't fix this issue in general. The constant field really needs to be a static_field as it should never change value under any operations, and we should have an init=False to ensure it does not change value upon initialisation e.g.,

from gpjax.mean_functions import Constant
from gpjax.base import static_field
from dataclasses import dataclass
from jaxtyping import Float, Array
import jax.numpy as jnp

@dataclass
class Zero(Constant):
    constant: Float[Array, "1"] = static_field(jnp.array([0.0]), init=False)

Other options available here too.

[Thomas-Christie]

Thanks for the advice @daniel-dodd , I will work on this.

[trsav]

#339 feels slightly similar to this having just put in the request. Parameter (c) defined as a function of a static field which should remain constant, however is being 'trained' (in this case resulting in an invalid Polar GP).

[Thomas-Christie]

Fixed in #358