Add `find_MAP` with close JAX integration and fix bug with Laplace fit #385

jessegrabowski · 2024-10-27T06:29:30Z

Closes #376

This PR adds code to run find_MAP using JAX. I'm using JAX for gradients, because I found the compile times were faster. Open to suggestions/rebuke.

It also adds a fit_laplace function, which is bad because we already have a fit_laplace function. This one has slightly different objective though -- it isn't meant to be used as a step sampler on a subset of model variables. Instead, it is meant to be used on the MAP result to give an approximation to the full posterior. My function also lets you do the Laplace approximation in the transformed space, then do sample-wise reverse transformation. I think this is legit, and lets you obtain approximate posteriors that respect the domain of the prior. Tagging @theorashid so we can resolve the differences.

Last point is that I added a dependency on better_optimize. This is a package I wrote that basically rips out the wrapper code used in PyMC find_MAP and applies it to arbitrary optimization problems. It is more feature complete than the PyMC wrapper -- it supports all optimizer modes for scipy.optimize.minimize and scipy.optimize.root, and also helps get keywords to the right place in those functions (who can ever remember if an argument goes in method_kwargs or in the funciton itself?). I plan to add support for basinhopping as well, which will be nice for really hairy minimizations.

I could see an objection to adding another dependency, but 1) it's a lightweight wrapper around functionality that doesn't really belong in PyMC anyway, and 2) it's a big value-add compared to working directly with the scipy.optimize functions, which have gnarly, inconsistent signatures.

theorashid · 2024-10-28T05:40:31Z

Hey, nice one, yeah I agree, we should only have one fit_laplace function.

it isn't meant to be used as a step sampler on a subset of model variables

The current fit_laplace isn't either. It isn't a step sampler. (The INLA stuff #340 still has a few blockers so that's separate and not yet in the library.) The implementation was made by a user following Statistical Rethinking, where McElreath fits some models using the Laplace approximation of all parameters.

Current behaviour of fit_laplace is:

Take the MAP estimate of the parameters (https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L109)
Move to untransformed space (https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L112)
Get the Hessian (https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L114)
cov = np.linalg.inv(hessian) https://github.com/pymc-devs/pymc-experimental/blob/main/pymc_experimental/inference/laplace.py#L119
simulate from this posterior approx rng.multivariate_normal(mean, cov, size=(1, draws)) and turn into idata. (Looking at this now actually, we should probably transform back into the constrained space? Otherwise I can see e.g. standard deviation variables with samples below zero.)

The behaviour when you only pass a subset of variables isn't really desirable in my opinion (see #345 (comment)), so we put a warning. So as you say:

Instead, it is meant to be used on the MAP result to give an approximation to the full posterior.

Agree, that's the best plan for fit_laplace.

Judging by your docs and a quick glance at your code, I think you're basically doing the same thing. The current implementation is few lines of code and a few docs, so I reckon

make sure you can pass the test case with your method, which is an example from BDA3 https://github.com/pymc-devs/pymc-experimental/blob/main/tests/test_laplace.py
throw any of the useful code and docs into your method

Then it should be safe to delete the existing code and we can go back to one fit_laplace.

I could see an objection to adding another dependency

I would love a generic optimiser in p u r e pytensor, but I can see looking at your code that there a lot of fancy extras that would take a large effort to write in pytensor. Still, if we want to go back to one of our efforts with a fixed point operator (pymc-devs/pytensor#978 and pymc-devs/pytensor#944), we could probably write find_MAP with that in some form, with fewer bells and whistles though.

Happy to look at your code and review properly later in the week if you'd like me to. Let me know. Otherwise, I'll leave to the core devs.

pymc_experimental/inference/jax_find_map.py

ricardoV94 · 2024-10-31T10:52:38Z

Happy to look at your code and review properly later in the week if you'd like me to. Let me know. Otherwise, I'll leave to the core devs.

That would be appreciated

ricardoV94 · 2024-10-31T10:55:33Z

Agree with what @theorashid said. This fit_laplace is going for the same goal as the previous one. Happy to replace it, if it's not married to JAX backend. Still fine to allow using JAX for the autodiff. What you're offering is very similar to nutpie gradient_backend kwarg, so we could use the same terminology

ricardoV94 · 2024-10-31T10:55:53Z

No objections about your custom library wrapper

jessegrabowski · 2024-12-03T18:47:40Z

tagging @theorashid -- I couldn't pick you as a reviewer?

I did a major refactor of this. I broke the marriage to jax and generalized the find_MAP function. Files have been renamed to reflect this. I also merged the two laplace approaches. The biggest change is that I removed the ability to choose vars. I think the idea here was to be able to partially marginalize some variables in a model? But I think this would require a somewhat different approach.

theorashid · 2024-12-04T01:05:18Z

yea sorry I'm just a normal, but I'll give it a review. Will do it at some point in the next 2 weeks.

pymc_experimental/inference/find_map.py

pymc_experimental/inference/laplace.py

ricardoV94

Minor suggestions, PR looks amazing!

ricardoV94 · 2024-12-04T10:54:28Z

@jessegrabowski can we close #376 with this PR?

Do you have a test that covers something like it?

- Rename function `laplace` -> `sample_laplace_posterior`

…mization

jessegrabowski · 2024-12-04T11:11:15Z

Yes, I think this test and this test should cover that issue

pymc_experimental/inference/find_map.py

theorashid · 2024-12-04T16:23:33Z

sweet, all done?

jessegrabowski · 2024-12-04T16:25:47Z

For now, though I'd still appreciate it if you could have a look and open issues on any bugs/shortcomings you find

theorashid

I managed to follow the code through and it looks good to me. Happy you got rid of the option to fit on a subset of variables, which didn't make sense to me anyway. If it passes the original test then it should be good. You can do something about the other comments if you want, but maybe not because we are e x p e r i m e n t a l

theorashid · 2024-12-09T03:01:49Z

pymc_experimental/inference/find_map.py

+    return f_loss_and_grad, f_hess, f_hessp
+
+
+def _compile_functions(


NIT: Maybe _compile_functions and _compile_jax_gradients are slightly too generic function names. I found it a little tricky to remember exactly what they were doing when reading through the code

theorashid · 2024-12-09T03:33:05Z

pymc_experimental/inference/find_map.py

+    use_hess = use_hess if use_hess is not None else method_info["uses_hess"]
+    use_hessp = use_hessp if use_hessp is not None else method_info["uses_hessp"]
+
+    if use_hess and use_hessp:


I was going through all the methods thinking when you would need hess and hessp and then came back to this. I would probably warn the user / not let them pass both use_hess and use_hessp

theorashid · 2024-12-09T03:36:21Z

pymc_experimental/inference/laplace.py

+    return idata
+
+
+def fit_mvn_to_MAP(


fit_mvn_at_MAP? I mean technically this function just fits a MVN at a point, the user doesn't necessarily have to pass the MAP

theorashid · 2024-12-09T03:39:56Z

pymc_experimental/inference/laplace.py

+        H_inv = get_nearest_psd(H_inv)
+        if on_bad_cov == "warn":
+            _log.warning(
+                "Inverse Hessian is not positive semi-definite at the provided point, using the closest PSD "


For my understanding, what sort of scenarios/models would get a not PSD hessian. And is using closest PSD a good ideas?

theorashid · 2024-12-09T03:40:36Z

pymc_experimental/inference/laplace.py

+
+    Parameters
+    ----------
+    mu


add docs here

theorashid · 2024-12-09T03:48:58Z

pymc_experimental/inference/laplace.py

+        and 1).
+
+        .. warning::
+            This argumnet should be considered highly experimental. It has not been verified if this method produces


theorashid · 2024-12-09T03:51:43Z

pymc_experimental/inference/laplace.py

+    gradient_backend: str, default "pytensor"
+        The backend to use for gradient computations. Must be one of "pytensor" or "jax".
+    chains: int, default: 2
+        The number of sampling chains running in parallel.


I'd probably add something here reiterating that this isn't a sampling inference method. This is just sampling from the approximated posterior. There was already people in the forum asking about the differences in these methods

theorashid · 2024-12-09T03:53:21Z

tests/test_find_map.py

+
+
+@pytest.mark.parametrize(
+    "method, use_grad, use_hess",


any use_hessp tests? or are we just testing if scipy works here

jessegrabowski added 2 commits October 27, 2024 14:17

Add JAX-based find_MAP

6aa20f7

add better_optimize to CI envs

7ed3b2f

jessegrabowski requested review from fonnesbeck and ricardoV94 October 27, 2024 06:29

jessegrabowski added 5 commits October 27, 2024 14:37

Fix relative import

e412f6f

Remove find_MAP import from module-level __init__.py

f9b6258

Update docstring

ad3abd9

Allow calling find_MAP inside model context without model argument

be1d790

Required patched better_optimize

923eb26

jessegrabowski mentioned this pull request Oct 29, 2024

Laplace approximation not handling non-scalar parameters #376

Closed

ricardoV94 reviewed Oct 31, 2024

View reviewed changes

pymc_experimental/inference/jax_find_map.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Oct 31, 2024

View reviewed changes

pymc_experimental/inference/jax_find_map.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Oct 31, 2024

View reviewed changes

pymc_experimental/inference/jax_find_map.py Outdated Show resolved Hide resolved

jessegrabowski added 4 commits December 1, 2024 02:21

in-progress refactor

f705d43

More refactor

a23762b

Generalize code to use any pytensor backend

2d21403

Reconcile the two laplace approximation functions

4c2529d

jessegrabowski requested a review from ricardoV94 December 3, 2024 18:43

Use absolute import in doctest

07ebe40

jessegrabowski added 2 commits December 4, 2024 11:43

Fix imports

b40e101

Fix unrelated statespace test

bc340c2

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pymc_experimental/inference/find_map.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pymc_experimental/inference/find_map.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pymc_experimental/inference/find_map.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pymc_experimental/inference/laplace.py Outdated Show resolved Hide resolved

ricardoV94 approved these changes Dec 4, 2024

View reviewed changes

jessegrabowski added 4 commits December 4, 2024 19:04

- Rename argument use_jax_gradients -> gradient_backend

da338bf

- Rename function `laplace` -> `sample_laplace_posterior`

Fix typo introduced by rename refactor

3ebbf20

use mode=FAST_COMPILE to get unobserved_value_vars after MAP opti…

2035202

…mization

Rename test_jax_find_map.py -> test_find_map.py

f2504e9

ricardoV94 added the enhancements New feature or request label Dec 4, 2024

ricardoV94 changed the title ~~Add JAX-based find_MAP function~~ Add find_MAP with close JAX integration and fix bug with Laplace fit Dec 4, 2024

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pymc_experimental/inference/find_map.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pymc_experimental/inference/find_map.py Outdated Show resolved Hide resolved

jessegrabowski added 3 commits December 4, 2024 19:19

Improve docstring for fit_laplace

a81079b

Update tests to match new signature

4d88343

Update docstring

9b1cd0e

jessegrabowski merged commit 5055262 into pymc-devs:main Dec 4, 2024
7 checks passed

theorashid reviewed Dec 9, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `find_MAP` with close JAX integration and fix bug with Laplace fit #385

Add `find_MAP` with close JAX integration and fix bug with Laplace fit #385

jessegrabowski commented Oct 27, 2024 •

edited

Loading

theorashid commented Oct 28, 2024

ricardoV94 commented Oct 31, 2024

ricardoV94 commented Oct 31, 2024

ricardoV94 commented Oct 31, 2024

jessegrabowski commented Dec 3, 2024

theorashid commented Dec 4, 2024

ricardoV94 left a comment

ricardoV94 commented Dec 4, 2024 •

edited

Loading

jessegrabowski commented Dec 4, 2024

theorashid commented Dec 4, 2024

jessegrabowski commented Dec 4, 2024

theorashid left a comment

theorashid Dec 9, 2024

theorashid Dec 9, 2024

theorashid Dec 9, 2024

theorashid Dec 9, 2024

theorashid Dec 9, 2024

theorashid Dec 9, 2024

theorashid Dec 9, 2024

theorashid Dec 9, 2024

		return f_loss_and_grad, f_hess, f_hessp


		def _compile_functions(

Add find_MAP with close JAX integration and fix bug with Laplace fit #385

Add find_MAP with close JAX integration and fix bug with Laplace fit #385

Conversation

jessegrabowski commented Oct 27, 2024 • edited Loading

theorashid commented Oct 28, 2024

ricardoV94 commented Oct 31, 2024

ricardoV94 commented Oct 31, 2024

ricardoV94 commented Oct 31, 2024

jessegrabowski commented Dec 3, 2024

theorashid commented Dec 4, 2024

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 commented Dec 4, 2024 • edited Loading

jessegrabowski commented Dec 4, 2024

theorashid commented Dec 4, 2024

jessegrabowski commented Dec 4, 2024

theorashid left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Add `find_MAP` with close JAX integration and fix bug with Laplace fit #385

Add `find_MAP` with close JAX integration and fix bug with Laplace fit #385

jessegrabowski commented Oct 27, 2024 •

edited

Loading

ricardoV94 commented Dec 4, 2024 •

edited

Loading