gp(likelihoods-advanced): implement StudentT, multi-class, heteroscedastic, and multi-latent likelihood wrappers
- Dominant language
- Python
- Stars
- 1
- Forks
- 0
- Avg merge
- 19h 2m
- Merged PRs (30d)
- 16
Description
## Problem / Request
`pyrox.gp._likelihoods` already ships `GaussianLikelihood`, `BernoulliLikelihood`, `PoissonLikelihood`, `StudentTLikelihood`, `SoftmaxLikelihood`, and `HeteroscedasticGaussianLikelihood`. The only remaining gap from the original "advanced likelihoods" scope is a **multi-latent likelihood wrapper** — a likelihood whose log-density depends on more than one latent process at the same input (e.g. a Gaussian whose mean and log-variance are both GP-distributed; a multi-class likelihood where each class has its own latent process).
## User Story
> As a user fitting heteroscedastic GP regression with separate latent processes for the mean and log-variance (or a multi-class GP with one latent per class), I want a `MultiLatentLikelihood` wrapper that takes $K$ latent draws $(f_1, \dots, f_K)$ at the same input and a `log_prob(y, fs)` callback, so I can reuse the existing inference paths (`gauss_kl`, `cvi`, `markov` smoothing) without forking them per likelihood.
## Motivation
Concrete consumer paths today:
- **Heteroscedastic regression with separate GP for log-noise** — requires two latents per $x$; the existing `HeteroscedasticGaussianLikelihood` ties noise to a single latent via a fixed link, which is too restrictive.
- **Multi-class GP** — natural shape is $K$ class-specific latents per $x$, currently doable only by wrapping `SoftmaxLikelihood` with manual broadcasting.
- **Bayesian Neural Field with separate uncertainty head** (#73 demo) — outputs two latents (mean + log-variance) over the spatiotemporal grid; the multi-latent wrapper is the right interface.
Without this, every multi-latent model has to bypass the `Likelihood` protocol and open-code its own integration / log-prob path.
## Mathematics
A multi-latent likelihood factorises as
$$
p(y \mid f_1, \dots, f_K) \;=\; \ell(y;\, g(f_1, \dots, f_K)),
$$
where $g$ is a deterministic link from the $K$ latents to the likelihood's natural parameters (e.g. $g(f_1, f_2) = (f_1, \exp f_2)$ for mean/log-variance Gaussian). Expectations against $q(f_1, \dots, f_K)$ that the inference path needs:
$$
\mathbb{E}_q[\log p(y \mid f)] \;=\; \int \ell(y;\, g(f_{1:K}))\, q(f_{1:K})\, df_{1:K},
$$
and its derivative in $(m_k, v_k)$ — the cavity / variational parameters of latent $k$. Under the independent-latents-given-the-input assumption $q(f_{1:K}) = \prod_k q(f_k)$ (the standard VFE / SVGP factorisation), the $K$-dim integral reduces to a tensor product of 1-D cubatures, identical to the existing `Likelihood.expected_log_prob` machinery — just vectorised over $K$.
## Proposed API
```python
class MultiLatentLikelihood(Likelihood):
"""Likelihood depending on K latent processes at the same input.
Args:
K: Number of latent processes.
log_prob_fn: ``(y, fs) -> log_prob`` where ``fs.shape == (K,)``.
link_fn: Optional link from ``(K,)`` latents to likelihood natural params.
integrator: Cubature rule for ``E_q[log p(y | f_{1:K})]``.
"""
K: int # number of latent processes
log_prob_fn: Callable # static — closure
link_fn: Callable | None # static
integrator: AbstractIntegrator # default UnscentedIntegrator(K)
def log_prob(self, y: Array, fs: Float[Array, " K"]) -> Float[Array, ""]: ...
def expected_log_prob(
self,
y: Array,
q_fs: tuple[GaussianState, ...], # one GaussianState per latent
) -> Float[Array, ""]: ...
```
Reuses the existing `Likelihood` protocol — slots into `gauss_kl`, `cvi`, and the markov inference paths without changes.
## Example usage
```python
import jax.numpy as jnp, pyrox.gp as pgp, gaussx
# Heteroscedastic regression with two independent latents:
# f_1 ~ GP_mean(x); f_2 ~ GP_logvar(x)
# y | f_1, f_2 ~ Normal(f_1, exp(f_2)^2)
def log_prob(y, fs):
mean, log_sigma = fs[0], fs[1]
return -0.5 * ((y - mean) / jnp.exp(log_sigma)) ** 2 - log_sigma - 0.5 * jnp.log(2 * jnp.pi)
lik = pgp.MultiLatentLikelihood(
K=2,
log_prob_fn=log_prob,
integrator=gaussx.UnscentedIntegrator(),
)
# Slots into the existing inference paths — same call shape as scalar likelihoods
elbo = pgp.gauss_kl_elbo(model, q_fs=(q_mean, q_logvar), y=y, likelihood=lik)
```
## Tasks / sub-tasks
- [ ] **Spike** — confirm the tensor-product cubature shape works with the existing `gauss_kl` / `cvi` paths; the assumption is that `expected_log_prob` returns a scalar and `q_fs` is a tuple of `GaussianState`s.
- [ ] **`MultiLatentLikelihood`** — extend `src/pyrox/gp/_likelihoods.py`. Default integrator: `UnscentedIntegrator` (3rd-order, 2K+1 points).
- [ ] **Wire to inference paths** — verify `gauss_kl_elbo` / `cvi_recipe` accept a tuple-shaped variational distribution; widen the signature if not.
- [ ] **Tests** — (i) `K=1` case agrees with existing single-latent likelihoods to `1e-6`; (ii) heteroscedastic Gaussian example: ELBO matches the analytical form for a linear log-variance toy; (iii) `jit` / `vmap` / `grad` compatibility.
- [ ] **Docs** — `pyrox.gp.likelihoods` reference entry; cross-link from the heteroscedastic-regression tutorial.
## References
- Saul, A. D., Hensman, J., Vehtari, A. & Lawrence, N. D. (2016). *Chained Gaussian processes*. AISTATS — multi-latent likelihood formulation.
- Existing likelihoods: [`src/pyrox/gp/_likelihoods.py`](https://github.com/jejjohnson/pyrox/blob/main/src/pyrox/gp/_likelihoods.py).
- Existing integrators: `gaussx.UnscentedIntegrator`, `gaussx.GaussHermiteIntegrator`.
## Relationships
- Parent (wave epic): #43
- Blocked by: #
- Blocks: #
- Related: #73
Contributor guide
Assessment
This issue has not been assessed yet.