Zygote very slow relative to ForwardDiff?

General Usage Performance
1

I’m trying to speed up fitting a model with DynamicHMC.jl, so I was experimenting with different autodiff packages. I’ve seen Zygote mentioned quite a bit as a promising autodiff package, but the performance seems to be several orders of magnitude slower. Am I not using it correctly? Is there any way to avoid all those allocations?

Simple example:

julia> const xs = randn(10^5);

julia> function f(θ)
        sum(logpdf.(Normal(θ[1], θ[2]), xs))
       end
f (generic function with 2 methods)

julia> @time ForwardDiff.gradient(f, [0.0, 1.0]);
  0.356206 seconds (381.72 k allocations: 27.519 MiB, 99.14% compilation time)

julia> @time ForwardDiff.gradient(f, [0.0, 1.0]);
  0.003052 seconds (7 allocations: 2.289 MiB)

julia> @time Zygote.gradient(f, [0.0, 1.0]);
  1.380850 seconds (5.43 M allocations: 176.099 MiB, 12.80% gc time, 8.46% compilation time)

julia> @time Zygote.gradient(f, [0.0, 1.0]);
  1.280075 seconds (5.40 M allocations: 174.718 MiB, 14.57% gc time)
2

Backwards-mode autodiff is designed for the case of a huge number of parameters (hundreds, thousands, or more); with just 2 parameters you are better off with forwards-mode autodiff (like ForwardDiff).

1 Like