Julia: Custom Backpropagation on Zygote

Zygote is a package in Julia that realizes automatic differentiation and automatic derivation, and @adjointmacros are an important part of the Zygote interface. @adjointThe backpropagation of the function can be customized using .

Pullbacks

To understand, @adjointyou must first understand the lower-level functions pullback. gradientIt's actually pullbacksyntactic sugar for .

julia> y, back = Zygote.pullback(sin, 0.5)
(0.479425538604203, Zygote.var"#41#42"{Zygote.ZBack{ChainRules.var"#sin_pullback#1430"{Float64}}}(Zygote.ZBack{ChainRules.var"#sin_pullback#1430"{Float64}}(ChainRules.var"#sin_pullback#1430"{Float64}(0.8775825618903728))))

julia> y
0.479425538604203

Give pullbackthe input two parameters sinand 0.5represent the function to be derived and the value to be derived respectively, and two outputs will be obtained: the result of the given function sin(0.5)and one pullback, which is the variable in the above code back. backGradient calculation is performed on the function sin, accepting a derivation, and generating a new variable. . Mathematically, it is the realization of the vector-Jacobian product. where $y=f\left(x\right)$ and the gradient$\frac{\partial l}{\partial x}$written as $\bar{x}$，pullback ${\mathcal{B}}_y$ 如下计算：
$$\bar{x}=\frac{\partial l}{\partial x}=\frac{\partial l}{\partial y}\frac{\partial y}{\partial x}={\mathcal{B}}_y(\bar{y})$$
More specifically, taking the above code as an example, the function$y=\sin (x)$. $\frac{\partial y}{\partial x}=\cos \left(x\right)$ , so pullback is$\bar{y}\cos \left(x\right)$，其中$\bar{y}=\frac{\partial l}{\partial y}$. In other words, equivalent pullback(sin, x)to dsin(x) = (sin(x), ȳ -> (ȳ * cos(x),)).

gradientIn the function $l=f\left(x\right)$ and assume$\bar{l}=\frac{\partial l}{\partial l}=1$ , and feed it into the pullback. Insinthe example of ,

julia> dsin(x) = (sin, ȳ -> (ȳ * cos(x),))
dsin (generic function with 1 method)

julia> function gradsin(x)
           _, back = dsin(x)
           back(1)
       end
gradsin (generic function with 1 method)

julia> gradsin(0.5)
(0.8775825618903728,)

julia> cos(0.5)
0.8775825618903728
                
julia> back(1)
(0.8775825618903728,)

Personal understanding, why add an item $\frac{\partial l}{\partial y}$, which is to implement the chain rule. For example, suppose the final loss is $l$ , the function$y\left(x\right)$ , to get the loss function$l$ for parameterxxDifferentiation of$x$$\frac{\partial l}{\partial x}$, according to the chain rule is the loss function to the function $Differentiation\; of\; y$ multiplication function with respect to parameter$x$ 的微分，即 $\frac{\partial l}{\partial y}\frac{\partial y}{\partial x}$. function $y$ ispullbackthe loss function to the function$Differentiation\; of\; y$ (with$\bar{y}$means) multiplied by the function pair $Differentiation\; of\; x$ .

For the above example, pullbackthe first result returned by the function is: Suppose the function $y=\sin \left(x\right)$ is the loss function$When\; l$ ,$x=$The result at$0.5$$,$$that$$\cos (0.5)$ , and the returnedbackis a$\frac{\partial l}{\partial y}$function, which can be regarded as $\mathcal{B}(\frac{\partial l}{\partial y})=\frac{\partial l}{\partial y}\cos (0.5)$。

If $l=0.5y=0.5\sin \left(x\right)$ , we can get$\frac{\partial l}{\partial y}=0.5$，那么 $\frac{\partial l}{\partial x}=\mathcal{B}(\frac{\partial l}{\partial y})=\mathcal{B}(0.5)$。

---

reference:

[1] Custom Adjoints • Zygote