Actually perform iteration with DynamicHMC

[devmotion]

So far, sampling with DynamicHMC just calls DynamicHMC and obtains all samples before any sampling step in Turing is performed. In each sampling step in Turing, we then just pop one sample and save it in a Transition.

This approach has some clear disadvantages, such as being memory inefficient, invalidating any progress logging in Turing, not being able to resume sampling, and also not being able to perform sampling until a convergence criterion is met (not supported in Turing yet but supported in AbstractMCMC).

The disadvantage of making use of the iteration interface in DynamicHMC is that the interface is not part of the official API and hence allowed to break at any time.

In principle, it would be trivial to update the log density function (which is now implemented more clearly as a callable struct that wraps the model and the sampler instead of wrapping a closure) in the sampling step to support (some kind of) Gibbs sampling. It is not included yet since I'm not sure how reasonable it would be to that without adapting the step size and/or the metric of the HMC algorithm (or how they could be updated using DynamicHMC).

[devmotion]

Pkg.installed is deprecated on Julia 1.4. I assume that it is unlikely that anyone will run the tests with DynamicHMC < 2 and according to the comment the bug (?) was observed on 32bit with DynamicHMC < 2, so I guess it should be safe to remove the check completely.

[yebai]

Happy to remove this check.

[devmotion]

not being able to resume sampling

I just noticed that it won't be possible even with these changes since we deliberately only save the samples and their log probability (the full transition is only kept until the next step is performed). However, that means that the resulting Chains object is not sufficient for resuming sampling (the same is true for Gibbs sampling BTW). In general, I think to base resuming on a Chains object is not the most intuitive approach (it's not even the only possible output format); I guess it would be more reasonable to have a general implementation that resumes sampling based on a model, a sampler, and a transition. This general approach could be even be added to AbstractMCMC I assume and might allow us to remove all/most checks for resume_from === nothing and set_resume! in the initializations of the samplers. We would just somehow need the possibility to obtain the final transition from sample, then we could save it (maybe even as part of a Chains object).

[devmotion]

Maybe we could just mcmcsample (and mcmcpsample?) to return the final transition together with the bundled samples. In the standard sample (and psample?) implementation in AbstractMCMC we could just throw it away (?), but in Turing, e.g., we could try to save it somehow.

[cpfiffer]

In the standard sample (and psample?) implementation in AbstractMCMC we could just throw it away (?), but in Turing, e.g., we could try to save it somehow.

Maybe chains.info? It's a little hacky but it might work.

[devmotion]

We could do that but I'm a bit worried that I guess we would have to silently remove the transition if the last sample of the chain is modified or removed. Maybe it should be more explicit, maybe we need an alternative to sample that returns the last transition as well.

[cpfiffer]

Perhaps it could be added to chains as a first-class citizen, such that we introduce a new field chain.state that holds whatever thing (Transition) you need to resume sampling.

[devmotion]

I like that! But still it wouldn't resolve the issue that you might lose the ability to resume sampling at some point without noticing, would it?

[cpfiffer]

Maybe the only time would be chain saving/loading, but that needs to be overhauled to use JLD2 or something better than what we have now. Chain subsetting would just carry over chain.state, so the resumed sampling might not mean much.

One "pie in the sky" way to solve this, which is a very big way to solve this, is to just remove the entire axis array setup, and just store all the raw Transition structs. As long as Transitions support getproperty/getfield/getindex and some metadata, you should be able to retrieve all the individual parameter values quite quickly. This way you can always resume no matter what you do to the chain, but it would be a lot of work and might not be very robust.

[devmotion]

This way you can always resume no matter what you do to the chain, but it would be a lot of work and might not be very robust.

Sure, that's something one could do, similar to just specifying chain_type = nothing. But the main issue with Gibbs and DynamicHMC (and possibly other samplers in the future) is that one wants to include method-specific stuff in the transition for only the next step (so things that are neither samples nor statistics of samples) (more like a state that can't be computed a priori and is modified during sampling) but then just save the samples and their statistics preferably. The same setup, e.g., could allow to remove the state-related hacks and type instabilities in the other HMC samplers as well.

Maybe that shouldn't be included in the transitions at all but rather we should have something like SamplerWithState that is, similar to the transitions, not initialized and computed a priori when Sampler is defined (since we haven't computed a sample and hence a depending state yet) but only after the first sample step. Basically, in all subsequent steps we would then not call step! with a Sampler but with a SamplerWithState. Then it would also make sense to return both the samples and the SamplerWithState which could be used for resuming sampling again.

[cpfiffer]

Did you want to park this until #31 goes through, or merge it over an revisit it if the tooling there makes it into AbstractMCMC? I'm fine with this code as-is.

[cpfiffer]

Also do you have any idea what the performance drop is due to not running all the samples at once?

[yebai]

Shall we give this PR another try?

[devmotion]

I'd say it's more reasonable to open a new PR that upgrades the DynamicHMC sampler (and probably the other samplers at the same time?) to the proposed interface in TuringLang/AbstractMCMC.jl#42, if we agree on it. I wanted to update the AbstractMCMC integration in DynamicPPL and Turing first before merging that PR, to see if it causes any problems (shouldn't be the case but I assumed it would be better to check it first). I kept this PR open mainly as a reminder to fix the integration.

[devmotion]

Closed in favour of #1497.