# Speeding up gradient of logpdf
**URL:** https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708
**Category:** Machine Learning
**Tags:** question, performance, autodiff
**Created:** [February 4, 2024, 7:59pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708 "2024-02-04T19:59:46Z")
**Posts on this page:** 20
**Page:** 1
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 4, 2024, 7:59pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/1 "2024-02-04T19:59:46Z")
Somehow, I seem to be unable to speed up gradient calculations of a simple logpdf. As a minimal working example, consider the following computation:
```julia
using Distributions
import AbstractDifferentiation as AD
import ForwardDiff
import ReverseDiff
N = 1000
c = rand(Poisson(1.2), N)
λ₀ = rand(N)
function loss(c, λ)
sum(logpdf.(Poisson.(λ), c))
end
# Same, but non-allocating
function loss2(c, λ)
sum(((c, λ),) -> logpdf(Poisson(λ), c), zip(c, λ))
end
function adtest(ad, loss, c, λ)
AD.gradient(ad, Base.Fix1(loss, c), λ)
end
```
Benchmarking gives
```julia-repl
julia> using BenchmarkTools
julia> @btime loss2($c, $λ₀);
73.047 μs (0 allocations: 0 bytes)
julia> @btime adtest($(AD.ReverseDiffBackend()), loss2, $c, $λ₀);
1.082 ms (24012 allocations: 968.33 KiB)
```
with `loss2` slightly faster than `loss` and `ForwardDiff` or `Zygote` not competitive at `N = 1000`. Scaling with N seems as expected and there is a huge gap between direct and gradient calculations, i.e., ReverseDiff being about 15-30 times slower than computing the loss depending on `N`.
Trying the same in Jax, it’s scaling seems to be almost unreal
```Python
In [57]: %timeit fwd(c, lam, jit = True).block_until_ready()
27 µs ± 45.1 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)
In [58]: %timeit fwd(c, lam, jit = False).block_until_ready()
1.83 ms ± 603 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [59]: %timeit bwd(c, lam, jit = True).block_until_ready()
9.01 µs ± 25.4 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)
In [60]: %timeit bwd(c, lam, jit = False).block_until_ready()
11.9 ms ± 279 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
```
with the gradient calculation even faster than computing the loss when using `jit`!
Any ideas on how to speed this up a bit, in particular, closing the gap between direct and gradient calculations?
In any case, for completeness the Python code:
```python
import jax.numpy as jnp
import jax.random as random
import jax.scipy as jsp
from jax import jit, grad
N = 1000
key = random.PRNGKey(0)
key, subkey = random.split(key)
c = random.poisson(subkey, 1.2, shape = (N,))
key, subkey = random.split(key)
lam = random.uniform(key, shape = (N,))
def loss(c, lam):
return jnp.sum(jsp.stats.poisson.logpmf(c, lam))
jloss = jit(loss)
def fwd(c, lam, jit=False):
if jit:
return jloss(c, lam)
else:
return loss(c, lam)
dloss = grad(loss, argnums = 1)
jdloss = jit(dloss)
def bwd(c, lam, jit = False):
if jit:
return jdloss(c, lam)
else:
return dloss(c, lam)
```
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 4, 2024, 8:07pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/2 "2024-02-04T20:07:56Z")
Hey there! A few remarks:
- AbstractDifferentiation.jl introduces closures, which can [slow things down a lot](https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/54) (the price to pay for generality). If you care about performance, right now I would recommend calling each AD backend directly.
- For gradients, don’t expect forward mode autodiff (like ForwardDiff.jl) to scale well with N. You need to bet on reverse mode anyway.
- To close the gap I would either try Enzyme.jl or ReverseDiff.jl with a precomputed [tape](https://juliadiff.org/ReverseDiff.jl/api/#the-abstracttape-api).
---
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 4, 2024, 8:35pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/3 "2024-02-04T20:35:15Z")
Thanks, will experiment a bit:
1. That’s a bit unfortunate as I like how easily different autodiff algorithms can be tried when using `AbstractDifferentiation`.
2. Yes, reverse mode is the way to go when `N` gets large.
3. Indeed, precompiling the tape gets it down to 5x slower … when using mini-batched I will need to make sure that all batches are the same size, right? What else do I need to keep fixed in order to get correct gradients with a precompiled tape?
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 4, 2024, 8:35pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/4 "2024-02-04T20:35:33Z")
Setup:
```julia
using BenchmarkTools
using Distributions
using Enzyme
using ReverseDiff
N = 1000;
c₀ = rand(Poisson(1.2), N);
λ₀ = rand(N);
function loss(c, λ)
l = zero(promote_type(eltype(c), eltype(λ)))
@inbounds @simd for i in eachindex(c, λ)
l += logpdf(Poisson(λ[i]; check_args=false), c[i])
end
return l
end
f = Base.Fix1(loss, c₀)
tape = ReverseDiff.GradientTape(f, ones(N));
```
Results
```julia
julia> @btime f($λ₀);
30.151 μs (1 allocation: 16 bytes)
julia> @btime ReverseDiff.gradient!(dλ₀, tape, $λ₀) setup = (dλ₀ = zeros($N));
257.153 μs (0 allocations: 0 bytes)
julia> @btime Enzyme.autodiff(Reverse, f, Active, Duplicated(λ₀, dλ₀)) evals = 1 setup = (dλ₀ = zeros($N));
142.595 μs (5 allocations: 144 bytes)
```
---
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 4, 2024, 9:19pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/5 "2024-02-04T21:19:22Z")
Great, will need to see if I can get Enzyme to work on my actual and somewhat larger code. Still impressive what Jax is doing.
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 5, 2024, 8:38am UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/6 "2024-02-05T08:38:29Z")
> [@bertschi](#):
>
> That’s a bit unfortunate as I like how easily different autodiff algorithms can be tried when using `AbstractDifferentiation`.
My bad, the [issue I quoted](https://github.com/JuliaDiff/AbstractDifferentiation.jl/issues/54) involved some bad benchmarking, and things might also have improved since then.
Still, that’s part of the reason why @hill and I have been working on a very simple prototype for high-performance cross-backend gradients. Feel free to take a look and tells us what you think!
> **[GitHub - gdalle/DifferentiationInterface.jl: An experimental redesign for...](https://github.com/gdalle/DifferentiationInterface.jl)**
>
> An experimental redesign for AbstractDifferentiation.jl. - GitHub - gdalle/DifferentiationInterface.jl: An experimental redesign for AbstractDifferentiation.jl.
> [@bertschi](#):
>
> when using mini-batched I will need to make sure that all batches are the same size, right? What else do I need to keep fixed in order to get correct gradients with a precompiled tape?
Yes, and in general avoid any flow control that depends on the precise inputs you give the function. Basically make sure that the same code path as in the tape is taken every time
---
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 5, 2024, 6:16pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/7 "2024-02-05T18:16:51Z")
Thanks, will have a look.
Seems like I found another slight edge:
```julia
loss5(c, λ) = mapreduce(ReverseDiff.@forward((c,λ) -> logpdf(Poisson(λ), c)), +, c, λ)
f5 = Base.Fix1(loss5, c₀)
tape5 = ReverseDiff.GradientTape(f5, ones(N));
```
```julia-repl
# Running your loss function and gradient for comparison
julia> @btime f($λ₀);
37.671 μs (1 allocation: 16 bytes)
julia> @btime ReverseDiff.gradient!(dλ₀, $tape, $λ₀) setup = (dλ₀ = zeros($N));
264.666 μs (0 allocations: 0 bytes)
julia> @btime f5($λ₀);
39.995 μs (4 allocations: 8.00 KiB)
julia> @btime ReverseDiff.gradient!(dλ₀, $tape5, $λ₀) setup = (dλ₀ = zeros($N));
158.477 μs (0 allocations: 0 bytes)
```
From what I understand, `ChainRules` has similar rules, but I could not (yet) figure out how to use them with ReverseDiff. Further, `@forward` seems to be limited to at most 2-args and my actual code is slightly more complicated …
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 6, 2024, 6:39am UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/8 "2024-02-06T06:39:59Z")
> [@bertschi](#):
>
> From what I understand, `ChainRules` has similar rules, but I could not (yet) figure out how to use them with ReverseDiff.
Here the function is so simple that you shouldn’t need to define custom rules. Those are mostly useful when you can gain a clear performance boost with manual specification , or when the derivative computation fails (eg due to mutation).
---
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 6, 2024, 6:06pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/9 "2024-02-06T18:06:37Z")
Sure, but that’s not what I was after. Basically, when using ReverseDiff on `sum(f.(x))` or `mapreduce(f, +, x)` it ends up with lots of scalar operations on the tape. Instead, using forward-mode for the applications of `f` only stores a single `sum` or `mapreduce` call on the tape and handles the scalar calls of `f` much better (that’s why the @forward declaration gave a nice speedup in my example).
The ChainRules.jl docs have a similar example in the section " [Writing rules that call back into AD](https://juliadiff.org/ChainRulesCore.jl/dev/rule_author/superpowers/ruleconfig.html#Writing-rules-that-call-back-into-AD)" and I was wondering if there exist already rules of such type for `mapreduce` or even better `sum(broadcasted, ...)` and if yes, how to tell ReverseDiff to make use of them?
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 6, 2024, 6:21pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/10 "2024-02-06T18:21:40Z")
> [@bertschi](#):
>
> The [ChainRules.jl](https://juliahub.com/ui/Packages/ChainRules) docs have a similar example in the section " [Writing rules that call back into AD](https://juliadiff.org/ChainRulesCore.jl/dev/rule_author/superpowers/ruleconfig.html#Writing-rules-that-call-back-into-AD)" and I was wondering if there exist already rules of such type for `mapreduce` or even better `sum(broadcasted, ...)` and if yes, how to tell ReverseDiff to make use of them?
I suppose there is indeed a rule that takes care of broadcasting, but I’m not 100% sure. In any case you can tell ReverseDiff.jl to use ChainRules.jl with a macro:
>
However ReverseDiff.jl only accepts one argument, so it won’t save you if you need to manage several. Maybe ComponentArrays.jl + the `@forward` macro would be the solution?
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [February 7, 2024, 5:51pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/11 "2024-02-07T17:51:20Z")
ReverseDiff will selectively use ForwardDiff for certain operations where it can be more efficient. Broadcasting is one of them, so we can make use of that in the loss function:
```julia
poiss_logpdf(c, λ) = logpdf(Poisson(λ; check_args=false), c)
function loss_bc(c, λ)
l = sum(poiss_logpdf.(c, λ))
return l
end
```
Although this does incur some overhead, it makes ReverseDiff gradients faster.
My timings with `loss` from [https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/4:](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/4:)
```julia
18.848 μs (0 allocations: 0 bytes) # forward
139.533 μs (0 allocations: 0 bytes) # ReverseDiff gradient
52.640 μs (0 allocations: 0 bytes) # Enzyme gradient
```
And with `loss_bc`:
```julia
19.942 μs (1 allocation: 7.94 KiB)
59.501 μs (0 allocations: 0 bytes)
81.432 μs (2 allocations: 15.88 KiB)
```
Checking the tape for `loss_bc` shows two operations (`broadcast` + `sum`) as expected.
---
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 7, 2024, 7:46pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/12 "2024-02-07T19:46:09Z")
Thanks, that’s what I was hoping for. Still a bit tricky to get down the right path, i.e.,
```julia
function loss_bc1(c, λ)
sum(logpdf.(Poisson.(λ), c))
end
function loss_bc2(c, λ)
sum(poiss_logpdf.(c, λ))
end
```
makes a huge difference:
```julia-repl
julia> tape_bc1 = ReverseDiff.GradientTape(λ -> loss_bc1(c₀, λ), ones(N));
julia> tape_bc2 = ReverseDiff.GradientTape(λ -> loss_bc2(c₀, λ), ones(N));
julia> @btime ReverseDiff.gradient!(dλ₀, $tape_bc1, $λ₀) setup = (dλ₀ = zeros($N));
4.315 ms (16489 allocations: 523.27 KiB)
julia> @btime ReverseDiff.gradient!(dλ₀, $tape_bc2, $λ₀) setup = (dλ₀ = zeros($N));
99.035 μs (0 allocations: 0 bytes)
```
Yet another example of a function barrier?
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [February 7, 2024, 9:49pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/13 "2024-02-07T21:49:11Z")
Not quite, but a similar idea. The choice to pull `poiss_logpdf` into its own function was intentional. It avoids [Performance regression for BernoulliLogit · Issue #1934 · TuringLang/Turing.jl · GitHub](https://github.com/TuringLang/Turing.jl/issues/1934#issuecomment-1382244113), where trying to broadcast a type’s constructor (without parameters, so Julia thinks it’s a `UnionAll` instead of a proper, fully qualified type) can be extremely slow when dealing with ADs such as ReverseDiff.
---
### Author:  [@bertschi](https://discourse.julialang.org/u/bertschi)
#### Post date: [February 7, 2024, 10:59pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/14 "2024-02-07T22:59:04Z")
Thanks, this is helps a lot. Now, I can also understand all the red parts in the profile flame graphs much better and hopefully fix some more of them in my actual code.
Also Zygote is now very fast on your function:
```julia-repl
julia> @btime Zygote.gradient(λ -> loss_bc2($c₀, λ), $λ₀);
71.454 μs (48 allocations: 48.98 KiB)
```
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 8, 2024, 9:30am UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/15 "2024-02-08T09:30:33Z")
> [@ToucheSir](#):
>
> The choice to pull `poiss_logpdf` into its own function was intentional.
The following repo contains all the `logpdf`s you may want for this trick:
> **[GitHub - JuliaStats/StatsFuns.jl: Mathematical functions related to statistics.](https://github.com/JuliaStats/StatsFuns.jl)**
>
> Mathematical functions related to statistics. Contribute to JuliaStats/StatsFuns.jl development by creating an account on GitHub.
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [February 10, 2024, 3:39pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/16 "2024-02-10T15:39:18Z")
On current Enzyme#main I get the following timings (rerunning all of the above on the same system to more easily compare):
The critical thing for performance on Enzyme just now was marking logabsgamma as having correct memory semantics/custom rule. There’s also several internal buffers Enzyme creates that can [and should] be optimized away which should further improve perf, but Enzyme computing the derivative as fast as the primal seems a good enough start for now.
```julia
julia> using Enzyme
julia> using BenchmarkTools
julia> using Distributions
julia> using Enzyme
julia> using ReverseDiff
julia> N = 1000;
julia> c₀ = rand(Poisson(1.2), N);
julia> λ₀ = rand(N);
julia> function loss(c, λ)
l = zero(promote_type(eltype(c), eltype(λ)))
@inbounds @simd for i in eachindex(c, λ)
l += logpdf(Poisson(λ[i]; check_args=false), c[i])
end
return l
end
loss (generic function with 1 method)
julia> f = Base.Fix1(loss, c₀)
(::Base.Fix1{typeof(loss), Vector{Int64}}) (generic function with 1 method)
julia> tape = ReverseDiff.GradientTape(f, ones(N));
julia> @btime f($λ₀);
25.021 μs (1 allocation: 16 bytes)
julia> @btime ReverseDiff.gradient!(dλ₀, tape, $λ₀) setup = (dλ₀ = zeros($N));
193.504 μs (0 allocations: 0 bytes)
julia> @btime Enzyme.autodiff(Reverse, f, Active, Duplicated($λ₀, dλ₀)) evals = 1 setup = (dλ₀ = zeros($N));
25.530 μs (5 allocations: 144 bytes)
julia> @btime Enzyme.autodiff(Reverse, loss, Active, Const($c₀), Duplicated($λ₀, dλ₀)) evals = 1 setup = (dλ₀ = zeros($N));
25.110 μs (0 allocations: 0 bytes)
julia> function loss_bc(c, λ)
l = sum(poiss_logpdf.(c, λ))
return l
end
loss_bc (generic function with 1 method)
julia> f_bc = Base.Fix1(loss_bc, c₀)
(::Base.Fix1{typeof(loss_bc), Vector{Int64}}) (generic function with 1 method)
julia> tape_bc = ReverseDiff.GradientTape(f_bc, ones(N));
julia> @btime ReverseDiff.gradient!(dλ₀, tape, $λ₀) setup = (dλ₀ = zeros($N));
131.613 μs (0 allocations: 0 bytes)
julia> function loss_bc1(c, λ)
sum(logpdf.(Poisson.(λ), c))
end
loss_bc1 (generic function with 1 method)
julia> function loss_bc2(c, λ)
sum(poiss_logpdf.(c, λ))
end
loss_bc2 (generic function with 1 method)
julia> tape_bc1 = ReverseDiff.GradientTape(λ -> loss_bc1(c₀, λ), ones(N));
julia> tape_bc2 = ReverseDiff.GradientTape(λ -> loss_bc2(c₀, λ), ones(N));
julia> @btime ReverseDiff.gradient!(dλ₀, $tape_bc1, $λ₀) setup = (dλ₀ = zeros($N));
2.908 ms (16489 allocations: 523.27 KiB)
julia> @btime ReverseDiff.gradient!(dλ₀, $tape_bc2, $λ₀) setup = (dλ₀ = zeros($N));
49.711 μs (0 allocations: 0 bytes)
julia> @btime Zygote.gradient(λ -> loss_bc2($c₀, λ), $λ₀);
58.461 μs (49 allocations: 49.00 KiB)
```
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 11, 2024, 12:27pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/17 "2024-02-11T12:27:24Z")
The speedups are impressive! But I don’t understand what you mean by
> The critical thing for performance on Enzyme just now was marking logabsgamma as having correct memory semantics/custom rule
Is this something you fixed in Enzyme yourself before rerunning? If not, what else is different on `Enzyme#main`?
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [February 11, 2024, 2:47pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/18 "2024-02-11T14:47:01Z")
The impactful change was [Enzyme.jl/ext/EnzymeSpecialFunctionsExt.jl at 3db5fe49a20855cb78e6b7c9f825a4dbe631768e · EnzymeAD/Enzyme.jl · GitHub](https://github.com/EnzymeAD/Enzyme.jl/blob/3db5fe49a20855cb78e6b7c9f825a4dbe631768e/ext/EnzymeSpecialFunctionsExt.jl#L7) which allows Enzyme to rewrite the calling convention from Julia of the function in a way that makes calling and optimizing it much faster. In particular logabsgamma had a sret (aka returning multiple arguments via memory), which meant that it could not get effectively optimized. This marking allowed our compiler/rewriting system to (correctly) rewrite it in a way that allowed it to be marked as not touching memory/potentially get deleted/not alias other allocations/etc.
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [February 12, 2024, 9:52am UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/19 "2024-02-12T09:52:05Z")
I see that this is the first Enzyme extension, does that mean many package would need similar tweaks?
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [February 12, 2024, 3:08pm UTC](https://discourse.julialang.org/t/speeding-up-gradient-of-logpdf/109708/20 "2024-02-12T15:08:06Z")
So we use that Enzyme.jl-internal table as a way of specifying that a Julia function behaves like a function Enzyme already knows how to handle and marking it as read-only. For example the table also contains sin.
The only thing that was necesary here was marking it as read-only, but while I was at it I also let it use the Enzyme-internal rule.
Specifically, this is an alternate way to specify the differential behavior of something to Julia-level rules. Specifically, Enzyme proper (in contrast to Enzyme.jl) has a table of functions it knows how to handle, that are common to different languages/etc so they do not need to be re-implemented ([Enzyme/enzyme/Enzyme/InstructionDerivatives.td at main · EnzymeAD/Enzyme · GitHub](https://github.com/EnzymeAD/Enzyme/blob/main/enzyme/Enzyme/InstructionDerivatives.td)). This is one such function.
This function does need a custom rule as it is implemented as an approximation to the true value. Differentiating an approximation != the approximation of the derivative (e.g. d^4/dx^2 [mysin(x) = x-x^3/3!] != mysin’‘’'(x) ). For this reason it should have a custom derivative.
However, again, the custom derivative behavior here wasn’t the necessary part (since actually this function isn’t differentiated wrt if I recall, per Enzyme’s activity analysis proving it wasn’t part of the differentiable graph). The only part was marking it as read-none and using a faster calling convention than Julia emits normally. We do that already for all internal rules, so I just did that (and added the rule) while I was at it.