链式编程_微分编程（三）：可逆微分-贝塞尔函数的命运

这个章节是微分编程的最后一个章节，前面两个章节细数了微分编程的“三宗罪”，可逆编程，可逆编程那部分把很多干巴巴的理论甚至硬件都扯完了，所以今天就可以很轻松的从例子出发讲讲几种AD机制的不一样的地方，为什么会有这么大的性能差异。今天会涉及的几种AD机制都是源到源的AD，包括基于checkpointing的adjoint mode AD，基于可逆编程的adjoint mode AD和tangent mode AD三种。

相信很多人都知道第一类贝塞尔函数，它在物理里面是重要的函数。一个

阶贝塞尔函数可以用泰勒展开来求解

其中Gamma函数

。本文将会给这个函数的多种自动微分实现进行benchmark，最终汇总为下表

性能测试总结表

表 1: 不同自动微分机制下贝塞尔函数求导的 benchmark,括号内的数字为可逆性检查开关关闭的结果，以多次测试的最少时间为准。

其中第一行Julia的不带AD的forward实现是基准，NiLang是可逆编程嵌入式语言（eDSL），NiLang.AD是该eDSL上的Adjoint mode AD的实现。ForwardDiff是Julia下Tangent mode AD的实现。Zygote是Julia下源到源的Adjoint mode AD的实现。Manual是手动绑定（或者说operator overloading）的Adjoint mode AD，定义见文中。

---

正文

1. 基于计算图的AD：Zygote

上述贝塞尔函数的表达式是个求和，每个求和项可以表达为

。我们很容易写下如下的Julia实现

function besselj(ν, z; atol=1e-8)
    k = 0
    s = (z/2)^ν / factorial(ν)
    out = s
    while abs(s) > atol
        k += 1
        s *= (-1) / k / (k+ν) * (z/2)^2
        out += s
    end
    out
end

如果调用了自动微分库Julia的源到源自动微分库Zygote，它会劫持了这段代码的SSA表达式，并把它表达成了计算图。这个函数的计算图被画在了图一，它把数据转换成了箭头，函数转换成了节点，并展开了流控制令计算图变得很长。它没有内存更没有inplace函数的概念，像极了数学表达式，所以很容易通过链式法则来求导。

图一：第一类贝塞尔函数的计算图。带箭头的线表示数据，圆圈代表函数，灰色区域是被展开的循环体。

不费吹灰之力，就可以通过如下短短一行代码得到第一类贝塞尔函数的导数

julia> using Zygote
julia> Zygote.gradient(besselj, 2, 1.0)
(-0.07635849493801533, 0.21024361585183118)

源到源自动微分真是太有趣了，它继tensorflow，pytorch这些基于operator overloading的机制之后，让自动微分变得更加自动了，大大减少获得导数的劳动力成本。同样的源到源自动微分还有Tapenade等。正当我们感慨源到源的美妙时，我们还关心的问题是，它性能可以吗？于是我们做了个数值小实验，

julia> using BenchmarkTools

julia> @benchmark besselj(2, 1.0)
BenchmarkTools.Trial: 
  memory estimate:  0 bytes
  allocs estimate:  0
  --------------
  minimum time:     23.363 ns (0.00% GC)
  median time:      23.731 ns (0.00% GC)
  mean time:        26.207 ns (0.00% GC)
  maximum time:     89.767 ns (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     996


julia> @benchmark gradient(besselj, 2, 1.0)
BenchmarkTools.Trial: 
  memory estimate:  13.47 KiB
  allocs estimate:  411
  --------------
  minimum time:     37.788 μs (0.00% GC)
  median time:      38.771 μs (0.00% GC)
  mean time:        44.812 μs (5.75% GC)
  maximum time:     9.084 ms (98.76% GC)
  --------------
  samples:          10000
  evals/sample:     1
end{
    lstlisting}
end{
    minipage}

我们发现Zygote比前向计算慢了1500倍以上，不仅如此，还多了13KB的heap上的内存分配，而前向过程根本没有heap上的内存分配。造成这个结果的原因有两部分，一部分是Zygote过于依赖Julia的Compiler去推导类型并去除运行时的overhead，但是Julia的type inference并没有做到完美。另一个比较本质的原因是，这样的自动微分用了naive的checkpointing，造成了很多内存分配，虽然自动了，但是性