# Automatic/Black-box variational inference replication vs Autograd

**URL:** <https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020>\
**Category:** Machine Learning\
**Tags:** flux, zygote\
**Created:** [December 13, 2021, 12:15pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020 "2021-12-13T12:15:23Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![lf28](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lf28/32/31200_2.png) [@lf28](https://discourse.julialang.org/u/lf28)\
**Post date:** [December 13, 2021, 12:15pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/1 "2021-12-13T12:15:23Z")

</div>

Hello all,

I have been trying to replicate a simple example (Fig 2) presented in the paper: Black-box stochastic variational inference in five lines of Python ([https://www.cs.toronto.edu/~duvenaud/papers/blackbox.pdf](https://www.cs.toronto.edu/~duvenaud/papers/blackbox.pdf)). The python code can be accessed via a direct link here. [https://github.com/HIPS/autograd/blob/master/examples/black\_box\_svi.py](https://github.com/HIPS/autograd/blob/master/examples/black_box_svi.py).

I have just translated all the code from Python to Julia; and basically wanted to use Zygote instead of Autograd.

```nohighlight
function gaussian_entropy(log_std)
	dim = length(log_std)
	return 0.5 * dim * (1.0+ log(2*pi)) + sum(log_std)
end

function unpack_params(params, dim)
	μ = params[1:dim]
	logσ = params[(dim+1):end]
	return μ, logσ
end

function log_density(x)
	mu = x[1,:]
	log_sigma = x[2,:]
	sigma_density = logpdf.(Normal(0, 1.35), log_sigma)
	mu_density = logpdf.(Normal.(0, exp.(log_sigma)), mu)
	return sigma_density .+ mu_density
end

function variational_objective(logprob, params, nsample, dim)
	μ, logσ = unpack_params(params, dim)
	samples = randn(dim, nsample) .* exp.(logσ) .+ μ
	lower_bound = gaussian_entropy(logσ) + mean(logprob(samples))
	return -lower_bound
end

```

```julia
    ws = [-1., -1.0, -5.0, -5.0]
	mc_ = 5000
	dim_ = 2
	loss() = variational_objective(log_density, ws, mc_, dim_)
	θ = Flux.params(ws)
	grads = gradient(() -> loss(), θ)
	opt = ADAM(0.1, (0.9, 0.999))
	for t in 1:2000
		Flux.Optimise.update!(opt, ws, grads[ws])
	end

```

It seems grads[ws] initially (for t=1) returns the correctly estimated gradient but later in the loop the gradients are wrong. I have tried to replace grads[ws] with ForwardDiff.gradient; that works fine.

I have also tried to use pullback within the loop (the code used in the VAE example), which works fine. However, it becomes excruciatingly slow.

```julia
	 loss2, back = Flux.pullback(Flux.params(ws)) do
             variational_objective(log_density_, ws, 5000, 2)
         end
	 grad = back(1f0)
	 Flux.Optimise.update!(opt, ws, grad[ws])

```

I am very new to Flux and Zygote. So it is quite likely that I have done something silly in the replication. Any help or suggestion will be really appreciated.

To make the code complete, the following code was used to plot the target density and approximating density. The converging parameters are pretty off.

```julia
begin
	x₁ = range(-2, stop=2, length=151)
	x₂ = range(-4, stop=2, length=151)
	μ_, logσ_=unpack_params(ws, 2)
	plot(x₁,x₂, (x, y) -> exp(log_density([x, y])[1]), st=:contour)
	plot!(x₁,x₂, (x, y) -> pdf(MvNormal(μ_, Diagonal(exp.(2 .*logσ_))), [x,y]), st=:contour)
end

```

Best regards,  
Lei

---

<div class="post-metadata">

**Author:** ![trahflow](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/trahflow/32/30585_2.png) [@trahflow](https://discourse.julialang.org/u/trahflow)\
**Post date:** [December 13, 2021, 2:03pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/2 "2021-12-13T14:03:52Z")

</div>

> [@lf28](#):
>
> It seems grads[ws] initially (for t=1) returns the correctly estimated gradient but later in the loop the gradients are wrong.

Note that in contrast to `autograd.grad()` which returns a _function_ that represents the gradient, `Zygote.gradient(l() -> loss(), θ)` returns the gradient _evaluated_ at `θ` (i.e. a concrete number/array of numbers)!

So what’s wrong here, is that you’re not updating your gradient. You just calculate the gradient once in the beginning and use the same gradient all the time when calling `Flux.Optimise.update!`  
Instead, you have to move the gradient calculation inside your training loop.

```julia
for t in 1:2000
    grads = gradient(() -> loss(), θ)
    Flux.Optimise.update!(opt, ws, grads[ws])
end

```

> [@lf28](#):
>
> I have also tried to use pullback within the loop […] However, it becomes excruciatingly slow.

How slow? Do you have any benchmarks?

---

<div class="post-metadata">

**Author:** ![lf28](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lf28/32/31200_2.png) [@lf28](https://discourse.julialang.org/u/lf28)\
**Post date:** [December 13, 2021, 3:01pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/3 "2021-12-13T15:01:56Z")

</div>

That makes a lot of sense! Thanks a million.

I have changed the code accordingly (also wrap the code in a function); now it converges to the correct approximating distribution. But For 2000 iterations, it takes around 300s to finish; comparing with 1.2 s with ForwardDiff.gradient (which returns the gradient as a function).

Here is the benchmark result (I have set the iteration to 3; otherwise it might take forever):

 ![Screenshot 2021-12-13 at 15.00.40](https://global.discourse-cdn.com/julialang/original/3X/b/5/b5acd7a16b0e308bcbb3ccb527357efe59ac93f9.png)

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [December 13, 2021, 5:57pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/4 "2021-12-13T17:57:07Z")

</div>

> [@lf28](#):
>
> ForwardDiff.gradient (which returns the gradient as a function).

[Differentiation API · ForwardDiff](https://juliadiff.org/ForwardDiff.jl/stable/user/api/#ForwardDiff.gradient) returns the gradient value and not a function, were you perhaps thinking of a something else?

That aside, if ForwardDiff works for you and runs quickly, there’s no reason to stop using it in favour of Zygote. `Flux.Optimise.update!` doesn’t care where the gradients come from, as long as they match the shape of the parameters.

---

<div class="post-metadata">

**Author:** ![lf28](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lf28/32/31200_2.png) [@lf28](https://discourse.julialang.org/u/lf28)\
**Post date:** [December 13, 2021, 6:40pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/5 "2021-12-13T18:40:26Z")

</div>

Thanks for your reply. Mmm… I have basically created a gradient function by using ForwardDiff; I assume ReverseDiff would work as well. i.e. something like

```julia
loss = (x) -> variational_objective(log_density, x, 5000, 2)
gradfun = (x) -> ForwardDiff.gradient(loss, x)

for t in 1:2000
    Flux.Optimise.update!(opt, ws, gradfun(ws))
end

```

I am not sure whether Zygote can be used in a similar fashion?

Yes. I can choose a specific automatic differentiation package to use. But ultimately, I would like to do variational inference on some reasonably complicated models and I want something more uniform; and it seems Flux/Zygote have better constructs/APIs.

I am not sure which part of the variational\_objective makes Zygote slow. I am going to try auto-differentiation on log\_density only and add back the Monte Carlo gradient estimation afterwards to see whether the problem lies in the Gaussian sampling step.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [December 13, 2021, 7:23pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/6 "2021-12-13T19:23:16Z")

</div>

> [@lf28](#):
>
> I am not sure whether Zygote can be used in a similar fashion?

`gradfun(x) = Zygote.gradient(loss, x)` 😉

It’s unclear what the bottlenecks might be for Zygote, but here is one candidate. Slicing of arrays in `unpack_params` and `log_density`). Since `x` is never used as a whole, `μ` and `logσ` could be passed in as separate parameters.

```julia
function log_density(mu, log_sigma)
	sigma_density = logpdf.(Normal(0, 1.35), log_sigma)
	mu_density = logpdf.(Normal.(0, exp.(log_sigma)), mu)
	return sigma_density .+ mu_density
end

function variational_objective(logprob, μ, logσ, nsample)
	sample_μs = randn(nsample) .* exp.(logσ) .+ μ
	sample_logσs = randn(nsample) .* exp.(logσ) .+ μ
	lower_bound = gaussian_entropy(logσ) + mean(logprob(sample_μs, sample_logσs))
	return -lower_bound
end

loss(μ, logσ) = variational_objective(log_density, μ, logσ, 5000)
gradfun(μ, logσ) = Zygote.gradient(loss, μ, logσ)

# split ws into μ, logσ
for t in 1:2000
    ∇μ, ∇logσ = gradfun(μ, logσ)
    Flux.Optimise.update!(opt, μ, ∇μ)
    Flux.Optimise.update!(opt, logσ, ∇logσ)
end

```

---

<div class="post-metadata">

**Author:** ![lf28](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lf28/32/31200_2.png) [@lf28](https://discourse.julialang.org/u/lf28)\
**Post date:** [December 15, 2021, 12:11pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/7 "2021-12-15T12:11:53Z")

</div>

ahhh, thanks a lot! Should have tried that.

I have tried your method. It seems removing the packing/unpacking improves efficiency a bit; but it still takes around 280-290s.

I have tried doing Monte Carlo gradient estimation manually. That makes it even worse … now it takes forever. I think the problem is simply because reverse differentiation takes longer when the number of parameters is small.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [December 15, 2021, 5:04pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/8 "2021-12-15T17:04:25Z")

</div>

Edit: this is slower than the original, see post below for a better solution.

What do the per-iteration benchmarks look like now? I wonder if most of those 290s is actually just AD/compiler overhead on the first iteration.

Another thing to try would be using one big scalar function for `log_density`, just like you’ve done for `gaussian_entropy`:

```julia
function log_density(mu, log_sigma)
	sigma_density = logpdf.(Normal(0, 1.35), log_sigma)
	mu_density = logpdf.(Normal.(0, exp.(log_sigma)), mu)
	return sigma_density .+ mu_density
end

# ... mean(logprob(sample_μs, sample_logσs))

```

Would become

```julia
function log_density(mu, log_sigma)
	sigma_density = logpdf(Normal(0, 1.35), log_sigma)
	mu_density = logpdf(Normal(0, exp(log_sigma)), mu)
	return sigma_density + mu_density
end

# ... mean(logprob.(sample_μs, sample_logσs))

```

This will run a single broadcast and (hopefully) allocate one output instead of multiple intermediate values.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [December 15, 2021, 9:44pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/9 "2021-12-15T21:44:20Z")

</div>

Here’s a version that takes advantage of the optimized routines in Distributions.jl and [https://github.com/TuringLang/DistributionsAD.jl](https://github.com/TuringLang/DistributionsAD.jl).

```julia
using Zygote, Distributions, DistributionsAD, LinearAlgebra, FillArrays
using BenchmarkTools

function gaussian_entropy(log_std)
	dim = length(log_std)
	return 0.5 * dim * log1p(2pi) + sum(log_std)
end

function unpack_params(params, dim)
  μ = params[1:dim]
  logσ = params[(dim+1):end]
	return μ, logσ
end

function log_density_old(x)
	mu = x[1,:]
	log_sigma = x[2,:]
	sigma_density = logpdf.(Normal(0, 1.35), log_sigma)
	mu_density = logpdf.(Normal.(0, exp.(log_sigma)), mu)
	return sigma_density .+ mu_density
end

function variational_objective_old(logprob, μ, logσ, nsample)
	samples = randn(2, nsample) .* exp.(logσ) .+ μ
	lower_bound = gaussian_entropy(logσ) + mean(logprob(samples))
	return -lower_bound
end

function log_density_new(x)
	mu = x[1,:]
	log_sigma = x[2,:]
  sigma_density = logpdf(MvNormal(Diagonal(Fill(1.35^2, length(mu)))), log_sigma)
  mu_density = logpdf(MvNormal(Diagonal(abs2.(exp.(log_sigma)))), mu)
	return sigma_density .+ mu_density
end

function variational_objective_new(logprob, μ, logσ, nsample)
  samples = rand(MvNormal(μ, Diagonal(abs2.(exp.(logσ)))), nsample)
	lower_bound = gaussian_entropy(logσ) + mean(logprob(samples))
	return -lower_bound
end

ws = [-1., -1.0, -5.0, -5.0]
mc_ = 5000
dim_ = 2
μ, logσ = unpack_params(ws, dim_)
loss_old(μ, logσ) = variational_objective_old(log_density_old, μ, logσ, mc_)
loss_new(μ, logσ) = variational_objective_new(log_density_new, μ, logσ, mc_)

```

And the results:

```julia-repl
julia> @benchmark gradient(loss_old, $μ, $logσ)
BenchmarkTools.Trial: 34 samples with 1 evaluation.
 Range (min … max): 144.132 ms … 156.172 ms ┊ GC (min … max): 7.08% … 11.76%
 Time (median): 149.142 ms ┊ GC (median): 9.07%
 Time (mean ± σ): 149.351 ms ± 3.406 ms ┊ GC (mean ± σ): 8.82% ± 1.61%

   █ ▃ ██ ▃ ▃ ▃ ▃ ▃         
  ▇█▇▁▇▁▁▁▇▁▁▁▁█▇▁▇▁▁▁▁▁▁██▇█▁█▁▁▁▁▁▁▇█▇▁▁▇▁█▁▁▁▇▁▁▁▁▁▁█▁▁▁▁▇▁▇ ▁
  144 ms Histogram: frequency by time 156 ms <

 Memory estimate: 50.87 MiB, allocs estimate: 1319987.

julia> @benchmark gradient(loss_new, $μ, $logσ)
BenchmarkTools.Trial: 7053 samples with 1 evaluation.
 Range (min … max): 553.498 μs … 3.567 ms ┊ GC (min … max): 0.00% … 66.38%
 Time (median): 586.556 μs ┊ GC (median): 0.00%
 Time (mean ± σ): 703.969 μs ± 415.897 μs ┊ GC (mean ± σ): 12.82% ± 15.94%

  █▇▅▃▂ ▁▁ ▁▁▁ ▁
  █████████▇▇▆▆▅▃▇██▆▄▃▁▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄████▇▇▆▆▆▆ █
  553 μs Histogram: log(frequency) by time 2.58 ms <

 Memory estimate: 2.24 MiB, allocs estimate: 547.

```

The difference is, unexpectedly, quite drastic! More could probably be done about the allocations (I haven’t tested with StaticArrays, for example), but this is probably closer to what you’re seeing with ForwardDiff.

---

<div class="post-metadata">

**Author:** ![lf28](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/lf28/32/31200_2.png) [@lf28](https://discourse.julialang.org/u/lf28)\
**Post date:** [December 16, 2021, 10:25pm UTC](https://discourse.julialang.org/t/automatic-black-box-variational-inference-replication-vs-autograd/73020/10 "2021-12-16T22:25:26Z")

</div>

This is so interesting to know. I definitely have learnt a lot. Thanks a lot!

It seems the two log\_density functions are not the same. MvNormal will return the total log likelihood. I have changed the variational\_objective\_new accordingly (i.e. divide nsamples). It is probably worthwhile to note using joint log likelihood can speed up a lot; it does make the code harder to understand.