DynamicHMC is (~5 times) faster than Turing's HMC implementation

[yebai]

Time used for collecting 2 million samples

• Turing.NUTS: 3318.293540 seconds (33.35 G allocations: 1.299 TiB, 29.35% gc time)
• Turing.DynamicNUTS: 848.741643 seconds (7.86 G allocations: 251.076 GiB, 32.13% gc time)

using the following

using DynamicHMC, Turing, Test

@model gdemo(x, y) = begin
  s ~ InverseGamma(2,3)
  m ~ Normal(0,sqrt(s))
  x ~ Normal(m, sqrt(s))
  y ~ Normal(m, sqrt(s))
  return s, m
end

mf = gdemo(1.5, 2.0)

@time chn1 = sample(mf, DynamicNUTS(2000000));

@time chn2 = sample(mf, Turing.NUTS(2000000, 0.65));

[skanskan]

What about Stan and Klara?

[tpapp]

This is how I would code the above in DynamicHMC 1.0:

using DynamicHMC, TransformVariables, LogDensityProblems, Distributions

struct GaussianDemo{T,P}
    y::T
    x::T
    s_prior::P
end

function (p::GaussianDemo)((m, s))
    d = Normal(m, √s)
    logpdf(d, p.x) + logpdf(d, p.y) + logpdf(p.s_prior, s) + logpdf(Normal(0, √s), m)
end

p = GaussianDemo(1.5, 2.0, InverseGamma(2, 3))
P = TransformedLogDensity(as((m = asℝ, s = asℝ₊)), p)
∇P = ForwardDiffLogDensity(P)
@time chain, tuned = NUTS_init_tune_mcmc(∇P, 10_000)

which on my laptop yields

  0.430758 seconds (3.74 M allocations: 211.077 MiB, 18.41% gc time)

Note that I changed the sample size since

julia> using MCMCDiagnostics
julia> mapslices(effective_sample_size, get_position_matrix(chain); dims = 1)
1×2 Array{Float64,2}:
 3841.64  4080.44

so I don't think you need 2 million samples. Finally, for real-world problems with higher dimensions, I would have used

∇P = FluxGradientLogDensity(P)

since reverse-mode AD is faster as the dimension increases (YMMV).

[mohdibntarek]

On the type stability branch using 100_000 samples:

• NUTS: 20.769626 seconds (166.79 M allocations: 9.301 GiB, 19.40% gc time)
• DynamicHMC: 27.478044 seconds (230.15 M allocations: 13.436 GiB, 18.68% gc time)

[xukai92]

@mohamed82008 you mean a type stable Turing with Turing's HMC or this?

[mohdibntarek]

No this is before this package, just using Turing HMC.

[xukai92]

Wow then we would have a lot more potential important still. I will do a simple integration of this package in the same way as DynamicHMC in Turing after I port all adaptation codes into this package. We will see how much we can get by doing that :D

[cpfiffer]

That's a pretty amazing improvement!

[xukai92]

Just did some quick test against DynamicHMC with 3ac4db1. We are among the same speed now.

#samples	DynamicHMC	AdavancedHMC
10_000	14.57 / 5.59	12.47 / 2.58
100_000	24.16 / 13.36	23.84 / 11.96
500_000	80.29 / 55.00	88.38 / 61.43

Numbers are in seconds with the first number being wrapping @time over Turing's sample function and the second number being wrapping @time over the packages' core sample function. AHMC is using dense adaption with Stan's three phase schedule. I'm not sure if DHMC is doing the same?

Also the numbers are slightly dependent on initial value and adaption phase, but it seems that DHMC scales better than us. I have few thoughts on why this happens and will create a seperate PR on possible TODOs from an algorithmic point of view.

[trappmartin]

Could you re-run this using @btime? This way the time estimates would be more reliable as they are estimated for multiple calls while time is only a snapshot of a single call.

The macro is from BenchmarkTools.

[yebai]

Updated comparison based on AHMC (7a0e42c) and Turing (TuringLang/Turing.jl@92a0e2a). Collecting 20k samples take

• DynamicHMC + Turing: 18.991198 seconds (55.29 M allocations: 3.017 GiB, 6.91% gc time)
• AdvancedHMC + Turing: 11.790677 seconds (66.83 M allocations: 3.063 GiB, 11.18% gc time)

julia> mf = gdemo(1.5, 2.0)
Turing.Model{Tuple{:s,:m},Tuple{:x,:y},getfield(Main, Symbol("###inner_function#382#10")),NamedTuple{(:x, :y),Tuple{Float64,Float64}},NamedTuple{(:x, :y),Tuple{Symbol,Symbol}}}(getfield(Main, Symbol("###inner_function#382#10"))(Core.Box(getfield(Main, Symbol("###inner_function#382#10"))(#= circular reference @-2 =#))), (x = 1.5, y = 2.0), (x = :nothing, y = :nothing))

julia> @time chn1 = sample(mf, DynamicNUTS(20000))
 18.991198 seconds (55.29 M allocations: 3.017 GiB, 6.91% gc time)

julia> @time chn2 = sample(mf, Turing.NUTS(20000, 0.65));
┌ Info: Finished 1000 adapation steps
│   typeof(adaptor) = AdvancedHMC.Adaptation.StanHMCAdaptor{AdvancedHMC.Adaptation.DensePreconditioner{Float64,AdvancedHMC.Adaptation.WelfordCov{Float64}},AdvancedHMC.Adaptation.NesterovDualAveraging{Float64}}
│   (τ.integrator).ϵ = 0.26506169649879274
└   h.metric = [0.397546, 0.850986]
┌ Info: Finished 20000 sampling steps in 6.903967957 (s)
│   typeof(h.metric) = AdvancedHMC.Adaptation.DenseEuclideanMetric{Float64,Array{Float64,1},Array{Float64,2},LinearAlgebra.UpperTriangular{Float64,Array{Float64,2}}}
│   typeof(τ) = AdvancedHMC.NUTS{AdvancedHMC.Leapfrog{Float64}}
│   EBFMI(Hs) = 19663.70513482061
└   mean(αs) = 0.7903777752674378
 11.790677 seconds (66.83 M allocations: 3.063 GiB, 11.18% gc time)

Note: This test is a very rough benchmark targeting Turing's internals. The test is only meant for getting a feel of speed improvements for Turing. A more thorough comparison of real-world problems is required for a fair evaluation.