Shaping ApproxInferenceBase.jl

[francescoalemanno]

text copied from JuliaApproxInference/LikelihoodfreeInference.jl#5
i think this is a better place for this issue, i took the liberty of copying your markdown post @jbrea

ApproxInferenceBase.jl: a common API and some basic utilities

My proposition here is to write together a very light-weight ApproxInferenceBase.jl package that serves as a primary dependency of ABC packages. See for example DiffEqBase.jl or ReinforcementLearningBase.jl for how this is done in other eco-systems. I would include in ApproxInferenceBase.jl

Ingredients

• Basic API to coordinate all algorithms
• everything related to prior distributions
• everything related to summary statistics
• everything related to metrics
• testing (and possibly assertion) utilities
• a well-written documentation of the common API

API

My proposition for the API is the following (I am biased of course, and I am very open to discussion!)

Additional to everything related to priors, summarys stats and metrics,
ApproxInferenceBase.jl exports a function fit! with the following signature

fit!(setup, model, data; verbosity = 0, callback = () -> nothing, rng = Random.GLOBAL_RNG)

Every ABC package that relies on ApproxInferenceBase.jl extends this fit! function, e.g.

ApproxInferenceBase.fit!(method::RejectionABC, model, data; kwargs...) = "blabla"

The user provides models as callable objects (functions or functors) with one argument.
Constants are recommended to be handled with closures.
Extraction of summary statistics is done in the model.
For example

model(params) = "do something with params"

my_complex_model(params, constants) = "do something with params and constants"
model(params) = let constants = "blabla" my_complex_model(params, constants) end

my_raw_model(params) = "returns some raw data"
model(params) = extract_summary_statistics(my_raw_model(params))

struct MyFunctorModel
    options
end
(m::MyFunctorModel)(params) = "do something with m and params"

ABC methods/plans/setups are specified in the form

setup = method(metric = my_metric, kwargs...)
setup = method(prior = my_prior, kwargs...) # if method has a prior

One master packages to access all methods

Similar in spirit to DifferentialEquations.jl we could create one package that aggregates all packages and gives unified access. The dependency graph would be something like

     ApproxInferenceBase.jl
                |
     -----------------------
    |           |          |
 ABCPkg1     ABCPkg2      etc.
    |           |          |
    ------------------------
                |
              ABC.jl

This package does nothing but reexport all the setups/methods defined in the
different packages and the fit! function. The name of this package should of course be discussed.

ABCProblems.jl

I think it would be nice to have a package with typical ABC benchmark problems,
like the stochastic lotka-volterra problem, the blowfly problem etc. Maybe we
could collect them in a package ABCProblems.jl.

New methods to be implemented

Here is an incomplete list of methods that I would love to see implemented in
julia. Together with a collection of benchmark problems one would get a nice box
to benchmark new methods we do research on.

• Expectation Propagation for Likelihood-Free Inference, Barthelmé & Chopin 2014
• Parameter estimation in hidden markov models with intractable likelihoods using sequential monte carlo, Yıldırım et al. 2015
• Approximate Bayesian computation with the Wasserstein distance, Bernton et al. 2019
• Statistical Inference for Generative Models with Maximum Mean Discrepancy, Briol et al. 2019
• Approximate Bayesian Computation with Kullback-Leibler Divergence as Data Discrepancy, Jian et al. 2018
• Automatic Posterior Transformation for Likelihood-Free Inference, Greenberg et al. 2019

Conclusions and Questions

Who would be up for such a collaborative effort?
How do you like my proposition for ApproxInferenceBase.jl? What would you change?
Shall we create ApproxInferenceBase.jl, ABCProblems.jl and ABC.jl? Or something similar with different names?

[torfjelde]

I've been pointed to this issue by @jbrea due to the fact that I'm currently trying to integrate some simple ABC methods into Turing.jl so that it can be used as a drop-in replacement when MCMC methods are impractical to use.

Have you guys had a look at AbstractMCMC.jl? At the time of writing, I cannot claim to have a good overview of the ABC-literature, but from what I can tell it seems like ABC methods would fit nicely into the AbstractMCMC.jl "interface"/framework. The docs are here: https://turing.ml/dev/docs/for-developers/interface.

The one thing that might not fit as nicely is when the transitions are bundled into a Chains object from MCMCChains.jl, since some of the diagnostic statistics aren't as applicable to samples from ABC samplers. But even this part is easily customizable (see for example the HMC sampler which includes a lot of additional stats).

The way I've currently done it in TuringLang/Turing.jl#1334 for simple ABC methods, e.g. random-walk, is to overload the sample_init! and step! function:

• step!: this is the meat of the implementation and essentially defines each step for the ABC method.
• sample_init!: here I pre-compute the statistics for the observations so we don't have to recompute at every call to step!. In additon, I extract the observations from the Turing model, but when working outside of a Turing.jl model you could simple define a AbstractMCMC.AbstractModel as is done here with a data field or something.
• sample_end!: here I just drop the rejected samples (though this could easily be directly encoded into the impl of step! by just sampling until we accept)

Essentially, you'd just define your own subtype of AbstractMCMC.AbstractModel which contains the observations and the prior (say, ABCModel) and then packages such as Turing.jl could incorporate your implemetations by overloading the sample_init! or step! for a Turing.Model by just wrapping in a ABCModel and calling your implementations, followed by an optional overload of sample_end!. This is what we do when interacting with other packages implementing the interface in AbstractMCMC, .e.g Elliptical Slice Sampler (https://github.com/TuringLang/Turing.jl/blob/master/src/inference/ess.jl).

I'd love to hear what you guys think!:)

(Also, if this is not suitable for ABC I'd love to hear that too! That would save me from implementing it in Turing only to realize that this approach will be insufficient later on!)

[jbrea]

Hi @torfjelde. Thanks a lot for your input! I definitely want to have a look at AbstractMCMC.jl for the sampling based ABC methods. Your implementation for non-adaptive acceptance thresholds looks neat! Did you also think about adding to Turing.jl alternative ABC methods like KernelABC? To avoid too much duplicate work, it may also make sense at some point to think about where the core ABC methods should be: in this org or in Turing.jl? If they stay in this org, I think it would just make a lot of sense to have them such that they are very easy to integrate with Turing.jl.

[jbrea]

I looked a little bit more into AbstractMCMC.jl. It seems that the few ABC-MCMC methods (e.g. Marjoram et al. 2003) can nicely be fit in this framework. But I don't really see how the ABC-SMC (Sequential Monte Carlo) methods fit in (e.g. Beaumont et al. 2009, Del Moral et al. 2011, Drovandi and Pettit 2011, Turner and Sederberg 2012). These methods operate iteratlvely on a set of particles. AFAIK they are considered superior to the ABC-MCMC methods. Chapter 4 of the Handbook of Approximate Bayesian Computation (edited by Sisson, Fan and Beaumont) gives a nice overview.

I guess it would be nice to implement the ABC-MCMC methods with AbstractMCMC.jl, but for the ABC-SMC methods, the kernel-based methods and all the more recent methods based on function approximation (like the paper by Greenberg et al. I mentioned above) we need alternative approaches.

---

EDIT:
Thinking a bit more about this, it may make sense to split the packages in this org as follows:

• ABC-MCMC (not yet started, @torfjelde's PR in Turing.jl comes closest)
• ABC-SMC (I think @francescoalemanno's KissABC.jl moves in this direction; maybe it would make sense to integrate also the Beaumont et al. 2009 and the del Moral et al. 2011 methods that are currently in LikelihoodfreeInference.jl)
• KernelABC (some methods exist already in LikelihoodfreeInference.jl)
• ...

What do you think?

[francescoalemanno]

as far as i know, Turing has an SMC implementation which exploits the AbstractMCMC interface, we must look deeper

[torfjelde]

Sorry for the late reply!

So it's not quite clear to me why ABC-SMC would be an issue? As @francescoalemanno pointed out, we already have SMC samplers implemented in the AbstractMCMC interface. Is there a particular isse you have in mind @jbrea?
And recently I was made aware that there are even SMC-like ABC samplers implemented in the AbstractMCMC interface: TuringLang/AdvancedPS.jl#9. Speaking of which, maybe @itsdfish have some thoughts on issues when implementing ABC methods in AbstractMCMC.jl?

And when it comes to function approximation ("emulation") methods, it should be fairly straight-forward to support by performing the initial fitting procedure in sample_init! and then if there's additional online adjustments made to the approximator, that's just done in the step! method, no?

Thinking a bit more about this, it may make sense to split the packages in this org as follows:

I still need to get a better overview of the ABC-literature before I can properly judge which structure is a good idea, but it seems to me that the way to go would be:

1. Define an AbstractMCMC.AbstractModel, say ABCModel, which contains the information of the model (e.g. Turing.Model or something else) necessary to run ABC on it. As far as I can understand, this seems to be the reference statistics obtained from observations + a way of evaluating the prior. Then downstream packages which want to make use of the ABC samplers, e.g. Turing, simply need to wrap this information in an instance of a ABCModel, and then call the corresponding methods for this.
2. Figure out exactly what functionality is shared between the different methods. If there's a lot of information, then maybe partitioning by families of algorithms, e.g. ABC-SMC, is a good idea. Otherwise, (1) could just form the skeleton and all ABC samplers are their own packages where they simply overload sample_init! and step! for the ABCModel.

Calling some more of the Turing folks: @cpfiffer @devmotion @xukai92. Know you guys have a lot on your hands, so feel free to ignore this ping right now; just figured it might be useful to bring someone with a bit more AbstractMCMC.jl experience into the discussion.

[itsdfish]

@torfjelde, I should open with the disclaimer that I do not have a lot of experience implementing MCMC samplers. Nonetheless, I think AbstractMCMC has promise as a basic framework for ABC methods. One advantage is that AbstractMCMC provides the option of interfacing with Turing. In some cases, I overloaded some of the methods in AbstractMCMC. However, I think someone with more experience with MCMC samplers might be able to use default methods more creatively or define generic methods for particle samplers.

[devmotion]

I'll chime in and add some more comments.

AbstractMCMC is not a Turing-specific thing, and actually designed in a way that packages without any Turing relation at all can get a nice interface and some default methods for sampling MCMC chains for free. This is achieved by defining AbstractModel and AbstractSampler types and providing default implementations of StatsBase.sample([rng::AbstractRNG, ]model::AbstractModel, sampler::AbstractSampler[, nsamples::Int; kwargs...]) and StatsBase.sample([rng::AbstractRNG, ]model::AbstractModel, sampler::AbstractSampler, parallelalg::AbstractMCMCParallel, nsamples::Int, nchains::Int[; kwargs...]) (note that sample is not owned by AbstractMCMC) that call into different methods (which apart from the sampling step all have some default implementation). So if you implement the sampling step, in particular you get progress bars and parallel sampling (multithreaded and multicore) for free. In addition you also get an iterator and a transducer for free.

Currently, we're preparing AbstractMCMC 2.0 (TuringLang/AbstractMCMC.jl#42). The user-facing sample API stays the same but developers are expected to switch to immutable samplers. This means that you don't have to implement the functions mentioned above but (most likely only) AbstractMCMC.step for your model and sampler types that returns a tuple of the next sample and the next state (same as Base.iterate). I'm quite excited about this change since hopefully it will allow us to get rid of all the hacks and type instabilities caused by having to initialize samplers without knowing the initial state in Turing. I think since you're still discussing the interface you should target AbstractMCMC 2.0 directly. In particular, sample_init! and sample_end! won't exist anymore.

For Turing integration, it should actually not matter if you use the AbstractMCMC interface or some other API. In any case we will have to transform/wrap the Turing models in a way that your samplers can deal with. The advantage of having a united and general API would be that probably less code and effort would be needed for Turing support of your samplers. I agree with @torfjelde's comments above that probably you would want to have a general model structure that is used by the different ABC methods and maybe also a abstract type for the ABC samplers. Additionally, I would also imagine that the different implementations of AbstractMCMC.step (or any sampling procedure with another API) could call into some basic methods that are reused throughout the different packages.

There's nothing special about SMC at all in Turing regarding the AbstractMCMC interface, although I admit that it's a little unintuitive at first. The number of iterations in SMC just corresponds to the number of particles, and we just return the first particle as first sample etc. Hence the SMC steps (i.e., propagating and reweighting particles) do not actually correspond to the AbstractMCMC.step function. Hence currently all the propagation of the particles happens actually in the first AbstractMCMC step, and in the remaining steps we just pop from the array of weighted particles at the final time point. In principle one could also just implement sample for SMC in a completely different way, not using the default implementation in AbstractMCMC at all, to make it more intuitive - but then one would lose all/most goodies provided by AbstractMCMC for free. I could imagine a general interface for particle methods being added to AdvancedPS at some point (it's just an empty repo so far but the plan is to move the particle filter methods from Turing there and refactor them such that they don't require Turing at all).

As an example, in https://github.com/TuringLang/EllipticalSliceSampling.jl/blob/abstractmcmc2/src/abstractmcmc.jl you can see how EllipticalSliceSampling implements the AbstractMCMC 2.0 interface and thereby gets access to all the goodies and the sampling interface in AbstractMCMC.

[francescoalemanno]

Thank you @devmotion , the new interface looks really neat, i look forward to updating my package to use it, as soon as AbMCMC 2 is released

[jbrea]

Thanks a lot @torfjelde, @itsdfish and @devmotion! Things are much clearer now. Let's go ahead with the AbstractMCMC 2.0 interface 😄.