Related: #38
I come from a modelling background. I'm familiar with Kalman filters but less familiar with the other inference methods here. I'd like some help in the docs. I know cuthbert is not a tool for defining models, but it would be good to have some docs on what I can use where.
If I were organising the docs, I would have some examples index page explaining that inference has two parts: an inner loop (state estimation given fixed $\theta$) and an outer loop (learning $$\theta$$ using the inner loop's log marginal likelihood $$\log p(\mathbf{y}_{1:T} | \theta)$$). cuthbert focuses on the inner loop.
Then, here are the inner loop methods (I think?):
| |
HMM |
LGSSM |
Nonlinear Gaussian |
Generalised |
| Discrete filter |
Exact |
-- |
-- |
-- |
| Kalman filter |
-- |
Exact |
-- |
-- |
| EKF / UKF |
-- |
Overkill |
Approximate |
Poor* |
| Particle filter (SMC) |
Overkill |
Overkill |
Yes |
Yes |
And then eventually a link to an example for each. The way dynamax has a section for each of these models is very helpful. But cuthbert is a better design as it does not bundle the outer loop of parameter estimation in with the inner loop.
Then, we could have some outer loop examples which I expect would pass off to a different library:
-
MLE / MAP: using optax. Can probably be part of the examples
-
EM: E-step runs filter + smoother for expected sufficient statistics; M-step updates $\theta$. Already have an example
-
MCMC/SVI: Use with e.g. blackjax or numpyro. I think this is what @DanWaxman is doing with dynestx, but I need to go through that library.
Aside: outer loop discussion and numpyro integration
cc: @juanitorduz
Not sure where else to put this. I can move this to a separate issue or discussion elsewhere.
Yes, with dynamax you can specify SSMs, but:
- I want to have these SSMs within a numpyro model so I can specify complicated priors or the SSM as part of a larger model
- dynamax does not separate the inner loop and outer loop. e.g.
model.fit_em and model.fit_sgd
- Is dynamax stale?
I think we could make a new, small package to specify models in a dynamax-like way that summarises the inner loop table above and nothing else. The model can then spit out a log-likelihood value to be used in some outer loop (optax, numpyro).
I'm not sure this is at all possible for every case (again, I'm really only familiar with Kalman stuff), but as an end user, it would be great if this new package could
- specify a model (initial distribution, dynamics, observation) like below
- then autogenerate all the
get_*_ functions
- user never has to touch
build_filter and can just do infer(model, data, method="ekf")
Something like
model = new_package.NonlinearGaussianSSM(
init_mean=m0,
init_cov=P0,
dynamics_fn=jnp.sin,
dynamics_cov=Q,
observation_fn=lambda x: x**2,
observation_cov=R,
)
which contains enough information to automatically produce:
- Kalman: not compatible (nonlinear), raise an error or suggest EKF
- EKF (Taylor):
log_density = lambda x_prev, x: logpdf(x, dynamics_fn(x_prev), chol_Q) - you can construct this mechanically from dynamics_fn and dynamics_cov
- UKF (Moments):
mean_and_chol_cov = lambda x: (dynamics_fn(x), chol_Q) - same
- Particle filter:
propagate_sample = lambda key, x, _: dynamics_fn(x) + chol_Q @ normal(key) -- same
The dispatch logic would be full of if/elses.
I think @DanWaxman is working on the numpyro integration (I need to review that code, it might be exactly what I've written above) and if so I think dynestx should cover less complicated models in dynamax's docs like LGSSM/ nonlinear / generalised SSM.
Related: #38
I come from a modelling background. I'm familiar with Kalman filters but less familiar with the other inference methods here. I'd like some help in the docs. I know cuthbert is not a tool for defining models, but it would be good to have some docs on what I can use where.
If I were organising the docs, I would have some examples index page explaining that inference has two parts: an inner loop (state estimation given fixed$\theta$ ) and an outer loop (learning $$\theta$$ using the inner loop's log marginal likelihood $$\log p(\mathbf{y}_{1:T} | \theta)$$ ). cuthbert focuses on the inner loop.
Then, here are the inner loop methods (I think?):
And then eventually a link to an example for each. The way dynamax has a section for each of these models is very helpful. But cuthbert is a better design as it does not bundle the outer loop of parameter estimation in with the inner loop.
Then, we could have some outer loop examples which I expect would pass off to a different library:
Aside: outer loop discussion and numpyro integration
cc: @juanitorduz
Not sure where else to put this. I can move this to a separate issue or discussion elsewhere.
Yes, with dynamax you can specify SSMs, but:
model.fit_emandmodel.fit_sgdI think we could make a new, small package to specify models in a dynamax-like way that summarises the inner loop table above and nothing else. The model can then spit out a log-likelihood value to be used in some outer loop (optax, numpyro).
I'm not sure this is at all possible for every case (again, I'm really only familiar with Kalman stuff), but as an end user, it would be great if this new package could
get_*_functionsbuild_filterand can just doinfer(model, data, method="ekf")Something like
which contains enough information to automatically produce:
log_density = lambda x_prev, x: logpdf(x, dynamics_fn(x_prev), chol_Q)- you can construct this mechanically fromdynamics_fnanddynamics_covmean_and_chol_cov = lambda x: (dynamics_fn(x), chol_Q)- samepropagate_sample = lambda key, x, _: dynamics_fn(x) + chol_Q @ normal(key)-- sameThe dispatch logic would be full of if/elses.
I think @DanWaxman is working on the numpyro integration (I need to review that code, it might be exactly what I've written above) and if so I think dynestx should cover less complicated models in dynamax's docs like LGSSM/ nonlinear / generalised SSM.