Cholesky decomposition performed for every sample(?)

[jonny-so]

Hi, unless I have missed something, it looks as though (e.g.) the NUTS sampler implementation creates a new Euclidean metric object for every sample

blackjax/blackjax/mcmc/nuts.py

Line 123 in 65ae00e

metric = metrics.default_metric(inverse_mass_matrix)

which in turn performs a Cholesky decomposition of the (inverse) mass matrix. Is that correct? It should only be necessary to do this at the end of the adaptation windows. In principle the compiler could be eliding the redundant calls, but I would guess that is asking too much of it in this case.

[junpenglao]

Not sure i understand, that line is a non-op for NUTS because inverse_mass_matrix would be a Metric type

[junpenglao]

OK i think i see what you mean, I think adding another

inverse_mass_matrix = metrics.default_metric(inverse_mass_matrix)

in

blackjax/blackjax/mcmc/nuts.py

Line 150 in 65ae00e

def as_top_level_api(

should fix it?

[jonny-so]

I think it would fix it for the top level API, but not during warmup which calls kernel() directly. Maybe it makes sense to add a check inside kernel that it is in fact a metric object, to avoid a silent performance pitfall. It looks like barker.py also has the same issue.

Fixing it during warmup looks a little trickier. Ideally WindowAdaptationState would store inverse_mass_matrix as a Metric rather than an Array, but Metric is not a jax type because it contains various function-valued fields.

[jonny-so]

hmc and dynamic_hmc also actually. One option to resolve this for window adaptation would be to have separate EuclideanMetric and RiemannianMetric types, with the former being a proper jax type (so it can be used inside scan), as only the latter really needs to contain closures. Then the Metric type would be removed altogether, and MetricTypes would be

MetricTypes = Union[EuclideanMetric, Array, RiemannianMetric, Callable[[ArrayLikeTree], Array]]

or going further, maybe it makes sense to disallow Array and Callable being used as metrics directly. e.g. it may be cleaner to have

MetricTypes = Union[EuclideanMetric, RiemannianMetric]

and only allow explicit construction of those types from Array/Callable, rather than allowing a kind of implicit conversion. This is just one suggestion -- if you have a preferred way to fix the window adaptation I'm happy to do it and submit a PR.

[junpenglao]

Thanks for the thoughtful reply! I agree MetricTypes = Union[EuclideanMetric, RiemannianMetric] is a nice solution, but is there a way to extend to other API that depends on Metric as well?

[junpenglao]

I see that there is 2 route:

1. low level api should only accept a MetricTypes
2. play with hashing so if the input array is identical we can reuse the converted Metric.

[junpenglao]

Converted to an issue #760

[AdrienCorenflos]

Sorry this took me a bit of time to look at (vacation etc). I don't really like the idea of implementing a solution to a problem we have not confirmed exists.

I'm not 100% confident what regime we are worried about though so I can't really go in and check: can @jonny-so provide an example (JITTED) where this is happening and shouldn't. I don't really mind if this is a minimal example or a convoluted for this one, just maybe link to a colab/repo where this is happening and it can be confirmed by using JAX debug calls in the stack. If you can then explain with the example why you think the recomputation of the cholesky shouldn't be happening it would be even better.

[AdrienCorenflos]

To be clear, I'm happy to go in and debug once you give me a starting point, there is no need to isolate the issue yourself.

[jonny-so]

Sure, I've forked the repo at https://github.com/jonny-so/blackjax-issue760. I added a single debug trace to metrics.py, and added a small example of the problem in test.py -- changes can be seen here jonny-so@c7a330e.

If you run (e.g.):

git clone https://github.com/jonny-so/blackjax-issue760
cd blackjax-issue760
python -m venv venv
source venv/bin/activate
pip install -e .
python test.py

then you should see something like:

warming up
cholesky!
... (repeated 200 times)
warmed up, pos={'loc': Array(8.949477, dtype=float32), 'log_scale': Array(3.0040097, dtype=float32)}
cholesky!
... (repeated 100 times)

You can ignore the "pos={..}" part of the "warmed up" line, I just included that to force the print traces to happen in the right order.

If you look in test.py, you can see that I have copied the example given in the quickstart page of the blackjax docs, with 2 modifications, 1) I have used is_mass_matrix_diagonal=False, and 2) I have wrapped the whole thing in jax.jit. The code uses a window adpatation with 200 steps, and then draws 100 samples using the fitted parameters. The mass matrix Cholesky should only need to happen a few times during the window adaptation phase (at the end of "middle" adaptation windows, when the mass matrix actually changes), and shouldn't happen at all after warmup. The example shows it is happening for every sample, both during and after warmup.

Let me know if there is anything else I can do.