Chain rule and Jacobians

Chain rule

• 1-var function:
  • $z=3y$
  • $y=x^2$

$$\frac{dz}{dx}=\frac{dz}{dy}\cdot \frac{dy}{dx}=3\cdot 2x=6x$$

• multi-var function:
  • $h=f(z)$
  • $z=W(x)+b$

$$\frac{\partial h}{\partial x}=\frac{\partial h}{\partial z}\cdot \frac{\partial z}{\partial x}$$

Every variable depends on the chain rule.

Example: if $h=f(z)$ and $h_i=f(z_i)$, what is $\frac{\partial h}{\partial z}$?

$$\frac{\partial h}{\partial z}=\left(\begin{array}{cccc}{f}^{\prime }(z_1)& 0& \cdots & 0\\ 0& f^{\prime }(z_2)& \cdots & 0\\ ⋮& ⋮& \ddots & ⋮\\ 0& 0& \cdots & f^{\prime }(z_n)\end{array}\right)=\mathrm{diag}(f^{\prime }(z))$$

Jacobians

$$\frac{\partial f}{\partial x}=\left[\begin{array}{cccc}\frac{\partial f_1}{\partial x_1}& \frac{\partial f_1}{\partial x_2}& \cdots & \frac{\partial f_1}{\partial x_n}\\ \frac{\partial f_2}{\partial x_1}& \frac{\partial f_2}{\partial x_2}& \cdots & \frac{\partial f_2}{\partial x_n}\\ ⋮& ⋮& \ddots & ⋮\\ \frac{\partial f_m}{\partial x_1}& \frac{\partial f_m}{\partial x_2}& \cdots & \frac{\partial f_m}{\partial x_n}\end{array}\right]=(\frac{\partial f}{\partial x}{)}_{ij}=\frac{\partial f_i}{\partial xj}$$

The jacobians are:

• $\frac{\partial }{\partial u}(U^{\top }h)=h$

$$\frac{\partial }{\partial z}f(z)=\mathrm{diag}(f^{\prime }(z))=\left[\begin{array}{cccc}{f}^{\prime }(z_1)& 0& \cdots & 0\\ 0& f^{\prime }(z_2)& \cdots & 0\\ ⋮& ⋮& \ddots & ⋮\\ 0& 0& \cdots & f^{\prime }(z_n)\end{array}\right]$$

• $\frac{\partial }{\partial b}(Wx+b)=I$

Exercise

Given the functions:

• $s=U^{\top }h$
• $h=f(Wx+b)$
• $x$

$s$ is the score, h is the the function that transforms input $x$, with $W$ weights and $f$ is the activation

activation means if we need to activate that neuron

---

Step 1: Rewrite

Express $z$ explicitly and rewrite h:

• $z=Wx+b$
• $h=f(z)$
• $x$
• $s=U^{\top }h$

---

Step 2: Apply the chain rule

We want to find the derivative of $s$ with respect to $x$:

$$\frac{\partial s}{\partial x}=\frac{\partial s}{\partial h}\cdot \frac{\partial h}{\partial z}\cdot \frac{\partial z}{\partial x}$$

Calculate each term:

• $\frac{\partial s}{\partial h}=U^{\top }$
• $\frac{\partial h}{\partial z}=\mathrm{diag}\left(f^{\prime }(z)\right)$, a diagonal matrix with derivatives of $f$ applied element-wise
• $\frac{\partial z}{\partial x}=W$

Putting it all together:

$$\frac{\partial s}{\partial x}=U^{\top }\cdot \mathrm{diag}\left(f^{\prime }(z)\right)\cdot W$$