Gradients of Neural Networks

 

Computing Neural Networks Gradients

Vectorized Gradients

함수 $f$가 $f: \mathbb{R}^n \to \mathbb{R}^m$ 즉 길이가 $n$인 벡터를 길이가 $m$인 벡터로 매핑할 때 Jacobian은 다음과 같이 행렬의 형태로 표현할 수 있다.

Jacobian Matrix of $f$

즉, Jacobian matrix의 $(i,j)$는 $\left( \frac{\partial f}{\partial x} \right)_{ij}=\frac{\partial f_i}{\partial x_j}$ 이다. Jacobian matrix를 사용할 때의 이점은, chain rule을 이용할 때 단순히 Jacobian의 곱하기로 표현할 수 있다는 것이다.

 

예를 들어, $f(x)=[f_1(x), f_2(x)]$이고, $g(y)=[g_1(y_1, y_2), g_2(y_1, y_2)]$라 하자. 그리고 $f$와 $g$를 합성하여 $g(x)=[g_1(f_1(x), f_2(x)), g_2(f_1(x), f_2(x))]$일 때의 $g$의 Jacobian을 구하면 다음과 같다,

Jacobian Matrix of $g$

이는 2개의 Jacobian matrix를 곱한 것과 동일한 것을 알 수 있다.

Multiplying the two Jacobians

 

Useful Identities

(1) $z=x$일 때, $\cfrac{\partial z}{\partial x}$ ?

$z_i = x_i$이므로 $(\cfrac{\partial z}{\partial x})_{ij}=\cfrac{\partial z_i}{\partial x_j}=\mathbf{1}_{\{i=j\}}$ 따라서

\[ \cfrac{\partial z}{\partial x}=I \]

 

(2) $z=f(x)$일 때, $\cfrac{\partial z}{\partial x}$ ?

$z_i = f(x_i)$이므로 $\cfrac{\partial z_i}{\partial x_j}=\cfrac{\partial}{\partial x_j}f(x_i)=f'(x_i)_{\{i=j\}}$. 따라서

\[ \cfrac{\partial z}{\partial x} = \text{diag}\left( f'(x) \right) \]

 

(3) $z=Wx$일 때, $\cfrac{\partial z}{\partial x}$ ?

$W \in \mathbb{R}^{n \times m}$ 이고 $x \in \mathbb{R}^{m}$인 column vector라 하자. Jacobian은 당연히 $n \times m$이 된다. 그리고 

\[ z_i = \sum_{k=1}^{m} W_{ik}x_k \]

이므로 

\[ \cfrac{\partial z_i}{\partial x_j} = \cfrac{\partial}{\partial x_j} \sum_{k=1}^{m} W_{ik}x_k = \sum_{k=1}^{m}W_{ik}\cfrac{\partial}{\partial x_j}x_k = W_{ij} \]

따라서

\[ \cfrac{\partial z}{\partial x} = W \]

 

(4) $z=Wx$, $\delta = \frac{\partial J}{\partial z}$일 때, $\frac{\partial J}{\partial W}$ ?

$\cfrac{\partial J}{\partial W} = \cfrac{\partial J}{\partial z} \cfrac{\partial z}{\partial W} = \delta \cfrac{\partial z}{\partial W}$이므로, $\cfrac{\partial z}{\partial W}$를 구하면 된다.

$z$가 벡터이기 때문에 $\cfrac{\partial z}{\partial W}$는 $n \times m \times n$ 텐서가 된다.

\[z_k = \sum_{l=1}^{m} W_{kl}x_l, \ \cfrac{\partial z_k}{\partial W_{ij}} = \sum_{l=1}^{m} x_l \cfrac{\partial}{\partial W_{ij}}W_{kl} \]

$i=k$ 이고 $j=l$이면 $\frac{\partial}{\partial W_{ij}}W_{kl}=1$이고 그 외에는 $0$이다. 즉

따라서 

\[ \cfrac{\partial J}{\partial W_{ij}} = \cfrac{\partial J}{\partial z} \cfrac{\partial z}{\partial W_{ij}} = \delta_i x_j \]

\[ \cfrac{\partial J}{\partial W} = \delta^\top x^\top \]

 

(5) Cross-entropy loss w.r.t logits

$\hat{y} = \text{softmax}(\theta)$이고, $J=CE(y, \hat{y})$일 때,

\[ \cfrac{\partial J}{\partial \theta} = \hat{y} - y \]

 

Example: 1-Layer Neural Network

다음과 같은 구조를 같는 Neural Network를 생각해보자.

\begin{align} x &= \text{input} \\ z &= Wx+b_1 \\ h &= \text{ReLU}(z) \\ \theta &= Uh + b_2 \\ \hat{y} &= \text{softmax}(\theta) \\ J &= CE(y, \hat{y}) \end{align}

 

우리가 계산할 gradient는 다음과 같다.

\[ \cfrac{\partial J}{\partial U},\ \cfrac{\partial J}{\partial b_2},\ \cfrac{\partial J}{\partial W},\ \cfrac{\partial J}{\partial b_1},\ \cfrac{\partial J}{\partial x} \]

그리고 $\text{ReLU}'(x)=\text{sgn}(\text{ReLU}(x))$ 라 하자.

 

$\delta_1 = \cfrac{\partial J}{\partial \theta},\ \delta_2=\cfrac{\partial J}{\partial z}$라 하자. (2번 이상 반복하여 계산)

이들을 먼저 계산하면

\begin{align} \delta_1 = \cfrac{\partial J}{\partial \theta} = (\hat{y} - y)^T \end{align}

\begin{align} &= \cfrac{\partial J}{\partial z} = \cfrac{\partial J}{\partial \theta} \cfrac{\partial \theta}{\partial h} \cfrac{\partial h}{\partial z} \\ &= \delta_1 \cfrac{\partial \theta}{\partial h} \cfrac{\partial h}{\partial z} \\ &= \delta_1 U \cfrac{\partial h}{\partial z} \\ &= \delta_1 U \text{ReLU'}(z) \\ &= \delta_1 U \text{sgn}(h) \end{align}

 

나머지 gradient는 다음과 같이 계산할 수 있다.

Gradients for column vectors terms

References

https://web.stanford.edu/class/cs224n/readings/gradient-notes.pdf