Sequential Sampling Models (e.g., DDM, LBA) in RxInfer?

[DominiqueMakowski]

In psychology/neuroscience, sequential sampling models (such as drift diffusion, linear ballistic accumulators, etc.) are in theory very useful and desired, but in practice, there is an accessibility gap which prevents many researchers to actually use them; they are notoriously hard to implement, be it with stan, R, or Python.

The SSM package in Julia is one of the latest key steps towards making them more accessible, implementing many of the distributions in Julia. It is useable with Turing, but it is very slow, especially for hierarchical models (with random factors).

Here is a small example of a Turing LBA:

using Turing
using SequentialSamplingModels
using Random
using LinearAlgebra
using StatsPlots
using Random

# Generate some data with known parameters
dist = LBA(ν=[3.0, 2.0], A = .8, k = .2, τ = .3)
data = rand(dist, 100)

# Specify LBA model
@model function model_lba(data; min_rt = minimum(data.rt))
    # Priors
    ν ~ MvNormal(zeros(2), I * 2)
    A ~ truncated(Normal(.8, .4), 0.0, Inf)
    k ~ truncated(Normal(.2, .2), 0.0, Inf)
    τ  ~ Uniform(0.0, min_rt)

    # Data Generation
    data ~ LBA(;ν, A, k, τ )
end

If RxInfer was able to fit these types of models (faster than Turing), it would be a real game changer for the field of neuro/psych.

But is it feasible? If not, what would it require?

Thanks a lot!

PS: I just watched your talk, great stuff 👍

[albertpod]

Hi @DominiqueMakowski!

Thanks for your kind words! @bvdmitri has indeed done a tremendous job with RxInfer.jl!

At the moment, implementing this exact model you've provided is a challenge, but it's certainly possible.

Let me explain why it's challenging for now. As you've heard in @bvdmitri's talk, RxInfer.jl maps a probabilistic model into a factor graph. Every distribution you've used will have a corresponding node in the graph, e.g., MvNormal, truncated Gaussians, Uniform, and LBA. However, ReactiveMP.jl (the inference engine of RxInfer.jl) does not have truncated distributions and certainly doesn't have an LBA node. Besides, RxInfer.jl expects that every node in the graph would "know" how to communicate with the rest of the graph by sending so-called messages (updates). RxInfer.jl must also know how to compute the posterior resulting from these messages' collision. This requires some mathematical derivations, which can be challenging as they involve intractable integrals. You can find an example of how to derive the rules here.

Now, there are a few solutions that come to my mind on how to run the inference in this model without much coding and math involvement :) One approach is to cast every distribution into Gaussians and use a deterministic node wrapper over your LBA model. However, this is not ideal, as LBA is a distribution by itself and not deterministic by nature. @bvdmitri and @ismailsenoz are working on implementing an inference solution for distributions of non-exponential families, which could help in this case.

Alternatively, you can use amortization techniques to approximate LBA with Normalizing Flows and then use this approximated density inside the model. I believe @itsdfish had a similar problem here, but it might get complex too. I am not an expert on density estimation problems, so I can only offer limited help with this approach.

P.S. We are working hard now on RxInfer 3.0, which will make the inference for this type of model available almost out of the box.

[DominiqueMakowski]

Thanks for the thorough explanation!

P.S. We are working hard now on RxInfer 3.0, which will make the inference for this type of model available almost out of the box.

Do you suggest that it might be easier to just wait for 3.0 then ?

[bvdmitri]

Yeah, we are working on generic inference methods that would allow inference in a larger set of models (including arbitrary likelihoods). You example with LBA as a likelihood looks interesting, we should take it into account.

[itsdfish]

The LBA is a good choice for validating an amortized inference method like BayesFlow because it has a tractable likelihood function, which means you can use Dominique's Turing code as a ground truth. In my opinion, the biggest advantage of BayesFlow is that it allows fast inference for the vast array of simulation-based models without analytic likelihood functions. The only alternative is to use probability density approximation (e.g., simulating 10,000 data sets per evaluated parameter set) with Differential Evolution MCMC, which can take days or weeks, even in a fast language like Julia. If you decide the use the LBA as an example or test case, feel free to tag me if you have any questions.

[itsdfish]

P.S. We are working hard now on RxInfer 3.0, which will make the inference for this type of model available almost out of the box.

@albertpod, that sounds like a great development! I look forward to trying the new version once it is available. Thanks again for your help and patience.

[albertpod]

I just wanted to give you a quick heads-up on this discussion.
The release of RxInfer 3.0 is postponed to the next year.

@itsdfish we will be trying out models with intractable likelihoods. Hopefully, I can tag you once we start filing PRs :)

[itsdfish]

Thanks for the update! I look forward to seeing the PRs and version 3.0.

[DominiqueMakowski]

Excited to see the great work done on the v4 version to streamline the syntax.
Am I right in understanding that this new version does not yet come with the ability to use any distributions, such as the ones of interest in this thread? If so, is there a very rough roadmap/ETA for it? Kudos again for the hard work @albertpod and team!

[bvdmitri]

Hey @DominiqueMakowski ! There is a rough roadmap here https://github.com/ReactiveBayes/RxInfer.jl?tab=readme-ov-file#roadmap. In short, our plan is to work on the ability to use a broader class of distributions (perhaps even any) right after we finish with the new syntax.

[albertpod]

Thanks for bearing with us @DominiqueMakowski :) It did take more time than we expected.

[wmkouw]

Hi @DominiqueMakowski, I'd like to chime in and hopefully clarify further how/whether your interests and our roadmap align.

First of all, I should say that we (BIASlab) do not typically work with likelihood-free inference, i.e., models where evaluating the likelihood function is either computationally or analytically intractable. My own experience is limited to the literature and a few discussions with researchers on this topic. It is my understanding that the standard procedure, Approximate Bayesian Computation, is to sample parameters, simulate data, reject points (e.g., those that are too dissimilar to an oracle data set) and then use the simulated data to infer a posterior distribution over model parameters (fig). I've learned that more recent methods learn an amortized model of the likelihood to speed up data simulation from model parameters. Is this what you are interested in? @itsdfish 's comment gave me the impression that the SSM models you describe are in this realm.

If this is the case, then I can say two things about alignment with our roadmap:

1. The roadmap contains a Gaussian process node that provides an amortized model of an unknown (in this case intractable) function. There is also the Normalizing Flows demo that @albertpod already mentioned. Either of these nodes could take the place of your LBA likelihood.
2. Sampling procedures are not on our roadmap. So, you will not be able to sample model parameters and feed an amortized simulator with RxInfer alone. Your best bet here is to incorporate an RxInfer model + infer step inside a Turing model, which should be possible, I think, because RxInfer models are differentiable.

If this is not the case, then forget what I said. I have nothing to add to @albertpod's answer.

[itsdfish]

Hi @wmkouw,

My apologies for the delayed reply. I overlooked your message in my notifications. I think both Dominique and I have been looking for a Julia analogue to a Python package called BayesFlow. It is a type of normalizing flow which incorporates a special type of neural network used to optimize the summary statistics used in ABC (deepset or set transformer if I remember correctly). The end result is very accurate amortized inference which can be used for any model, including those without a known likelihood function. I think many of the pieces are available in Julia. The main problem is that I currently lack the background to assemble them into something useful. No problem if this type of inference method is outside of the goals of RxInfer.