Future Directions for SequentialSamplingModels.jl

[itsdfish]

A discussion about SSMs.

[kiante-fernandez]

Future Features

@itsdfish and @DominiqueMakowski: Finally getting around to drafting the "works in progress/future direction" post. I believe we have three exciting developments to discuss; I'd love to hear your thoughts.

Features to Consider

1. New Model Integration: I recommend introducing the following models:

   • Circular diffusion model
   • OU process
   • Ordinal response sequential sampling models/RTCON2-like models
   • Levy Flight models

2. Simulation-based Inference Approaches: Consider incorporating likelihood approximation methods or other simulation-based inference techniques. These could facilitate parameter estimation for complex models with Turing. Maybe using something like flux (I know @itsdfish has started looking into this).

3. Bayesian Generalized Linear Models: Utilize the regression formula syntax to draft models for Turing, akin to BRMS. I think the https://github.com/TuringLang/TuringGLM.jl has good examples of this.

I have some ideas and citations for those models I will post in the thread. Excited to keep working with you guys!

[itsdfish]

@kiante-fernandez, thanks for sharing these ideas. I think these are all good directions.

1. I think all of the models above would be nice additions and fit within the scope of the package. Another set of models to consider can be found in

Kvam, P. D. (2019). A geometric framework for modeling dynamic decisions among arbitrarily many alternatives. Journal of Mathematical Psychology, 91, 14-37

which generalizes the circular diffusion model to other geometries.

2. In my opinion, Julia is sorely lacking machine learning approaches for model inference. This is the feature that I would I would most like to see. Although I am far from an expert in the area, BayesFlow appears to be the best algorithm currently available. It's a type of neural network which combines a normalizing flow and a summary network to learn the marginal posterior distributions of a simulation based model given data and prior distributions. Once the network is trained, you can save it and re-use it. Given that it is applicable to a wide variety of models, I think the best approach would be to create a separate package.

I've had discussions with developers for a few different packages. We ran into limitations with RxInfer.jl. Some of the limitations were with the package itself. I also found it difficult to use. They are currently undergoing changes to make the package easier to use. So this might be an option to reconsider in the future.

I also had a lengthy discussion with a developer of [InvertibleNetworks.jl] (slimgroup/InvertibleNetworks.jl#83) who has expertise in this type of neural network approach. His solution was close but did not perform well on data sets of varying size. The solution is to use DeepSets or SetTransforms as the summary network (and probably set the right hyperparameters). He said he might look into it more at some point. Unfortunately, I don't know much about neural networks and finding a good entry point for non-experts has been difficult.

Another possibility is to inquire the Turing developers. I think they have the right combination of skills in programming, statistics and machine learning to develop such a package. They are currently developing NormalizingFlows.jl which is one of the main networks we would need. I wanted to see whether the other options were viable before opening a third discussion. I think it is a reasonable time to start that discussion.

3. I like this idea too. It could be included as a conditional dependency if it is not too large. I tried that initially with SSMPlots.jl, but I found the package extensions to be inflexible and of limited benefit, except for very small extensions. But feel free to try, and if it does not mesh well, it could go into a standalone package. I created a MWE to supplement the spare documentation for package extensions. It might be helpful if you go that route. Feel free to open a PR if you find that the MWE is missing an important feature.

[DominiqueMakowski]

I would see 3. as the crowning jewel for boosting accessibility to beginners, entry-level users, and Julia transitioners. However, it requires first to have a working framework that we're happy with.

A milestone when we could say okay the foundations are there would IMO be to be able to carry out hierarchical DDMs in reasonable times (comparable to, even if slightly slower than, stan / pymc3), as this is afaik the bulk of the "market" of users. ML-assisted solutions sounds great, though it's akin to dark magic given my understanding level. I also have (overly) high hopes for seeing game-changing improvements once Enzyme gets in Turing.

Also linking this discussion where RxInfer devs tease 3.0 that could make it suitable for our needs.

[itsdfish]

@DominiqueMakowski, these results look promising for Enzyme:

Code

using BenchmarkTools
using SequentialSamplingModels 
using Enzyme
using ForwardDiff
using Random
using ReverseDiff
using Zygote

Random.seed!(845)

data = rand(LBA(ν=[3.0,2.0], A = .8, k = .2, τ = .3), 10)

wrapper(data) = θ -> loglikelihood(LBA(ν=[θ[1],θ[2]], A = θ[3], k = θ[4], τ = θ[5]), data)

f = wrapper(data)

θ = [3.0,2.0,.8,.2,.3]

@btime Enzyme.gradient(Reverse, f, $θ)
@btime Zygote.gradient(f, $θ)
@btime ReverseDiff.gradient(f, $θ)
@btime ForwardDiff.gradient(f, $θ)

Results

  12.368 μs (14 allocations: 2.59 KiB)
  2.765 ms (15969 allocations: 546.36 KiB)
  242.962 μs (5492 allocations: 217.67 KiB)
  6.790 μs (12 allocations: 2.03 KiB)

The most appropriate comparison for hierarchical models will be ReverseDiff or Zygote. ForwardDiff is fast for models within a small number of parameters.

[itsdfish]

FYI, I saw that the Cambridge Machine Learning Group is looking for someone to work on the integration of Turing and Enzyme. AD packages are large projects, but it looks like steady progress is being made with Enzyme. My interpretation is this signals that Enzyme is becoming more mature and stable.

[kiante-fernandez]

I tried that initially with SSMPlots.jl, but I found the package extensions to be inflexible and of limited benefit, except for very small extensions. But feel free to try, and if it does not mesh well, it could go into a standalone package.

ArviZ now has a port over from python (or at least calls the python package)that could be useful for some worked examples. Have you played around with ArviZ @DominiqueMakowski

[DominiqueMakowski]

Not really, I actually discovered it recently alongside this tentative addition. But it looks super cool

[kiante-fernandez]

In my opinion, Julia is sorely lacking machine learning approaches for model inference. This is the feature that I would I would most like to see. Although I am far from an expert in the area, BayesFlow appears to be the best algorithm currently available.

While I consider the approach proposed by the BaysFlow group to be extremely viable, I believe there are additional avenues worth exploring. One of these involves using Likelihood Approximation Networks (LANS; Fengler 2021). LANS offers the advantage of being more straightforward to implement. The concept revolves around learning the mapping between parameters, data, and the generated synthetic or empirical likelihood. This trained network can then facilitate relativity more efficient posterior sampling (if you were considering traditional approximate Bayesian computation, refer to Turner 2018).
References:
Fengler, A., Govindarajan, L. N., Chen, T., & Frank, M. J. (2021). Likelihood Approximation Networks (LANs) for fast inference of simulation models in cognitive neuroscience. ELife, 10, e65074. https://doi.org/10.7554/eLife.65074
Turner, B. M., & Zandt, T. V. (2018). Approximating Bayesian Inference through Model Simulation. Trends in Cognitive Sciences, 22(9), 826–840. https://doi.org/10.1016/j.tics.2018.06.003

This strategy would enable us to leverage the capabilities of the ONNX format for translating differentiable likelihood approximators across various backends. The integration could be leverage existing models in HSSM (https://github.com/lnccbrown/HSSM), originally built in PyTorch. We could then just use Flux to interface with Turing.

Relevant Links:
ONNX.jl GitHub: https://github.com/FluxML/ONNX.jl
Turing Tutorial - Bayesian Neural Network: https://turing.ml/dev/tutorials/03-bayesian-neural-network/

[itsdfish]

Yeah. I agree. It would be worth considering other options. I dabbled with LANs before. You are right in that they are easier to implement on the developer side. There are two drawbacks worth considering: (1) they require more training data than BayesFlow (although sometimes less than used in the paper above), and (2) LANs only amortize the evaluation of the likelihood function, so you still have to perform MCMC sampling. Unlike LANS, BayesFlow amortizes the entire inference process, eliminating the need to performance MCMC sampling. The main drawback is that the prior is fixed for a given trained network, which is partially mitigated by the fact that BayesFlow requires less training data. Another drawback is that BayesFlow is more difficult to implement as a developer, but that is a one-time problem. We just need to find someone with the right expertise to help.

I like your idea of leveraging ONNX to reuse trained models across languages and different backends. It seems like it could be useful for any network architecture, including BayesFlow and LANs. I'm not sure how it works, but once you get it into a form that works with Flux, all you need to do to use it with Turing is define a custom distribution and embed the network within the struct.

[DominiqueMakowski]

Bayesian Generalized Linear Models: Utilize the regression formula syntax to draft models for Turing, akin to BRMS. I think the https://github.com/TuringLang/TuringGLM.jl has good examples of this.

Rather than adding that as a dependency here / create a separate package, it would probably be easier to add that directly into TuringGLM, so that it becomes the main interface for various Turing models. The first step would be probably to extend the @formula macro to accept multiple parameters à-la-brms:

@formula(drift ~ 1 + Condition,
         threshold ~ 1 + Condition,
         ndt ~ 1 + Condition)

which would then write the model with default priors

[itsdfish]

I think the problem is that SSMs are not GLMs. So adding SSM capability to GLMs is not within the scope of the package.

[DominiqueMakowski]

yeah their pkg name is not ideal but I wouldn't been surprised if they were interested in supporting other models like GAMs or SSMs since it would rely on the same internals nad mechanics. After all, they already support mixed models despite the fact that the name is not TuringGLMM 🙃

[itsdfish]

Feel free to open an issue and see if they would be interested in supporting other models and renaming the package. Alternatively, if someone can provide a minimum working example of package extensions, it can be supported here.

[itsdfish]

@DominiqueMakowski, I think the biggest challenge is defining a new formula syntax for SSMs and adapting the macros to work with the new syntax. In a simple linear model, formula(y ~ 1 + x) corresponds to

$\beta_0 \sim$ normal(0,1)
$\beta_1 \sim$ normal(0,1)
$\sigma \sim$ Gamma(1, 1)
$\mu = \beta_0 + \beta_1 x$
$y \sim$ normal($\mu$, $\sigma$)

So the software generates parameters ($\beta_0, \beta_1, \sigma$) and the user must specify the column names in a dataframe corresponding to the predictors and criterion variables. In an SSM, there are more dimensions of data, more parameters, and some parameters can be vectors. For example,

$\beta_{0_\nu} \sim$ normal(0,1)
$\beta_{1_\nu} \sim$ normal(0,1)
$\boldsymbol{\nu} = \beta_{0_\nu} + \beta_{1_\nu} x$
$\dots$
(c,rt) $\sim$ LBA($\boldsymbol{\nu}, \dots$)

What you have proposed might work, but there are some unresolved issues. First, there is no declaration of the data (c,rt). Perhaps you could assume it has a fixed name or perhaps you could have something like

@formula((c,rt) ~ LBA,
         drift ~ 1 + Condition,
         threshold ~ 1 + Condition,
         ndt ~ 1 + Condition)

Another small detail is whether $\sim$ should be used for parameters. In the standard notation above, $=$ is typically used because $\boldsymbol{\nu}$ is determined by other parameters. $\boldsymbol{\nu}$ is technically a random variable from a Bayesian perspective, but its distribution is not specified explicitly. So I am not sure what the right convention would be.

The other challenge would be working with macros. Macros are a very different way of programming. Unfortunately I don't have any experience with it currently.

[DominiqueMakowski]

Interesting, yeah, there are some additional considerations. What I had in mind was mostly copying a brms interface (which would align with TuringGLM's goals I believe), but you are right that we might need to do some bespoke things that might warrant a separate extension (or package, but you know my thoughts on the multiplication of small single-feature packages).

@formula((c,rt) ~ LBA,
         drift ~ 1 + Condition,
         threshold ~ 1 + Condition,
         ndt ~ 1 + Condition

The tuple is IMO not very pleasing to the eyes, but I don't see how else to declare the variable names indeed.

In case there is a parameter vector (for drifts or SDs), one could imagine that:

drift ~ 1 + Condition

Would mean repeat for vector length (i.e., same for all drifts), but we could allow something like:

drift[1] ~ 1 + Condition
drift[2] ~ 1 + Condition

In any case, this will involve a bit of thought about how to logically and elegantly support that.

whether ~ should be used for parameters.

It doesn't bother me. Everyone seems to converge on a brms-like way of expressing models, so we should facilitate that, even if there are some mathematical approximations. in brms, a hierarchical DDM is for instance:

formula <- bf(rt | dec(response) ~ condition + (0 + condition|p|id), 
               bs ~ condition + (condition|id), 
               ndt ~ condition + (condition|id),
               bias ~ condition + (condition|id)
)

The other challenge would be working with macros.

It's alright, I don't think this interface should be a priority anyway, better to focus on supporting the missing models (#29 ?) and making some nice tutorials & examples for now. I will soon(ish) get back to that and try to use SSM for my own data, so I might become a bit more active here to polish some of the docs.

[DominiqueMakowski]

Not really a future direction point but an interesting paper: https://osf.io/preprints/psyarxiv/h4fde

[itsdfish]

Thanks for sharing. I have not read through it in detail but two things caught my attention. First, as noted in the discussion:

"It is clear from the simulations that while researchers typically fix the diffusion coefficient in evidence accumulation
models, the diffusion coefficient is not just a scalar of the
other parameters. Rather, it has a substantive interpretation
as the degree of noise in information processing, and as such
it is of interest to ask how it is affected by experimental manipulations, and how it relates to other variables in individual
differences investigations"

If true, this is interesting because most sources say that the other parameters are a scalar multiple of the diffusion parameter. For example, increasing the diffusion parameter by a factor of 10 would increase the other parameters by a factor of 10.

The second point that was interest was

"Some algorithms do not easily find the
posterior distributions of classically unidentifiable models,
such as MCMC algorithms. However we used a package,
BayesFlow, that made this achievable with simulation-based
Bayesian inference"

I suspect the joint posterior is strongly determined by the choice of prior (in a non-bayesian approach I think the parameters would be a plane in 3D space rather than a cone). Having the ability to sample from the joint posterior in such cases seems to be another advantage of certain neural network approaches. I wish we had a BayesFlow analog in Julia.