Tracking potential performance improvement

[DominiqueMakowski]

As we know, performance is currently the main limitation. Maybe we can put in this thread potential directions, optimization tips, benchmarks etc. to discuss and track how things are evolving on the performance front.

• https://github.com/marius311/MuseInference.jl: I didn't test it but it sounds promising
• https://github.com/JaimeRZP/MicroCanonicalHMC.jl: haven't tested it either but it could also scrap some gains
• https://mlcolab.github.io/Pathfinder.jl/dev/examples/turing/: Pathfinder to "preheat" the MCMC

[itsdfish]

Nice finds! I think these would be worth exploring. If they seem useful I can add a performance tips page in the documentation. As far as I can tell, there are three varieties of performance tips:

1. Changes to code syntax and structure which improve speed
2. Choice of AD
3. Algorithmic improvments

The two packages you found fall into the third category. I think I should have some time tomorrow to look at one of them and do a comparision to NUTS.

[itsdfish]

@DominiqueMakowski, I am seeking clarification on AD requirements and parallel capabilities. Currently, I was only able to run MuseInference with Zygote AD on a single thread , which has comparable performance to NUTS with ReverseDiff AD (but we could get higher ESS by running multiple threads).

[DominiqueMakowski]

Just wanted to flag Tor's fantastic talk on what's new in Turing. Some questions that I have:

• Should we rewrite the docs to make use of the new model | (y=data, ) syntax? I'm not sure I'm yet 100% convinced, though I like the proximity with Bayes formula
• In the potential section on performance, we should probably showcase their new benchmark_model()

[itsdfish]

I'm not sure the benefit of model | (y=data, ) over model(y) (if any), warrents re-writting the documentation. I can mention the new syntax in the first Turing example.

benchmark_model() does seem useful. I will make note of that in the documentation.

[itsdfish]

One of the drawbacks of using the new syntax is that you cannot derive other variables like min_rt in a safe way. In the example below, data referenced the function signature is a global varaible because data is not passed directly with the new syntax. This can be problematic for two reasons: (1) global variables have performance penalties, and (2) the wrong data will be used if the global and local variables have different names. See the example below.

using LinearAlgebra
using SequentialSamplingModels
using Turing 

data = rand(LBA(), 50)
data1 = rand(LBA(; τ = .5), 50)

@model function model(; 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)
    println("min_rt $min_rt")

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

# correct: data global and data local are the same 
sample(model() | (; data), NUTS(), 1000)
# incorrect data1 != data
sample(model() | (; data=data1), NUTS(), 1000)

In addition, its not possible to define min_rt in the body of the function because data is not defined at the point when using the new syntax. So I would say that it is a potential footgun when variables must be processed from function inputs. It is OK in other cases.

[torfjelde]

Hi!

@DominiqueMakowski just made me aware of this pattern as a downside with the new condition syntax. IMHO it's generally not advised to write code a model like the above, i.e. having a kwarg depend on a global variable 😬

If you want to be the ability to pass in data as an argument but still use the condition syntax, you can just define this convenience constructor by hand:)

@model function model(; min_rt)
    ...
end

# Define a new constructor that takes data.
model(data) = model(; min_rt = minimum(data.rt))

It's a bit verbose, but it works 🤷

[itsdfish]

Thanks for the suggestion!

It looks like we could used model(data) and define min_rt within the function definition.

[kiante-fernandez]

I was just chatting with him over lunch! The syntax change seems cool. We also spoke about getting the start of TuringGLM. I was gonna play around with those syntax changes for a few examples.

[DominiqueMakowski]

Can you suggest that they rename TuringGLM (if it's someday to be able to accommodate non-GLM models like GAMs or whatnot) to something like TuringModeling.jl 😁

[DominiqueMakowski]

About the new syntax, my only worry is that it adds a layer of complexity, because arguably all new users of Julia will understand calling a model as a function and inputting the data as input parameters, but the new syntax will require explaining named tuples (and why is there a trailing comma , which is not the most streamlined and elegant looking thing) and the fact that the model definition uses variables that are not defined in the arguments

[DominiqueMakowski]

Could we extend support to dataframes? that would be neat I think:

df = DataFrame(y = rand(100))

# Condition on relevant data from dataframe
fit = model | df

Would it be doable with an extension?

[itsdfish]

Today I performed a quick benchmark with Enzyme. The first point of good news is that it ran without error and produced the correct output. The second point of good news is that it ran faster than ForwardDiff and ReverseDiff.

cd(@__DIR__)
using Pkg 
Pkg.activate("")
using Distributions
using Enzyme
using LinearAlgebra
using ReverseDiff
using SequentialSamplingModels
using Turing  

Enzyme.API.runtimeActivity!(true)

n_choices = 20
ν = fill(1, n_choices)
# Generate some data with known parameters
dist = LBA(; ν, A = .8, k = .2, τ = .3)
data = rand(dist, 100)

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

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

# 97.95 seconds
chains_forward = sample(model(data, n_choices), NUTS(1000, .85), 1000)

# 47.58 seconds
chains_enzyme = sample(model(data, n_choices), NUTS(1000, .85; adtype = AutoEnzyme()), 1000)

# compile = false ≈ 3960 seconds (early termination)
# compile = true ≈ 960 seconds (early termination, potentially unsafe caching)
chains_reverse = sample(model(data, n_choices), NUTS(1000, .85; adtype = AutoReverseDiff()), 1000)

The downside is that there are still some known problems with Distributions.jl. Nonetheless, progress is being made.