# Improving performance of item response model in Turing.jl

**URL:** <https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249>\
**Category:** Performance\
**Tags:** turing\
**Created:** [August 24, 2022, 8:52am UTC](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249 "2022-08-24T08:52:12Z")\
**Posts on this page:** 4\
**Page:** 2

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 31, 2022, 8:17pm UTC](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249/21 "2022-08-31T20:17:32Z")

</div>

> [@AndrewRadcliffe](#):
>
> The Stan default is `stepsize_jitter = 0`, but `Stan.jl` erroneously supplied `stepsize_jitter = 1.0` as the default for quite some time; it is possible that `CmdStan.jl` was also supplying the same, given that `CmdStan.jl` merged into `Stan.jl`. In any case, it would be interesting to see the benchmarks without the jittering, which, I suspect, may be slowing down Stan considerably.

That’s worth opening an issue

---

<div class="post-metadata">

**Author:** ![goedman](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/goedman/32/217_2.png) [@goedman](https://discourse.julialang.org/u/goedman)\
**Post date:** [September 1, 2022, 12:46pm UTC](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249/22 "2022-09-01T12:46:58Z")

</div>

Hmm, is this (stepwise\_jitter=1.0) true?

As far as I remember StanSample.jl (and thus Stan.jl) for at least 2 years (probably 5 years or so) use below settings by default w.r.t stepwise\_jitter

```julia
julia> bayesian_result_stan.model

Model name:
  name = parameter_estimation_model

C++ threads per forked process:
  num_threads = 8
  use_cpp_chains = false
  check_num_chains = true

C++ chains per forked process:
  num_cpp_chains = 1

No of forked Julia processes:
  num_julia_chains = 4

Actual number of chains:
  num_chains = 4

Sample section:
  num_samples = 1000
  num_warmups = 1000
  save_warmup = false
  thin = 1
  seed = -1
  refresh = 100
  init_bound = 2
Adapt section:
  engaged = true
  gamma = 0.05
  delta = 0.8
  kappa = 0.75
  t0 = 10
  init_buffer = 75
  term_buffer = 50
  window = 25

Algorithm section:

  algorithm = hmc

    NUTS section:
      engine = nuts
      max_depth = 10

  Metric section:
    metric = diag_e
    stepsize = 1.0
    stepsize_jitter = 0.0

Data and init files:
  use_json = true

Stansummary section:
  summary true
  print_summary false

Other:
  output_base = /Users/rob/.julia/dev/DiffEqBayesStan/examples/Fitzhugh-Nagumo/tmp/parameter_estimation_model
  tmpdir = /Users/rob/.julia/dev/DiffEqBayesStan/examples/Fitzhugh-Nagumo/tmp

```

```julia
`/Users/rob/.julia/dev/DiffEqBayesStan/examples/Fitzhugh-Nagumo/tmp/parameter_estimation_model sample num_samples=1000 num_warmup=1000 save_warmup=0 thin=1 adapt engaged=1 gamma=0.05 delta=0.8 kappa=0.75 t0=10 init_buffer=75 term_buffer=75 window=25 algorithm=hmc engine=nuts max_depth=10 metric=diag_e stepsize=1.0 stepsize_jitter=0.0 random seed=-1 init=2 id=1 data file=/Users/rob/.julia/dev/DiffEqBayesStan/examples/Fitzhugh-Nagumo/tmp/parameter_estimation_model_data_1.json output file=/Users/rob/.julia/dev/DiffEqBayesStan/examples/Fitzhugh-Nagumo/tmp/parameter_estimation_model_chain_1.csv refresh=100

```

This runs in about 8 seconds in a “normal” setup":

```julia
julia> include("/Users/rob/.julia/dev/DiffEqBayesStan/examples/Fitzhugh-Nagumo/fn.jl");
  8.037559 seconds (2.07 k allocations: 202.500 KiB)

```

I’ll rerun the timings from about a year ago later today, but an “old” plot can be found in DiffEqBayesStan.jl:

![performance_2021](https://global.discourse-cdn.com/julialang/original/3X/6/0/6022b3cc428872e54b90e6d8197b7c9b8064fe84.png)

---

<div class="post-metadata">

**Author:** ![torfjelde](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/torfjelde/32/206542_2.png) [@torfjelde](https://discourse.julialang.org/u/torfjelde)\
**Post date:** [January 18, 2023, 1:37pm UTC](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249/23 "2023-01-18T13:37:43Z")

</div>

So. After _a lot_ of digging, I’ve arrived at the following.

On Distributions@0.25.76 (and somewhere below), the model

```julia
@model function irt(y, i, p; I = maximum(i), P = maximum(p))
    theta ~ filldist(Normal(), P)
    beta ~ filldist(Normal(), I)
    y ~ arraydist(BernoulliLogit(theta[p] - beta[i]))

    return (; theta, beta)
end

```

achieves Stan-perf with ReverseDiff.jl in compiled mode ([GitHub - TuringLang/TuringBenchmarking.jl](https://github.com/TuringLang/TuringBenchmarking.jl)).

On more recent Distributions.jl versions, it’s 3X slower: [Performance regression for BernoulliLogit · Issue #1934 · TuringLang/Turing.jl · GitHub](https://github.com/TuringLang/Turing.jl/issues/1934). That issue made (eventually) made me realize what’s actually going wrong: fast vs. slow in the reverse AD scenario is just a matter of whether or not you hit the “fast” broadcasting pullback (which is just using ForwardDiff.jl under the hood) or not.

Whether or not you _do_ hit this “fast” path depends on a few things and the conditions are slightly different for different AD backends, but it really comes down to the following:

Broadcasting constructors is apparently not great for Zygote.jl and ReverseDiff.jl:

- ReverseDiff.jl sees a `UnionAll` in the broadcast statement, e.g. a `Normal` without it’s type-parameters is a `UnionAll`, and the type-inference breaks down causing a significant perf degradation (and, in this scenario, makes compilation of the tape useless).
  - On Distributions@0.25.76 there’s no explicit `BernoulliLogit` type, but it’s instead a _function_ which diverts to `Turing.BernoulliBinomial`, hence there’s no `UnionAll` involved in the actual broadcast statement and all is good.

- Zygote.jl will do some checks on the inputs to decide whether or not to take the ForwardDiff-path or not: [Zygote.jl/broadcast.jl at master · FluxML/Zygote.jl · GitHub](https://github.com/FluxML/Zygote.jl/blob/master/src/lib/broadcast.jl#L198-L199). If it decides to _not_ take the ForwardDiff-path, it’s _very_ slow. When a constructor is involved, e.g. `Normal`, the element-type `T` in this check will not be of type `Union{Real,Complex}` and so it fails and moves on.

Addressing these, by, for example, using [Faster `arraydist` with LazyArrays.jl by torfjelde · Pull Request #231 · TuringLang/DistributionsAD.jl · GitHub](https://github.com/TuringLang/DistributionsAD.jl/pull/231), you end up seeing the following:

```julia
julia> using LazyArrays # Necessary to trigger the faster paths using the above PR.

julia> # performant model
       @model function irt(y, i, p; I = maximum(i), P = maximum(p))
           theta ~ filldist(Normal(), P)
           beta ~ filldist(Normal(), I)
           y ~ arraydist(BroadcastArray(BernoulliLogit, theta[p] - beta[i]))

           return (; theta, beta)
       end
irt (generic function with 2 methods)

julia> # Instantiate
       model = irt(y, i, p);

julia> # Make the benchmark suite.
       suite = TuringBenchmarking.make_turing_suite(
           model,
           adbackends = [TuringBenchmarking.ForwardDiffAD{40}(), TuringBenchmarking.ReverseDiffAD{true}(), TuringBenchmarking.ZygoteAD()]
       );

julia> run(suite)
2-element BenchmarkTools.BenchmarkGroup:
  tags: []
  "linked" => 4-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "evaluation" => Trial(536.721 μs)
          "Turing.Essential.ReverseDiffAD{true}()" => Trial(780.870 μs)
          "Turing.Essential.ForwardDiffAD{40, true}()" => Trial(137.089 ms)
          "Turing.Essential.ZygoteAD()" => Trial(1.374 ms)
  "not_linked" => 4-element BenchmarkTools.BenchmarkGroup:
          tags: []
          "evaluation" => Trial(537.782 μs)
          "Turing.Essential.ReverseDiffAD{true}()" => Trial(780.371 μs)
          "Turing.Essential.ForwardDiffAD{40, true}()" => Trial(137.479 ms)
          "Turing.Essential.ZygoteAD()" => Trial(1.380 ms)

```

Stan is at ~1.2ms my laptop, so then it goes ReverseDiff \> Stan \> Zygote \> ForwardDiff, where the first 3 are fairly close.

Note that we’re not 100% certain [Faster `arraydist` with LazyArrays.jl by torfjelde · Pull Request #231 · TuringLang/DistributionsAD.jl · GitHub](https://github.com/TuringLang/DistributionsAD.jl/pull/231) is the way to go, but this is quite the performance improvement so we’ll make sure to get it through somehow. Might just mean that if you _really_ want to go “vrooom!” you need to write `y ~ arraydist(BernoulliLogit, x)` rather than `y ~ arraydist(BernoulliLogit.(x))` (we can automate this, but at the beginning I think it’s going to be opt-in by being slightly different).

---

<div class="post-metadata">

**Author:** ![p-gw](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/p-gw/32/210518_2.png) [@p-gw](https://discourse.julialang.org/u/p-gw)\
**Post date:** [January 18, 2023, 3:06pm UTC](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249/24 "2023-01-18T15:06:52Z")

</div>

That’s awesome, thanks for the effort!

I will certainly try it on a real world example and post the differences/improvement here.  
For now I will just mark your post as the answer as this seems to be the way forward.

[Previous page](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249.md?page=1)