ReverseDiff performance

[antoine-levitt]

(apologies if this is not the right place for such discussions)

What exactly can we expect of the performance of ReverseDiff? In particular, under what conditions exactly will it produce the gradient of a R^N -> R function in a time comparable to the evaluation of the function (say within a factor of 3-4)?

Below is a function (a simple arithmetic loop over an array), courtesy of @cortner, for which ReverseDiff is consistently ~30 times slower than the function on my machine. Moreoever, the recording and compilation step takes a large amount of time. I thought this was because the FLOPS/byte ratio is too low, as the manual suggests, but increasing it does not change the factor of ~30.

using ReverseDiff: GradientTape, GradientConfig, gradient, gradient!, compile

φ1(r) = r^2 + r^4 + exp(-r)
φ2(r) = -0.1 * r^3

function F(x)
    N = length(x)
    f = φ1(x[2]-x[1])
    for i = 3:N
        f += φ1(x[i]-x[i-1]) + φ2(x[i] - x[i-2])
    end
    return f
end

function benchmark()
    for N in (10000, 20000, 40000)

        # pre-record a GradientTape for `F` using inputs with Float64 elements
        @time const f_tape = GradientTape(F, rand(N))
        # compile `f_tape` into a more optimized representation
        @time const compiled_f_tape = compile(f_tape)
        # some inputs and work buffers to play around with
        gresult = zeros(N)
        x = rand(N)
        println("N = $N")
        println(" > Time for F")
        for n = 1:3
            @time F(x)
        end
        println(" > Time for DF")
        for n = 1:3
            @time gradient!(gresult, compiled_f_tape, x)
        end
    end
end
benchmark()

[jrevels]

Gradients taken via reverse-mode AD do cost O(1) operations w.r.t. the input size, but can have relatively high constant costs depending on how the target function is written.

What you're running into here is that ReverseDiff can't intercept unvectorizable loops as a single unit, so it has to record every scalar operation encountered as part of loop execution. When executing the tape forwards/backwards, updating buffers for the scalar operations can often be slower than performing the scalar operation itself. Thus, you get the performance you're seeing here.

You can achieve better performance by vectorizing your code, since ReverseDiff can intercept vectorized operations as individual nodes in the computation graph (also note that compilation isn't really benefiting you in this case):

using ReverseDiff: GradientTape, gradient!, compile, @forward
using BenchmarkTools: @btime

@forward φ1(r) = r^2 + r^4 + exp(-r)
@forward φ2(r) = -0.1 * r^3

function F(x)
    N = length(x)
    f = φ1(x[2]-x[1])
    for i = 3:N
        f += φ1(x[i]-x[i-1]) + φ2(x[i] - x[i-2])
    end
    return f
end

function Fvec(x)
    N = length(x)
    f = φ1(x[2] - x[1])
    tmp1 = x[3:N]
    tmp2 = x[2:(N-1)]
    tmp3 = x[1:(N-2)]
    f += sum(φ1.(tmp1 - tmp2) + φ2.(tmp1 - tmp3))
    return f
end

N = 40000
f_tape = GradientTape(F, rand(N))
compiled_f_tape = compile(f_tape)
fvec_tape = GradientTape(Fvec, rand(N))
compiled_fvec_tape = compile(fvec_tape)

x = rand(N)
results = zeros(N)

println("Time for F(x):")
@btime F($x)

println("Time for gradient!(results, f_tape, x):")
@btime gradient!($results, $f_tape, $x)

println("Time for gradient!(results, compiled_f_tape, x):")
@btime gradient!($results, $compiled_f_tape, $x)

println("Time for Fvec(x):")
@btime Fvec($x)

println("Time for gradient!(results, fvec_tape, x):")
@btime gradient!($results, $fvec_tape, $x)

println("Time for gradient!(results, compiled_fvec_tape, x):")
@btime gradient!($results, $compiled_fvec_tape, $x)

Output on my machine:

Time for F(x):
  1.729 ms (0 allocations: 0 bytes)
Time for gradient!(results, f_tape, x):
  35.348 ms (0 allocations: 0 bytes)
Time for gradient!(results, compiled_f_tape, x):
  37.608 ms (0 allocations: 0 bytes)
Time for Fvec(x):
  2.040 ms (16 allocations: 2.44 MiB)
Time for gradient!(results, fvec_tape, x):
  6.098 ms (0 allocations: 0 bytes)
Time for gradient!(results, compiled_fvec_tape, x):
  6.123 ms (0 allocations: 0 bytes)

[cortner]

Fantastic - thank you. I will start experimenting a lot more with this now

[antoine-levitt]

Very nice, thanks! To be clear, is this a fundamental limitation of reverse-mode AD, or something that could be optimized in the future?

[jrevels]

To be clear, is this a fundamental limitation of reverse-mode AD, or something that could be optimized in the future?

If an AD tool can prove (or allow users to assert) that a loop can be converted to a map/broadcast call, then it can perform such an optimization. Alternatively, it can provide limited differentiable control flow primitives to prevent the graph from blowing up.

I'm definitely planning on exploring these optimizations at some point.