Zygote should be preferred to ReverseDiff for reverse mode default

[theogf]

Since it is very typical for VI to repeatedly compute gradients and value of a given function it seems Zygote provide a 3x speedup over ReverseDiff. The reason is that apparently Zygote does some kind of caching.
Here is a quick benchmark:

using BenchmarkTools
using Distributions
using Zygote
using ReverseDiff

d = MvNormal(rand(50), rand(50, 50) |> x -> x * x')
f(x) = logpdf(d, x)
X = rand(d, 40)
@btime ReverseDiff.gradient($X) do x
    sum(f, eachcol(x))
end
# 15.941 ms (534899 allocations: 23.11 MiB)
@btime Zygote.gradient($X) do x
    sum(f, eachcol(x))
end
# 5.405 ms (22475 allocations: 4.51 MiB)

[yebai]

Interesting result - did you enable tape caching for ReverseDiff? From my experience, it provides substantial speedup over non-caching ReverseDiff.

@mohamed82008 @willtebbutt @devmotion

[devmotion]

Do you use DistributionsAD? Also d is a non-constant global, do you get the same results if you fix this? And is there any difference if you use

gradient(X) do x
    loglikelihood(d, x)
end

(with constant d)? IIRC DistributionsAD contains optimized differentiation rules for loglikelihood.

[theogf]

@yebai I don't know how to do the tape caching so I cannot try. But @devmotion here are the results with constant d and constant d from TuringMvNormal

using BenchmarkTools
using Distributions, DistributionsAD
using Zygote
using ReverseDiff

m = rand(50)
C = rand(50, 50) |> x -> x * x'
d = MvNormal(m, C)
const c_d = MvNormal(m, C)
ad_d = TuringDenseMvNormal(m, C)
const c_ad_d = TuringDenseMvNormal(m, C)
f(x) = logpdf(d, x)
c_f(x) = logpdf(c_d, x)
ad_f(x) = loglikelihood(ad_d, x)
c_ad_f(x) = loglikelihood(c_ad_d, x)
X = rand(d, 40)

## Previous result
@btime ReverseDiff.gradient($X) do x
    sum(f, eachcol(x))
end
# 18.051 ms (634979 allocations: 26.17 MiB)
@btime Zygote.gradient($X) do x
    sum(f, eachcol(x))
end
# 2.029 ms (14810 allocations: 4.15 MiB)

## With constant d
@btime ReverseDiff.gradient($X) do x
    sum(c_f, eachcol(x))
end
# 18.774 ms (634938 allocations: 26.17 MiB)
@btime Zygote.gradient($X) do x
    sum(c_f, eachcol(x))
end
# 1.481 ms (13480 allocations: 3.32 MiB)

## Same thing but using DistributionsAD/ loglikelihood
@btime ReverseDiff.gradient($X) do x
    sum(ad_f, eachcol(x))
end
# 16.594 ms (533099 allocations: 22.31 MiB)
@btime Zygote.gradient($X) do x
    sum(ad_f, eachcol(x))
end
# 2.650 ms (17890 allocations: 4.29 MiB)

## Same thing but using DistributionsAD/ loglikelihood and constant distribution
@btime ReverseDiff.gradient($X) do x
    sum(c_ad_f, eachcol(x))
end
# 16.560 ms (532938 allocations: 22.31 MiB)
@btime Zygote.gradient($X) do x
    sum(c_ad_f, eachcol(x))
end
# 2.188 ms (17174 allocations: 3.50 MiB)

[devmotion]

You don't have to use the Turing... types. If you load DistributionsAD, you also get the differentiation rules for MvNormal (see e.g. https://github.com/TuringLang/DistributionsAD.jl/blob/d6aaa6452c033312ebd26aaf4f241b3fac9bebb7/src/multivariate.jl#L309). So with this type piracy both approaches in your example should use the same differentiation rules.

[devmotion]

It also seems you still use sum and eachcol with loglikelihood? loglikelihood is already defined as the sum of logpdf for a set of samples (such as columns in a matrix), so both should be removed.

[theogf]

That's fair, but I just use the logpdf of a Gaussian as an example. In practice f might be much more complex and take vectors only. That's actually why I am posting this issue on the AdvancedVI repo

[trappmartin]

In your code, ReverseDiff builds a new Wengert list / gradient tape at every call. My guess is that, if you build and compile the tape once the results will look quite a bit different.

See: https://github.com/JuliaDiff/ReverseDiff.jl/blob/master/examples/gradient.jl

[theogf]

True, but Zygote has the advantage to do that automatically. The only annoying thing with the tape approach is that when samples are stored in a matrix, it forces you to output a vector of vector (I think), I will get back to my refactoring PR #25 soon and try to improve this aspect.

[Red-Portal]

I'll close this issue for now since the topic is quite outdated. Please re-open if anybody feals the issue is still relevant.