greta-dev / greta-dev/greta.dynamics

Add functions for implementing continuous relaxations of discrete stochastic transitions

Open
#31 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
R
Stars
6
Forks
2
PR merge metrics
No merged PRs in 30d

Description

### Background

Gradient-based inference (like the HMC greta uses) can only operate on continuous parameter spaces. That means it cannot learn the values of parameters with discrete support (e.g. no unobserved Poisson random variables).

But demographic stochasticity due to discrete stochastic variation in population sizes between timesteps in models of populations and infections (e.g. the number of new infectees is Poisson, the number of individuals surviving is binomial) is often important, especially when populations reach low numbers and near extinction. We cannot directly model the values of these discrete random variables, but we can apply continuous relaxations to _approximate_ them; keeping the state values continuous and replacing the discrete random variables with continuous random variables that matches the mean and variance (and ideally the full shape of the distribution) of the random variable we would like to model as stochastic.

### Continuous approximations to distributions

E.g. a poisson random variable can be approximated with an appropriately-shaped gamma distribution that exactly matches the PMF at discrete values, or by some other distribution that is a a close-enough approximation:

We might write a discrete stochastic growth-rate model like this:

1. $$x_t \sim Poisson(x_{t-1} \times r)$$

where $x$ takes integer values, $Poisson(\lambda)$ is the Poisson distribution, and $r$ is a positive-valued growth rate parameter. To estimate the posterior over the values represented by $x$ in this model, but using HMC, we could instead fit the model:

2. $$y_t \sim \pi(y_{t-1} \times r)$$

where $y_t$ is a (strictly positive) real-valued parameter, and $\pi(\lambda)$ is some probability distribution with support on positive real values that has similar moments (men, variance, skewness, etc.) to $Poisson(\lambda)$.

### Reparameterisation

If we structure these probability distributions such that they can be reparameterised in terms of the parameter and some latent 'innovation' or noise, with known distribution, we can significantly improve the ability to sample these models, since we can decorrelate the posterior distribution in a similar way to the reparameterisation trick for hierarchical models. This can also provide some computational advantages in greta/tensorflow by working with arrays rather than scalars.

If we know the quantile function of the continuous distribution (e.g. $q_{\pi}(p, \lambda)$ as the quantile function of $\pi(\lambda)$ , with $p$ the probability argument), we can reparameterise the innovations as $u \sim U(0, 1)$, and then plug them into the quantile function to sample poisson values. Ie. equation 2 above is equivalent to:

3.
```math
\displaylines{
y_t = q_{\pi}(u_t, \lambda_t) \\
\lambda_t = y_{t-1} \times r \\
u_t \sim U(0, 1)
}
```

The vector of $u$ values can then be computed in advance, and passed into the solvers to be chopped up appropriately. The dependency structure in the model means $y_t$ depends on $y_{t-1}$, and so they are _a priori_ (and therefore also _a posteriori_) correlated. But in this reparameterised formulation, HMC operates instead on $u$, and $u_t$ doesn't depend on $u_{t-1}$ so they are _a priori_ uncorrelated, which removes a lot of correlation in the posterior and makes sampling much easier.

Note that if the quantile function is expensive to compute, other approximations and reparameterisations may be more appealing. E.g. for the Poisson, the inverse of the incomplete gamma function gives the quantile of the gamma distribution whose PDF matches the Poisson PMF at discrete values, but the function has no analytic form and is expensive to compute. A lognormal approximation (with either uniform or standard normal innovations) is imperfect but much more computationally efficient.

Currently this reparameterisation trick is only applicable in the `greta_2` branch with `iterate_dynamic_function()`, since it requires functionality for indexing time-varying parameters.

### Implementation

We just need to provide functions for the relaxations, documentation, and examples of applying this approach. I have some, I just need to add them to the `greta_2` branch and work out the neatest user interface.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.