MOGPs in GPJax

[daniel-dodd]

To build on the excellent work of #360. @thomaspinder, @henrymoss, @ingmarschuster

Rough proposal

There's no free lunch when it comes to MOGP models. Therefore, it's essential to establish a solid foundation with structured abstractions that empower users to create their own models while ensuring seamless compatibility with other library modules.

We briefly outline the process of refactoring in alignment with GPJax's existing design principles to accomplish this goal. However, it's important to acknowledge that there may exist even more elegant approaches to achieve this objective. Regardless, our priority should be to seamlessly integrate these modifications consistently throughout the entire library.

In my view, it's beneficial to maintain a clear separation between kernel and likelihood computations as distinct entities. This decoupling simplifies the codebase, enhances extensibility, and maximises computational efficiency.

Existing API introduced by the PR:

# Define a kernel for the inputs.
in_kernel = gpx.kernels.RBF()

# Define a kernel for the ouputs.
out_kernel = gpx.kernels.CatKernel(... # args determine num outputs )

# Define mean function.
meanf = gpx.mean_functions.Constant(jnp.array([0.0, 1.0]))

# Define prior.
prior = gpx.Prior(mean_function=meanf, kernel=in_kernel, out_kernel=out_kernel)

# Define likelihood.
likelihood = gpx.Gaussian(num_datapoints = D.n)

# Define posterior.
likelihod * prior

Possible API:

Same as current GPJax.

# Define a kernel for the inputs.
in_kernel = gpx.kernels.RBF()

# Define a kernel for the outputs.
out_kernel = gpx.kernels.CatKernel(... # args determine num outputs )

# Define mean function
meanf = gpx.mean_functions.Constant(jnp.array([0.0, 1.0]))

# Build the Coregionalization Kernel.
kernel = gpx.CorKernel(in_kernel=in_kernel, out_kernel=out_kernel)

# Define prior.
prior = gpx.Prior(mean_function=meanf, kernel=kernel)

# Define likelihood.
likelihood = gpx.Bernoulli(D) # D is passed for abstract shapes.

# Construct posterior:
prior * likelihood

Additionally, this implies that we can introduce various MOGP models by altering both the likelihood and the kernel. For instance, to demonstrate this flexibility, consider we have a scenario with two output dimensions, where one input is binary while the other is continuous:

likelihood = gpx.ConditionallyIndependent(
	[
	gpx.Bernoulli(D), # Maybe annoying to pass D in twice.
	gpx.Gaussian(D),
	],
	active_dims=[0, 1],
)

Alternatively, if we require a different kernel structure, we can simply modify the Kernel. This flexibility allows us to establish efficient structures for different isotropic and heterotropic modelling scenarios.

Differences under the hood from the existing PR’s MOGP code

• Requires just restructuring the existing PR’s code slightly (split up between the kernel and the likelihood) — we don’t really need to change much. Just move a few things around. We would be able to retain full compatibility with the rest of the library and agnostic to the type of MOGP model:
• Really, just by restructuring existing PR’s code slightly - :
  1. We utilise kernel compute engine functionality to compute the covariance and cross covariance of MOGPs. One computation abstraction to handle this.
  2. We slightly alter the likelihood.
  3. Akin to the functionality provided, by ReshapedDistribution we can likely simplify this using tree_maps to extend the GaussianDistribution for computation/reshaping samples from MOGP.
  4. We would need slightly adjust the variational families to ensure the correct shapes are initialised on the sparse GP’s distributional parameters (easily determined by the prior/posterior object). And for the NonConjugate posterior’s latent function values. ~ two lines of code each. As you would have to do for the existing PR’s approach anyway.
• No other changes — critically code within the prior / conjugate posterior is kept minimal, and easy to ready — more didactic, in contrast with the current structure of the existing PR’s code, and we don’t have to start introducing new objects and having to refactor
• And given this modular code design, it is straightforward to develop new MOGP models, adjust parts of the code where desired for research, while the existing PR’s approach is coupled and hardcoded into the prior.

Alternatives:

• There may be better way to setup this infrastructure. The primary objectives should be to maintain coherence and minimalism in order to enhance clarity and simplicity to facilitate extensibility.

[ingmarschuster]

Hi @daniel-dodd, thanks for the proposal. I agree that minimalism/clarity and loose coupling are very important.

I think it's probably a good idea to introduce a gpx.CorKernel – how then would you specify the number of output dimensions, as this will be needed for realizing performance improvements by utilizing the Kronecker inversion formula (via Cola)? What I also don't understand clearly is how you propose to use the kernel compute engine functionality to compute the (cross-) covariance of MOGPs. These two questions are linked I guess.