4

自动微分(Automatic Differentiation):算法篇

 1 year ago
source link: https://lotabout.me/2023/Auto-Differentiation-Part-1-Algorithm/
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.

自动微分(Automatic Differentiation,下面简称 AD)是用来计算偏导的一种手段,在深度学习框架中广泛使用(如 Pytorh, Tensorflow)。最近想学习这些框架的实现,先从 AD 入手,框架的具体实现比较复杂,我们主要是理解 AD 的思想并做个简单的实现。

AD 能干什么?

AD 能用来求偏导的。

例如有一个 R2↦RR2↦R 的函数(函数有 2 个输入,1 个输出):f(x,y)f(x,y) ,对于 xx、yy 的偏导分别计为 ∂f∂x∂f∂x 和 ∂f∂y∂f∂y。通常我们不关心偏导的解析式,只关心具体某个 xixi, yiyi 取值下偏导 ∂f∂x|x=xi,y=yi∂f∂x|x=xi,y=yi 和 ∂f∂y|x=xi,y=yi∂f∂y|x=xi,y=yi 的值。

另外注意在神经网络在使用“梯度下降”学习时,我们关心的是“参数 ww”的偏导。而不是“输入 xx”的偏导。假设有 f(x)=ax2+bf(x)=ax2+b 这样的神经网络,损失函数是 l(f(x),y)l(f(x),y),现在给了一个样本标签对(x0,y0)(x0,y0),我们要计算的是 ∂l∂a|x=x0,y=y0,a=a0,b=b0∂l∂a|x=x0,y=y0,a=a0,b=b0 和 ∂l∂b|x=x0,y=y0,a=a0,b=b0∂l∂b|x=x0,y=y0,a=a0,b=b0。在对号入座时要牢记这点。

为什么用 AD?

求偏导有很多做法,例如 symbolic differentiation 使用“符号计算” 得到准确的偏导解析式,但对于复杂的函数,偏导解析式会特别复杂,占用大量内存且计算慢,并且通常应用也不需要解析式;再比如 numerical differentiation 通过引入很小的位移 hh,计算 f(x+h)−f(h)hf(x+h)−f(h)h 得到偏导,这种方法编码容易,但受 float 误差影响大,且计算慢(有几个输入就要算几次 ff)。

AD 认为所有的计算最终都可以拆解成基础操作(如加减乘除,exp, log, sin, cos 等基本函数)的组合。然后通过链式法则 逐步计算偏导。这样使用方只需要正常组合基础操作,就能自动计算偏导,且不受 float 误差的影响,还可以复用一些中间结果来减少计算量(等价于动态规划)。

链式法则回顾

AD 的数学基础就是链式法则(chain rule)

对于函数 z=h(x)z=h(x),如果有子函数 y=f(x)y=f(x),满足 z=h(x)=g(y)=g(f(x))z=h(x)=g(y)=g(f(x)),则求偏导有如下关系:

h′(x)=g′(f(x))f′(x)⟺∂z∂x∣∣∣x0=∂z∂y∣∣∣y=f(x0)∂y∂x∣∣∣x0h′(x)=g′(f(x))f′(x)⟺∂z∂x|x0=∂z∂y|y=f(x0)∂y∂x|x0

上述两种写法是一致的。另外如果涉及多个变量,例如 z=f(x,y)z=f(x,y),而 x=g(t),y=h(t)x=g(t),y=h(t),则有:

∂z∂t=∂z∂x∂x∂t+∂z∂y∂y∂t∂z∂t=∂z∂x∂x∂t+∂z∂y∂y∂t

这里之所以成立,应该是因为 x,yx,y 是独立的(没有深究)。

AD 具体是怎么做的?

AD 其实就是链式法则的具体实现。它有两种模式:前向模式(Forward accumulation)和反向模式(Reverse accumulation),我们只考虑反向模式。那么具体是怎么工作的呢?考虑下面的复杂函数[1]

y=f(x1,x2)=sinx1+x1x2=sinv1+v1v2=v3+v4=v5y=f(x1,x2)=sin⁡x1+x1x2=sin⁡v1+v1v2=v3+v4=v5

上述公式中,我们用了一些子函数来简化整个函数,画成图如下左图:

2023-04-AD-example-computation-graph.svg

于是为了求偏导 ∂f∂x1∂f∂x1 与 ∂f∂x2∂f∂x2 的值,我们可以先定义中间值 vi¯=∂f∂vivi¯=∂f∂vi,根据链式法则,有

vi¯=∂f∂vi=∂f∂vi+1∂vi+1∂vi=vi+1¯∂vi+1∂vivi¯=∂f∂vi=∂f∂vi+1∂vi+1∂vi=vi+1¯∂vi+1∂vi

于是计算时需要先“前向”计算一次,得到 v1,v2,⋯,v5v1,v2,⋯,v5 的值,之后再“后向”计算 v5¯,v4¯,⋯,v1¯v5¯,v4¯,⋯,v1¯ 的值(参考上右图),最终得到的 v1¯,v2¯v1¯,v2¯ 就是我们要计算的结果。而需要先“前向”计算一次,是因为后向计算时会用到前向的值,例如 v2¯=v4¯v1v2¯=v4¯v1 就需要用到前向的v1v1。

注意图里 v1¯v1¯ 的计算依赖了链式法则中多变量的情况,等于它所有后继节点偏导(即图中的 va1¯,vb1¯v1a¯,v1b¯)的和。

多输出情形

多输出的情况偏理论,跳过也影响不大。神经网络的输出,在训练时最终都会接入损失函数,得到 loss 值,一般都是一个标量,可以认为神经网络的学习总是单输出的。

在多输出的情况下,链式法则依然生效。

刚才都假设函数是 Rn↦RRn↦R,即 n 个输入,1 个输出。考虑 m 个输出,即 Rn↦RmRn↦Rm 的情况。假设输入是 x1,x2,⋯,xnx1,x2,⋯,xn,而输出是 f1(x1,⋯,xn),f2(x1,⋯,xn),⋯,fm(x1,⋯,xn)f1(x1,⋯,xn),f2(x1,⋯,xn),⋯,fm(x1,⋯,xn)。此时我们要计算的偏导就不是 n 个值了,而是一个 m×n 的矩阵[2],每个元素 Jij=∂fi∂xjJij=∂fi∂xj。这个矩阵一般称为 Jacobian Matrix

Jm×n=[∂f∂x1⋯∂f∂xn]=⎡⎣⎢⎢∇Tf1⋮∇Tfm⎤⎦⎥⎥=⎡⎣⎢⎢⎢⎢⎢⎢⎢∂f1∂x1⋮∂fm∂x1⋯⋱⋯∂f1∂xn⋮∂fm∂xn⎤⎦⎥⎥⎥⎥⎥⎥⎥Jm×n=[∂f∂x1⋯∂f∂xn]=[∇Tf1⋮∇Tfm]=[∂f1∂x1⋯∂f1∂xn⋮⋱⋮∂fm∂x1⋯∂fm∂xn]

其中 ∇Tfi∇Tfi 代表 fifi 对于所有输入的偏导(行向量)的转置。

考虑函数 g:Rn↦Rkg:Rn↦Rk,h:Rk↦Rmh:Rk↦Rm,而函数 ff 是二者的组合: f(x)=h∘g(x)=h(g(x))f(x)=h∘g(x)=h(g(x)),则有

J=Jh∘g=Jh(g(x))⋅Jg(x)J=Jh∘g=Jh(g(x))⋅Jg(x)

此时 JJ 中的每个元素:

Jij=∂fi∂xj=∑l=1k∂hi∂gl∂gl∂xj=[∂hi∂g1⋯∂hi∂gk]⎡⎣⎢⎢⎢⎢⎢⎢⎢∂g1∂xj⋮∂gk∂xj⎤⎦⎥⎥⎥⎥⎥⎥⎥Jij=∂fi∂xj=∑l=1k∂hi∂gl∂gl∂xj=[∂hi∂g1⋯∂hi∂gk][∂g1∂xj⋮∂gk∂xj]

可以看到和 Jh⋅JgJh⋅Jg 的结果是一致的。不过这些性质其实都是链式法则的内容,这里也只是扩充视野。

AD 把复杂的函数看成是许多小函数的组合,再利用链式法则来计算偏导。它有不同的模式,其中“后向模式”在计算偏导时先“前向”计算得到一些中间结果,之后再“反向”计算偏导。从工程的视角看,由于中间的偏导可以重复利用,能减少许多计算量。深度学习的反向传播算法(BP)是 AD 的一种特例。

所以回过头来,什么是 AD?AD 就是利用链式法则算偏导的一种实现。


  1. 例子取自维基百科,修改了其中的符号

  2. m×n 还是 n×m 取决于是行矩阵还是列矩阵,其实关系不大。


Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK