# 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:** 1\
**Showing post:** 23

<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).

---

_[View the full topic](https://discourse.julialang.org/t/improving-performance-of-item-response-model-in-turing-jl/86249)._