Unrealistically poor benchmark results

[dfdx]

White testing my own AD package, I tried to compare its performance to Zygote, but results turned to be unexpectedly poor for it.

Setup

Variational autoencoder with all dense layers on MNIST dataset. Nothing fancy, only built-in functions like matrix multiplication, elementwise operations and aggregation functions. I also defined derivative rules for both - Yota (my package) and Zygote. Full source can be found here.

It's also worth to mention that with the original cost function Zygote failed with segfault, so for simplicity I commended out guilty line for now.

Benchmark results

Compiling derivatives using Yota
  0.050908 seconds (819 allocations: 21.679 MiB, 5.64% gc time)
BenchmarkTools.Trial: 
  memory estimate:  21.68 MiB
  allocs estimate:  819
  --------------
  minimum time:     34.858 ms (0.00% GC)
  median time:      39.443 ms (6.69% GC)
  mean time:        42.610 ms (6.53% GC)
  maximum time:     155.245 ms (1.73% GC)
  --------------
  samples:          118
  evals/sample:     1

Compiling derivatives using Zygote
  1.573936 seconds (7.76 M allocations: 381.904 MiB, 35.54% gc time)
BenchmarkTools.Trial: 
  memory estimate:  381.90 MiB
  allocs estimate:  7762261
  --------------
  minimum time:     1.144 s (26.51% GC)
  median time:      1.337 s (33.87% GC)
  mean time:        1.301 s (32.50% GC)
  maximum time:     1.386 s (34.81% GC)
  --------------
  samples:          4
  evals/sample:     1

I use some aggressive optimizations (e.g. full buffering) so I actually expected 1.5-3x times difference, but not 30x. Note that both libraries give the same results and code doesn't have dead code branches.

What could go wrong?

1. I could make a mistake or misunderstand Zygote's README. Can you spot the problem with the code?
2. We hit some operation without efficient derivative. Maybe sum with keywords or mean?

[MikeInnes]

Thanks for the detailed issue. Zygote still does dumb things every so often (c.f. segfault); given that this code is pretty straightforward there's probably either some simple typing issue in broadcast or we're missing a derivative, like you suggest. It would be great to have a manifest for that benchmark as I might not look at it immediately, but I'll dig into it sometime soon.

One thing to note is that your Zygote adjoint definitions won't have any effect, since we use a fused forward mode for broadcasts (I assume Yota uses a simpler reverse mode?).

What do you mean by "full buffering"? It doesn't seem that you're compiling a tape in this script, otherwise I'd assume you meant something like allocating array storage up front -- but I could be missing something.

[dfdx]

One thing to note is that your Zygote adjoint definitions won't have any effect, since we use a fused forward mode for broadcasts (I assume Yota uses a simpler reverse mode?).

Yep, for broadcasted operations I simply record broadcasted derivatives of scalar functions, e.g. for sin.(x) I just record cos.(x) .* Δ.

What do you mean by "full buffering"? It doesn't seem that you're compiling a tape in this script,

grad(f, args...) is a caching layer for _grad(f, args...). Perhaps it's easier just to show the code:

function grad(f::Function, args...)
    # check if we have already calculated derivatives for this function and arguments 
    # of such types and sizes
    cache_key = (f, ([isstruct(arg) ? typeof(arg) : size(arg) for arg in args]...,))
    if haskey(GRAD_CACHE, cache_key)
        # if so, just use cached tape
        tape = GRAD_CACHE[cache_key]
        val = play!(tape, args...)
        return val, GradResult(tape)
    else
        # otherwise, calculate the tape
        val, g = _grad(f, args...)
        # compile it
        compile!(g.tape)
        # and put to the cache
        GRAD_CACHE[cache_key] = g.tape
        return val, g
    end
end

[...] but I'll dig into it sometime soon

No hurry! I think I'll do more benchmarks for Yota on CPU and GPU soon and cross-check them with Zygote to detect any unexpected performance issues.