Massive speedup for Turing model with DDM by enabling AutoReverseDiff(;compile=true)

[bfalandays]

Hey SequentialSamplingModels team, I wanted to share about some major performance gains I saw from making modifications to the DDM functions to enable the use of AutoReverseDiff(;compile=true) as the automatic-differentiation backend. I have a pretty complex hierarchical DDM built with Turing.jl, and I'm using NUTS as the sampler. When using the default AD backend (AutoForwardDiff) with 250 iterations, 50 adaptation steps, and target_accept = .65, the model took over 8 hours to run. After altering the code to get AutoReverseDiff(;compile=true) working, the same settings ran in ~12 minutes!

In the current version of SequentialSamplingModels.DDM, using compiled tape for the AD backend isn't an option due to branching control flow in two places:
(1) the pdf function has a step of reflecting the parameter values if choice == 1,
(2) the _pdf sub-function has branches for computing the number of terms for the small-time vs large-time expansion, and also for deciding which to use.

The first branch (reflecting parameters) is pretty simple to eliminate with some arithmetic. Note that I'm coding choice=1 as the upper boundary, choice=0 as lower boundary. So in pdf I compute ν_new = (1 - 2*choice)*ν to reflect ν, and z_new = (1-choice) * z + choice * (1-z) to reflect z.

The second was a bit trickier, but the approach I used was to always evaluate both expansions using the smaller value of K, then use ifelse to return the value I want--unlike the if branch, ifelse always evaluates both sides, so it works fine with compiled tape. I also leveraged the NaNMath.jl to deal with domain errors that arise when forcing every statement to be evaluated.

Here are the key bits of my code:

function pdf(d::DDM{T}, choice::Integer, rt::Real; ϵ::Real = 1.0e-4) where {T <: Real}
        rtT = oftype(d.ν, rt); ϵT = oftype(d.ν, ϵ)
        (ν, α, z, τ) = params(d)

        c = oftype(d.ν, choice)       # NOTE: coded as 0 (lower) or 1 (upper)
        ν_eff = (one(T) - T(2) * c) * ν
        z_eff = (one(T) - c) * z + c * (one(T) - z)
        
        return _pdf(ν_eff, α, z_eff, τ, rtT; ϵ = ϵT)
  end
  
  function _pdf(ν::T, α::T, z::T, τ::T, rt::T; ϵ::T) where {T <: Real}

        u_raw = (rt - τ) / α^2
        u     = max(u_raw, eps(T)) 
        piT   = oftype(u, π)

        # ----- small-time K_s -----
        K_s1 = T(2)
        K_s2 = NaNMath.max(
            sqrt(-T(2) * u * NaNMath.log(T(2) * ϵ * sqrt(T(2) * piT * u))),
            sqrt(u) + one(T)
        )
        cond_s = (T(2) * sqrt(T(2) * piT * u) * ϵ) < one(T)
        K_s    = ifelse(cond_s, K_s2, K_s1)

        # ----- large-time K_l -----

        K_l1 = one(T) / (piT * sqrt(u))
        K_l2 = NaNMath.max(
            sqrt((-T(2) * NaNMath.log(piT * u * ϵ)) / (piT^2 * u)),
            K_l1
        )
        cond_l = (piT * u * ϵ) < one(T)
        K_l    = ifelse(cond_l, K_l2, K_l1)

        p = exp((-α * z * ν) - (T(.5) * (ν^2) * (rt - τ))) / (α^2)

        cond = K_s < K_l
        Kmax = ifelse(cond, ceil(Int, K_s), ceil(Int, K_l)) # use the smaller K for both expansions to save time, since we'll ignore the expansion that needs the larger K anyway
        
        return p * ifelse(cond, _small_time_pdf(u, z, Kmax), _large_time_pdf(u, z, Kmax))  
end

[itsdfish]

@bfalandays, thank you for sharing your performance tip. That is a very large speed up indeed!

Have you tried the Mooncake AD? From what I gather, that might be the go-to reverse mode AD option in the future.