In MCMC warmup it can be useful to have different kinds of warmup steps. Right now AbstractMCMC only has a single step_warmup function, and if the sampler wants to control what each warmup step does, it has to figure out which step it's at, what it should be doing, and do that, which is kind of clunky, and that's probably why e.g. AdvancedHMC does not conform with the sampling warmup interface (basically, it's not implemented at all -- this is a known issue TuringLang/Turing.jl#2678).
DynamicHMC has a similar stage-by-stage warmup https://github.com/tpapp/DynamicHMC.jl/blob/021ffacc569e9ab051999a6dae0d1e32c9fd8877/src/mcmc.jl#L415-L425.
In both of these cases the design pattern has been to create a data structure that specifies this richer information about warmup steps. AdvancedHMC bundles it as part of the sampler; DynamicHMC just passes it as a separate argument to the equivalent of sample(). In both cases, it clashes with what AbstractMCMC expects which is a single integer passed as the num_warmup keyword argument.
Given my total lack of time I'm not sure it's worth the effort, but IMO it might be worth thinking about how AbstractMCMC's interface can support this sort of thing. There are two paths, which the packages above demonstrate.
-
One (the AdvancedHMC way) is to bundle it into the sampler, and have num_warmup take its cue from the sampler: if there is a mismatch (e.g. sampler says 'do X 10 times and Y 20 times', and num_warmup = 60, then the sampler takes precedence, and we issue a warning). This is what Turing does right now with AdvancedHMC. It's, frankly, quite ugly; but if there's a warning for mismatches this could be tolerable.
-
The other (the DynamicHMC way) is to change num_warmup to a different argument that takes this information directly. For example, it could be a list of (nsteps, step_kind), where step_kind is some kind of struct that step_warmup will dispatch on.
I believe the latter might conceivably be backwards compatible, and thus I would lean towards it: if an integer N is passed, we would just upgrade it to (N, AbstractMCMC.DefaultWarmup()), and step_warmup(args..., ::DefaultWarmup; kwargs...) can dispatch to step_warmup(args...; kwargs...), meaning that existing implementations continue working the way they did.
One would also have to think carefully about the implications on Turing's Gibbs sampler, which calls step_warmup on each component sampler in turn. It's not obvious how this information can be communicated easily. Perhaps there's no need to.
In MCMC warmup it can be useful to have different kinds of warmup steps. Right now AbstractMCMC only has a single
step_warmupfunction, and if the sampler wants to control what each warmup step does, it has to figure out which step it's at, what it should be doing, and do that, which is kind of clunky, and that's probably why e.g. AdvancedHMC does not conform with the sampling warmup interface (basically, it's not implemented at all -- this is a known issue TuringLang/Turing.jl#2678).DynamicHMC has a similar stage-by-stage warmup https://github.com/tpapp/DynamicHMC.jl/blob/021ffacc569e9ab051999a6dae0d1e32c9fd8877/src/mcmc.jl#L415-L425.
In both of these cases the design pattern has been to create a data structure that specifies this richer information about warmup steps. AdvancedHMC bundles it as part of the sampler; DynamicHMC just passes it as a separate argument to the equivalent of
sample(). In both cases, it clashes with what AbstractMCMC expects which is a single integer passed as thenum_warmupkeyword argument.Given my total lack of time I'm not sure it's worth the effort, but IMO it might be worth thinking about how AbstractMCMC's interface can support this sort of thing. There are two paths, which the packages above demonstrate.
One (the AdvancedHMC way) is to bundle it into the sampler, and have
num_warmuptake its cue from the sampler: if there is a mismatch (e.g. sampler says 'do X 10 times and Y 20 times', andnum_warmup = 60, then the sampler takes precedence, and we issue a warning). This is what Turing does right now with AdvancedHMC. It's, frankly, quite ugly; but if there's a warning for mismatches this could be tolerable.The other (the DynamicHMC way) is to change
num_warmupto a different argument that takes this information directly. For example, it could be a list of(nsteps, step_kind), wherestep_kindis some kind of struct thatstep_warmupwill dispatch on.I believe the latter might conceivably be backwards compatible, and thus I would lean towards it: if an integer
Nis passed, we would just upgrade it to(N, AbstractMCMC.DefaultWarmup()), andstep_warmup(args..., ::DefaultWarmup; kwargs...)can dispatch tostep_warmup(args...; kwargs...), meaning that existing implementations continue working the way they did.One would also have to think carefully about the implications on Turing's Gibbs sampler, which calls
step_warmupon each component sampler in turn. It's not obvious how this information can be communicated easily. Perhaps there's no need to.