Use mv normal instead of iid normals in Poisson

[xukai92]

This make the inference much faster.

[xukai92]

Time with this implementation.

summary_time.pdf

[itsdfish]

Thank you for this pull request. These results raise some interesting questions that I was hoping you might be able to answer.

1. Why is the multivariate normal faster in Turing than the normal?

2. Why did the new syntax ::Type{T}=Vector{Float64}) where {T} not make the model faster?

3. Do you recommend using a multivariate normal for Stan and DynamicHMC?

Thanks!

[xukai92]

Re 1.

The rand and logpdf for MvNormal is one operation and for Normal needs a loop or broadcasting (my understanding).

using Distributions
using BenchmarkTools
D = 10
d1 = Normal(0, 1)
f1() = begin
       x = Vector(undef, D)
       for i = 1:D
       x[i] = rand(d1)
       end
       logpdf.(Ref(d1), x)
end

@benchmark f1()

BenchmarkTools.Trial: 
  memory estimate:  1.47 KiB
  allocs estimate:  50
  --------------
  minimum time:     2.545 μs (0.00% GC)
  median time:      3.560 μs (0.00% GC)
  mean time:        4.215 μs (14.18% GC)
  maximum time:     4.334 ms (99.89% GC)
  --------------
  samples:          10000
  evals/sample:     9

d2 = MvNormal(zeros(D), 1)
f2() = begin
       x = rand(d2)
       logpdf(d2, x)
end

@benchmark f2()
BenchmarkTools.Trial: 
  memory estimate:  336 bytes
  allocs estimate:  3
  --------------
  minimum time:     182.080 ns (0.00% GC)
  median time:      241.653 ns (0.00% GC)
  mean time:        257.503 ns (5.20% GC)
  maximum time:     60.415 μs (99.51% GC)
  --------------
  samples:          10000
  evals/sample:     671

Re 2.

It should be faster than not having this.

Re 3.

For DHMC - sure.
For Stan - I believe Stan does this optimization already though not very sure.

[xukai92]

Also in case you want a numerically more stable verision, check this:

struct LogPoisson{T<:Real} <: Distributions.DiscreteUnivariateDistribution
    logλ::T
end

function Distributions.logpdf(lp::LogPoisson, k::Int)
    return k * lp.logλ - exp(lp.logλ) - lgamma(k + 1)
end

@model poisson(y, x, idx, N, Ns) = begin
    a0 ~ Normal(0, 10)
    a1 ~ Normal(0, 1)
    a0_sig ~ Truncated(Cauchy(0, 1), 0, Inf)
    a0s ~ MvNormal(zeros(Ns), a0_sig)
    
    for i ∈ 1:N
        logλ = a0 + a0s[idx[i]] + a1 * x[i]
        y[i] ~ LogPoisson(logλ)
    end
end

[itsdfish]

Thank you, Kai. The LogPoisson function appears to be more numerically stable so far. Judging by visual inspection, the number of numerical errors has dropped to about 1/3 of the previous number. I will incorporate it into the the MCMCBenchmark code base.

I have made some improvements to the interface so that we can test variants of the same sampler/model. I am currently running a benchmark comparing a Gaussian and Multivariate Gaussian versions of Turing and DynamicHMC as well as a Gaussian version of Turing that uses the new syntax. I think think this will be informative. Results to follow.

[itsdfish]

@xukai92, here is a summary of the results, which are on the branch called Poisson_Comparison. Please feel free to look over the code in case something seems off.

What I found is that the Multivariate Gaussian improved speed and allocations, but it is still about 5-8 times slower than the closest DynamicHMC configuration. Unless I did something wrong, the new array syntax did not provide a performance boost. Effective sample size is still about half of previous benchmarks of Stan.

I'll tag the original Turing issue and we can move the discussion over there if you would like.

Results:

Main_Results.zip

by(results,[:sampler,:Nd],:time=>mean)
15×3 DataFrame
│ Row │ sampler       │ Nd     │ time_mean │
│     │ Symbol⍰       │ Int64⍰ │ Float64   │
├─────┼───────────────┼────────┼───────────┤
│ 1   │ TuringG       │ 1      │ 58.5276   │
│ 2   │ TuringMVG     │ 1      │ 16.8669   │
│ 3   │ TuringGSyntax │ 1      │ 56.4256   │
│ 4   │ DHMCMVG       │ 1      │ 2.30684   │
│ 5   │ DHMCG         │ 1      │ 3.23581   │
│ 6   │ TuringG       │ 2      │ 87.3855   │
│ 7   │ TuringMVG     │ 2      │ 21.8264   │
│ 8   │ TuringGSyntax │ 2      │ 87.484    │
│ 9   │ DHMCMVG       │ 2      │ 3.63249   │
│ 10  │ DHMCG         │ 2      │ 5.21127   │
│ 11  │ TuringG       │ 5      │ 152.145   │
│ 12  │ TuringMVG     │ 5      │ 30.2751   │
│ 13  │ TuringGSyntax │ 5      │ 149.765   │
│ 14  │ DHMCMVG       │ 5      │ 6.62688   │
│ 15  │ DHMCG         │ 5      │ 10.3537   │

[xukai92]

The ESS issue for the Poisson model should be resolved when the generalised U-turn PR is merged (TuringLang/AdvancedHMC.jl#91).

It looks like this in my local

6×7 DataFrame
│ Row │ sampler     │ Nd    │ a0_ess_mean │ a1_ess_mean │ a0_sig_ess_mean │ tree_depth_mean │ epsilon_mean │
│     │ String      │ Int64 │ Float64     │ Float64     │ Float64         │ Float64         │ Float64      │
├─────┼─────────────┼───────┼─────────────┼─────────────┼─────────────────┼─────────────────┼──────────────┤
│ 1   │ CmdStanNUTS │ 1     │ 208.297     │ 208.776     │ 272.023         │ 7.24086         │ 0.0161274    │
│ 2   │ AHMCNUTS    │ 1     │ 203.079     │ 204.569     │ 274.711         │ 7.23374         │ 0.0157469    │
│ 3   │ CmdStanNUTS │ 2     │ 201.638     │ 202.592     │ 263.016         │ 7.46326         │ 0.0119092    │
│ 4   │ AHMCNUTS    │ 2     │ 213.495     │ 215.527     │ 252.686         │ 7.43646         │ 0.0121216    │
│ 5   │ CmdStanNUTS │ 5     │ 174.845     │ 174.25      │ 223.66          │ 7.71896         │ 0.00812359   │
│ 6   │ AHMCNUTS    │ 5     │ 200.267     │ 200.766     │ 230.273         │ 7.75796         │ 0.00801717   │