Implement Laplace (quadratic) approximation #345

carsten-j · 2024-06-01T05:22:10Z

This is an early version of q quadratic approximation implementation that I have developed while reading Statistical Rethinking by Richard McElreath.

There is a short discussion about this in the issue and maybe @theorashid can help with feedback of this draft PR.

This work is partly based on the Python package pymc3-quap but pymc3-quap is based on PYMC3 and a lot happend bewteen version 3 and 5 of PYMC. Optimizers works better when provided with a good initial guess and hence a (optional) starting point has been added to function arguments. Please see Github for a discussion about the differences between PYMC version 3 and 5 for computing the Hessian.

for more information, see https://pre-commit.ci

carsten-j · 2024-06-02T11:04:39Z

I am looking for the best way to return not just a posterior sample distribution but also the mean vector and covariance matrix of the Gaussian distribution. Any suggestion for this. So far my only idea is to add another section to the inferenceData returned containing this information. Thoughts on this?

zaxtax · 2024-06-02T13:25:22Z

I'm not sure the InferenceData is the best place to put it. We should copy whatever we do with Variational Inference

…

On Sun, 2 Jun 2024, 13:05 Carsten Jørgensen, ***@***.***> wrote: I am looking for the best way to return not just a posterior sample distribution but also the mean vector and covariance matrix of the Gaussian distribution. Any suggestion for this. So far my only idea is to add another section to the inferenceData returned containing this information. Thoughts on this? — Reply to this email directly, view it on GitHub <#345 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUM45EIPXABROOJPVCTZFL353AVCNFSM6AAAAABIT3F2F2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNBTHAYDGOBQGA> . You are receiving this because you are subscribed to this thread.Message ID: ***@***.***>

aloctavodia · 2024-06-02T14:24:15Z

Historically, Inferencedata has been focused on mcmc. But we have discussed a few times extend it to better handle other inference methods, like SMC or variational methods. It just that there has not been enough momentum to agree and implement and schema that works for those methods.

carsten-j · 2024-06-02T17:04:09Z

@zaxtax and @aloctavodia are you saying that I should not return inferencedata at all or just not return the gaussian mean and covariance in the inferencedata object? I am new to both PYMC and Bayesian statistics so I do not know the history of this package.
Best, Carsten

zaxtax · 2024-06-02T21:07:45Z

Oh, it's more that we haven't decided how to handle this within the library. Don't treat this as a blocker, though we should raise it for discussion more broadly

…

On Sun, 2 Jun 2024, 19:04 Carsten Jørgensen, ***@***.***> wrote: @zaxtax <https://github.com/zaxtax> and @aloctavodia <https://github.com/aloctavodia> are you saying that I should not return inferencedata at all or just not return the gaussian mean and covariance in the inferencedata object? I am new to both PYMC and Bayesian statistics so I do not know the history of this package. Best, Carsten — Reply to this email directly, view it on GitHub <#345 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAACCUKGG7KIW4SVLMG6ZLLZFNGB7AVCNFSM6AAAAABIT3F2F2VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDCNBTHE2DINBZGU> . You are receiving this because you were mentioned.Message ID: ***@***.***>

twiecki · 2024-06-03T12:32:57Z

CC @ferrine

aloctavodia · 2024-06-03T17:08:40Z

Oh, it's more that we haven't decided how to handle this within the library. Don't treat this as a blocker, though we should raise it for discussion more broadly

exactly, just saying that if necessary InferenceData can be extended.

ricardoV94 · 2024-06-05T08:59:23Z

Suggestion, include two groups in the returned inferencedata:

A fit group that includes the mean and covariance of the laplace fit, and
a posterior group that includes draws from this fit, and has all the bells and whistles like dimensions, deterministics, etc... Include the default extra groups like observed and constant data. This will look just like a fit from mcmc sampling. This can be disabled by the user by setting draws = 0

We could even try different fits from distinct initialization points (optionally) and save those as distinct "chains" in the fit and corresponding posterior groups. Although usually multiple initialization are used with the goal of finding the best fit, they could still be useful to detect multi-modality / pervasiveness of local optima.

pymc_experimental/inference/quadratic.py

ricardoV94 · 2024-06-05T09:10:15Z

@carsten-j PR looks great! I left some comment above

pymc_experimental/tests/test_quadratic.py

carsten-j · 2024-06-07T13:56:15Z

Thanks you @ricardoV94 and @twiecki for the review comments. I believe that all of them expect one has been fixed. I have not figured out how to use remove_value_transforms. I tried to browse through PYMC source code but that did not really help.

pymc_experimental/inference/laplace.py

ricardoV94 · 2024-06-10T14:52:39Z

Thanks you @ricardoV94 and @twiecki for the review comments. I believe that all of them expect one has been fixed. I have not figured out how to use remove_value_transforms. I tried to browse through PYMC source code but that did not really help.

The docs contains code example: https://www.pymc.io/projects/docs/en/stable/api/model/generated/pymc.model.transform.conditioning.remove_value_transforms.html

carsten-j · 2024-06-10T14:58:59Z

I should have mentioned that I did read the doc and looked at the example. But I have not been able to figure out how to apply it to my case. I will try again ...

ricardoV94 · 2024-06-10T16:04:31Z

To be able to use it inside the model context, it will need this change to get merged first: pymc-devs/pymc#7352

But you should be able to already test by doing the object way with pm.fit(..., model=model) outside of the model context

carsten-j · 2024-06-11T15:12:28Z

@ricardoV94 I figured out how to replace the for loop with remove_value_transforms. Is the PR ready for merge or are there additional review comments?

zaxtax · 2024-06-15T19:58:29Z

Looks good. Once the tests pass, I think it's good to merge

pymc_experimental/inference/laplace.py

theorashid · 2024-06-17T14:16:03Z

pymc_experimental/tests/test_laplace.py

+        logsigma = pm.Uniform("logsigma", 1, 100)
+        mu = pm.Uniform("mu", -10000, 10000)
+        yobs = pm.Normal("y", mu=mu, sigma=pm.math.exp(logsigma), observed=y)
+        vars = [mu, logsigma]


Question: say you only did vars=[mu], how would the variable logsigma be estimated?

I think find_MAP in that case uses the initial_point for the excluded variable(s). I never found that behavior useful tbh

Edit: Maybe it's fine. Either way it's documented here: https://github.com/pymc-devs/pymc/blob/05b557f6460a10c29c3db33690ee535f5b1ecde0/pymc/tuning/starting.py#L73-L75

Sounds like we may want to pass a similar start kwarg to laplace to set the value of variables that are not being optimized?

worth adding a test on this to confirm the behaviour

I am not sure I fully understand this. I will give it a second go with the documentation for find_MAP.

Hi Carsten, is there anything we can do to help get this over the line?

Hi @theorashid. I am not sure how to handle if only a subset of the variables are passed in, e.g. vars=[mu] and log_sigma is left out. If this should raise a warning I need some way of figuring out the number of model parameters and compare that with the number of parameters in vars. I am not sure how to determine the number of model parameters

model.free_RVs

@theorashid and @ricardoV94, I committed an update that will raise a warning in case number of variables in vars does not equal number of model variables.

pymc_experimental/inference/laplace.py

for more information, see https://pre-commit.ci

pymc_experimental/inference/laplace.py

for more information, see https://pre-commit.ci

…model variables

ricardoV94 · 2024-06-29T07:47:27Z

@carsten-j tests are no longer failing in main. You can rebase/merge into your branch

* Allow forward sampling of statespace models in JAX mode Explicitly set data shape to avoid broadcasting error Better handling of measurement error dims in `SARIMAX` models Freeze auxiliary models before forward sampling Bugfixes for posterior predictive sampling helpers Allow specification of time dimension name when registering data Save info about exogenous data for post-estimation tasks Restore `_exog_data_info` member variable Be more consistent with the names of filter outputs * Adjust test suite to reflect API changes Modify structural tests to accommodate deterministic models Save kalman filter outputs to idata for statespace tests Remove test related to `add_exogenous` Adjust structural module tests * Add JAX test suite * Bug-fixes and changes to statespace distributions Remove tests related to the `add_exogenous` method Add dummy `MvNormalSVDRV` for forward jax sampling with `method="SVD"` Dynamically generate `LinearGaussianStateSpaceRV` signature from inputs Add signature and simple test for `SequenceMvNormal` * Re-run example notebooks * Add helper function to sample prior/posterior statespace matrices * fix tests * Wrap jax MvNormal rewrite in try/except block * Don't use `action` keyword in `catch_warnings` * Skip JAX test if `numpyro` is not installed * Handle batch dims on `SequenceMvNormal` * Remove unused batch_dim logic in SequenceMvNormal * Restore `get_support_shape_1d` import

review-notebook-app · 2024-06-30T19:33:54Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

carsten-j · 2024-06-30T19:37:04Z

@ricardoV94 I have rebased the laplace branch but it looks like someone needs to approve Github worksflows.

zaxtax · 2024-07-01T12:56:31Z

Looks like there are still a few failing tests, but once those pass this is probably good to merge

carsten-j · 2024-07-01T14:54:54Z

@zaxtax failing test has been fixed. Can you approve the waiting workflow?

carsten-j · 2024-07-01T15:55:02Z

@zaxtax, all tests passed. Are you also able to merge the PR? Thanks.

twiecki · 2024-07-01T15:56:20Z

Congrats @carsten-j, this is a big one!

carsten-j · 2024-07-01T16:45:13Z

Thank you @twiecki. Really happy to contribute and thanks to all those that helped. After the summer I will try to work on documentation for building and running locally. I took me some time to figure out how this works!

zaxtax · 2024-07-01T20:36:50Z

Congrats @carsten-j this is really neat!

theorashid · 2024-07-02T08:34:42Z

Brilliant work @carsten-j . Hope to see you contribute to PyMC again!

carsten-j and others added 2 commits May 31, 2024 21:28

First draft of quadratic approximation

56abf7d

[pre-commit.ci] auto fixes from pre-commit.com hooks

8d3f0a1

for more information, see https://pre-commit.ci

ricardoV94 reviewed Jun 5, 2024

View reviewed changes