# Excessive allocations when using Zygote.gradient

**URL:** <https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296>\
**Category:** Performance\
**Tags:** memory-allocation, zygote\
**Created:** [January 13, 2021, 6:08pm UTC](https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296 "2021-01-13T18:08:40Z")\
**Posts on this page:** 5\
**Page:** 1

<div class="post-metadata">

**Author:** ![MrRobot](https://avatars.discourse-cdn.com/v4/letter/m/b38774/32.png) [@MrRobot](https://discourse.julialang.org/u/MrRobot)\
**Post date:** [January 13, 2021, 6:08pm UTC](https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296/1 "2021-01-13T18:08:40Z")

</div>

I have as a MWE a lennard-Jones force calculation. I would like to define only energy functions, and then use automatic differentiation to calculate the derivative, and thus, force.

However, for the LJ calculation, the zygote.gradient force calculation leads to several order of magnitude worse performance than the “by hand” analytical version. @btime gives 123.161 ms (910316 allocations: 40.93 MiB) for the zygote version and compared to the by hand force calculation which is 67.900 μs (1 allocation: 2.50 KiB).

A second problem I have is that I cannot exit the force calculation via an if statement in the energy function if two particles are far apart. I am unsure how to make this work. Zygote does work with if statements.

```julia
using BenchmarkTools
using StaticArrays
using Zygote

"Vector between two coordinate values, accounting for mirror image seperation"
@inline function vector1D(c1::Float64, c2::Float64, box_size::Float64)
    if c1 < c2
        return (c2 - c1) < (c1 - c2 + box_size) ? (c2 - c1) : (c2 - c1 - box_size)
    else
        return (c1 - c2) < (c2 - c1 + box_size) ? (c2 - c1) : (c2 - c1 + box_size)
    end
end

""" Potential energy between two atoms"""
@inline function pair_energy(r1::SVector{3, Float64}, r2::SVector{3, Float64}, box_size::SVector{3,Float64})

    # apply mirror image separation
    dx = vector1D(r1[1], r2[1], box_size[1])
    dy = vector1D(r1[2], r2[2], box_size[2])
    dz = vector1D(r1[3], r2[3], box_size[3])

    rij_sq = dx * dx + dy * dy + dz * dz

    ##### Would like to implement the below commented out section
    #if rij_sq > box_size[1] / 2
    # return 0.0
    #end
    
    sr2 = 1 / rij_sq
    sr6 = sr2^3
    sr12 = sr6^2
    e = 4 * (sr12 - sr6)
    return e
end

funct(x, y, b) = gradient(x -> pair_energy(x, y, b), x)

function analytical_total_force(r::Vector{SVector{3, Float64}}, box_size::SVector{3,Float64})
    n = length(r)
    forces = [SVector{3}(0.0, 0.0, 0.0) for i=1:n]

    for i = 1:(n-1)
        for j = (i+1):n
            dE_dr = funct(r[i], r[j], box_size)[1] # 
            forces[i] = forces[i] - dE_dr 
            forces[j] = forces[j] + dE_dr 
        end
    end
    return forces
end

natoms = 100
box_size = SVector{3}(6.0, 6.0, 6.0)
 r = [SVector{3}(rand(), rand(), rand()) .* box_size[1] for i = 1:natoms]
@btime test_analytical_force = analytical_total_force($r, $box_size)

```

---

<div class="post-metadata">

**Author:** ![contradict](https://avatars.discourse-cdn.com/v4/letter/c/ac91a4/32.png) [@contradict](https://discourse.julialang.org/u/contradict)\
**Post date:** [January 13, 2021, 6:37pm UTC](https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296/2 "2021-01-13T18:37:57Z")

</div>

I don’t have any advice about the performance part of the question, but I think the reason the conditional gives trouble is that AD systems have two ideas that look like 0. The first is “the derivative of the output with respect to that parameter at this point is zero”, Zygote uses 0 for this. The second idea is “the value at this point is independent of that parameter”, Zygote use `nothing` for this.

The conditional is exactly this second condition, so you have two possible choices. You can introduce a dependency to get back to the first case (`return 0.0*rij_sq`), or you can handle a gradient value of `nothing` when summing the forces. Surprisingly, the second option is a little faster on my machine, you should probably test with your actual code to see how it performs for you.

---

<div class="post-metadata">

**Author:** ![MrRobot](https://avatars.discourse-cdn.com/v4/letter/m/b38774/32.png) [@MrRobot](https://discourse.julialang.org/u/MrRobot)\
**Post date:** [January 13, 2021, 6:42pm UTC](https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296/3 "2021-01-13T18:42:20Z")

</div>

Thanks, yes that does make it work. The time gets worse with the if statement, and is now 200.074 ms (1273661 allocations: 52.80 MiB). I don’t know if this is typical for AD, but it makes me want to just revert to coding derivatives by hand ☹

---

<div class="post-metadata">

**Author:** ![marius311](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/marius311/32/3953_2.png) [@marius311](https://discourse.julialang.org/u/marius311)\
**Post date:** [January 13, 2021, 6:58pm UTC](https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296/4 "2021-01-13T18:58:16Z")

</div>

Given the input dimension in the function you’re taking a gradient of is really small (3), this is probably better suited for ForwardDiff. Try with

```julia
funct(x, y, b) = ForwardDiff.gradient(x -> pair_energy(x, y, b), x)

```

You’ll also need to remove the `::Float64` annotations on `vector1D` and change the `pair_energy` ones to just `SVector{3}` (or remove them completely). With that I get:

```julia
julia> @btime test_analytical_force = analytical_total_force($r, $box_size)
  924.744 μs (1 allocation: 2.50 KiB)

```

---

<div class="post-metadata">

**Author:** ![MrRobot](https://avatars.discourse-cdn.com/v4/letter/m/b38774/32.png) [@MrRobot](https://discourse.julialang.org/u/MrRobot)\
**Post date:** [January 13, 2021, 7:16pm UTC](https://discourse.julialang.org/t/excessive-allocations-when-using-zygote-gradient/53296/5 "2021-01-13T19:16:23Z")

</div>

Hi, thanks. With your changes I get 110.199 μs (1 allocation: 2.50 KiB) for ForwardDiff and 64.700 μs (1 allocation: 2.50 KiB) for the analytical “by-hand” version. That made a really big difference. I will have to profile it to see if factor of two can be improved.