Hi all, it has been a while as I have been busy! But I wanted to put here a suggestion I have had in my mind for a while and so i wondered what others would think about it as I was reminded of it this morning!
In my experience I have found that it can sometimes be tough to gauge how many samples are needed for a given modelling problem if the problem is quite tricky and you don’t have much of a frame of reference to work with, especially since many uninitiated academics/professionals can be more interested in the minimal amount of samples to get a good result (Maybe there is a convergence metric termination criteria could also be a good idea?).
In my experience, PyMC is a very linear approach in that you state your problem and state how many samples to run it for, but if you don’t know a good initial guess you can either overcut with your samples and spend time waiting and maybe get a worse idea of runtime because you are sampling more than is needed, or you can undercut and get a poor result and have to rerun from square 1 again since I remember hearing a while back it can be quite tough to reinitialise the sampler from where an old run finished?. Correct me if I’m wrong of course, but I did have this mental picture of a free running sampler as a kind of fix for this issue of maybe a GUI with the traceplot of metrics and you can continue adding samples as needed and it would dynamically update the plot and metrics so you could see how things were changing as they went, a useful illustration too.
Of course this is easy enough to do if you run the most samples and you can postprocess, or for simpler algorithms it is just a case of reinitialising the sampling algorithm at the last state for the live version, so really I suppose clarity on the matter depends on how easy/hard it is to reinitialise the samplers.
Either way open to seeing what other people think!
@bwengals was brainstorming ideas for a “tensorboard for GPs”, it sounds similar to what you are suggesting? Resumable samplers is something that has been requested many times, but nobody has stepped up to the plate to lead the refactor that would require. I don’t see any fundamental barriers to doing it, except that we don’t ever return a complete state that could be used to resume. Would be a good claude project that could land in pymc-experimental?
Having something like this would also really speed up development time. For slow sampling problems you could cut them off when you see problems showing up without having to wait for the N samples to finish every time.
1 Like
Indeed, I have my hands full in the short term with thesis corrections and a few other things so i can’t commit large sums of time to it currently at least until I get those corrections done but outside of that it sounds like a useful thing and ive felt quite strongly about it for a while but I don’t have much backend experience but I would still be pretty motivated to bring that to life as a concerted effort as I have more open time now post-PhD and would be a good exercise to build up!
Any notions off the top of your head as to what parts of the state are missing that prohibit reinitialisation? If not that is fine as that will probably be the starting grounds for the investigation!
For NUTS it would be stuff related to the internals of the Hamiltonian simulation. @aseyboldt is the undisputed top expert on that.
PyMC is complicated because you’re allowed to mix together arbitrary steps. So we’d have to work out some kind of serialized record that could collect each step, what variables its tracking, and any internal state it tracks for those variables. Then the steps themselves would have to be refactored so as to be constructable from such an object with the correct states loaded in. @ricardoV94 would have more suggestions and might push back on my characterization of the problem.
It’d be a big project but it definitely has steps that could be achieved (for example you could just get NUTS working alone first)
1 Like
PyMC step samplers have a dataclass with all the fields needed to store and resume sampling arbitrarily.
Just the plumbing wasn’t finished. Due to dev preference this was being implemented specifically for the zarrtrace backend: https://github.com/pymc-devs/pymc/pull/7687
If you use nutpie you can pass blocking=False (or smth) and get async sampling. Then at any time you can request the so far finished draws and see how it’s doing.
(Also both nutpie and pymc accept a callback function that’s called regularly…)
You still have to pre-commit to a number of tune and total draws but at least avoid @bwengals concern of waiting til the end just to find it was terrible.
I’ve also seen advice to just lower tune/draws during early iteration. Although then you don’t know if problems would have persisted with longer sampling