Introduction

본 포스팅에서는 Neural Network의 Gradients를 계산하는 방법을 알아보겠습니다. 해당 글은 cs224n lecture 5 material 자료를 정리한 글입니다.

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

Computing Neural Network Gradients

Vectorized Gradients

Neural Network의 gradient를 계산할 때는 single parameter(e.g., a single element in a weight matrix)과 관련하여 계산하는 것은 비효율적입니다. 실제로는 Matrix/Vector form으로 계산하지요(e.g., torch.autograd). Vector화된 gradients의 기본 building block은 Jacobian Matrix 입니다. 함수 $f:\mathbb{R}^n\rightarrow\mathbb{R}^m$을 길이가 $n$인 vector를 길이가 $m$인 vector로 매핑하는 함수라고 가정합시다. 즉, $f(x)=[f_1(x_1,\dots,x_n),f_2(x_1,\dots,x_n),\dots,f_m(x_1,\dots,x_n)]$. 이 때 Jacobian은 아래와 같은 $m\times n$ matrix입니다.

$$\cfrac{\partial f}{\partial x}=\begin{bmatrix} \frac{\partial f_1}{\partial x_1} & \cdots & \frac{\partial f_1}{\partial x_n}\\ \vdots & \ddots & \vdots\\ \frac{\partial f_m}{\partial x_1} & \cdots & \frac{\partial f_m}{\partial x_n}\\ \end{bmatrix}$$

즉, $\begin{pmatrix}\frac{\partial f}{\partial x}\end{pmatrix}_{ij}=\frac{\partial f_i}{\partial x_j}$ (just a standard non-vector derivative). Jacobian matrix는 Jacobian들을 곱하여 vector-valued function에 chain rule을 적용할 수 있기 때문에 매우 편리합니다.

단순화시켜서 생각해보겠습니다. 다음과 같이 함수 $f$와 $g$를 가지고 있다고 가정하겠습니다.

$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))]$

일반적인 chain rule을 사용하여 $g$의 derivative를 Jacobian으로 계산할 수 있습니다.

$$\cfrac{\partial g}{\partial x}=\begin{bmatrix}\frac{\partial}{\partial x}g_1(f_1(x),f_2(x))\\\frac{\partial}{\partial x}g_2(f_1(x),f_2(x))\end{bmatrix}=\begin{bmatrix}\frac{\partial g_1}{\partial f_1}\frac{\partial f_1}{\partial x}+\frac{\partial g_1}{\partial f_2}\frac{\partial f_2}{\partial x}\\\frac{\partial g_2}{\partial f_1}\frac{\partial f_1}{\partial x}+\frac{\partial g_2}{\partial f_2}\frac{\partial f_2}{\partial x}\end{bmatrix}$$

위 식은 아래의 두 Jacobian들의 곱으로 표현되지요.

$$\cfrac{\partial g}{\partial x}=\cfrac{\partial g}{\partial f}\cfrac{\partial f}{\partial x}=\begin{bmatrix}\frac{\partial g_1}{\partial f_1} & \frac{\partial g_1}{\partial f_2}\\\frac{\partial g_2}{\partial f_1}&\frac{\partial g_2}{\partial f_2}\end{bmatrix}\begin{bmatrix}\frac{\partial f_1}{\partial x}\\\frac{\partial f_2}{\partial x}\end{bmatrix}$$

Useful Identites

신경망의 기본 연산에 관련하여 Jacobian을 몇몇 간단한 함수로 계산하는 방법에 대하여 다룹니다.

Matrix times column vector with respect to the column vector

$$z=Wx,\;\text{what is }\frac{\partial z}{\partial x}\text{?}$$

Matrix에 column vector를 곱하는 경우!

가중치 행렬을 $W\in\mathbb{R}^{n\,\times\,m}$이라고 가정하겠습니다. 이 때 $z$를 m차원 vector를 입력으로 n차원 vector를 출력해주는 함수로 생각할 수 있고 Jacobian은 $n\times m$이 될 것입니다.

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

Jacobian의 각 entry $\begin{pmatrix}\frac{\partial z}{\partial x}\end{pmatrix}_{ij}$는 아래와 같이 쓸 수 있습니다.

$$\begin{pmatrix}\cfrac{\partial z}{\partial x}\end{pmatrix}_{ij}=\frac{\partial z_i}{\partial x_j}=\frac{\partial}{\partial x_j}\sum_{k=1}^{m}W_{ik}x_k=\sum_{k=1}^{m}W_{ik}\frac{\partial}{\partial x_j}x_k=W_{ij}$$

$$\because\;\frac{\partial}{\partial x_j}x_k=1\;\text{if}\;k=j\;\text{and}\;0\;\text{if otherwise}$$

Hence, $\boxed{\cfrac{\partial z}{\partial x}=W}$.

Row vector times matrix with respect to the row vector

$$z=xW,\;\text{what is }\frac{\partial z}{\partial x}\text{?}$$

row vector에 matrix를 곱하는 경우!

가중치 행렬을 $W\in\mathbb{R}^{m\,\times\,n}$이라고 가정하겠습니다. 위의 경우와 유사하게 $z$를 아래처럼 쓸 수 있습니다.

$$z_i=\sum_{k=1}^{m}x_k W_{ki}$$

위와 유사하게 Jacobian의 각 entry $\begin{pmatrix}\frac{\partial z}{\partial x}\end{pmatrix}_{ij}$를 아래와 같이 쓸 수 있습니다.

$$\begin{pmatrix}\cfrac{\partial z}{\partial x}\end{pmatrix}_{ij}=\frac{\partial z_i}{\partial x_j}=\frac{\partial}{\partial x_j}\sum_{k=1}^{m}x_kW_{ki}=\sum_{k=1}^{m}\frac{\partial}{\partial x_j}x_kW_{ki}=W_{ji}=W_{ij}^T$$

$$\because\;\frac{\partial}{\partial x_j}x_k=1\;\text{if}\;k=j\;\text{and}\;0\;\text{if otherwise}$$

Hence, $\boxed{\cfrac{\partial z}{\partial x}=W^T}$.

A vector with itself

$$z=x,\;\text{what is }\frac{\partial z}{\partial x}\text{?}$$

사실 trivial합니다. Identity겠지요.

우선 위의 식에서 $z_i=x_i$입니다. 즉,

$$\begin{pmatrix}\cfrac{\partial z}{\partial x}\end{pmatrix}_{ij}=\cfrac{\partial z_i}{\partial x_j}=\cfrac{\partial}{\partial x_j}x_i= \begin{cases} 1 & \text{if}\ i=j \\ 0 & \text{otherwise} \end{cases}$$

즉, Jacobian $\frac{\partial z}{\partial x}$은 diagonal matrix이며 대각성분이 모두 1입니다.

Hence, $\boxed{\cfrac{\partial z}{\partial x}=I}$.

Chain rule을 적용할 때 Identity matrix는 항등원이기 때문에 이 term을 대부분 생략합니다.

An elementwise function applied a vector

$$z=f(x),\;\text{what is }\frac{\partial z}{\partial x}\text{?}$$

vector에 elementwise function을 적용하는 경우! (e.g., sigmoid, log, exp, softmax)

$f$가 elementwise로 적용되기 때문에 $z_i=f(x_i)$입니다. 즉,

$$\begin{pmatrix}\cfrac{\partial z}{\partial x}\end{pmatrix}_{ij}=\cfrac{\partial z_i}{\partial x_j}=\cfrac{\partial}{\partial x_j}f(x_i)= \begin{cases} f^\prime(x_i) & \text{if}\ i=j \\ 0 & \text{otherwise} \end{cases}$$

위에서 볼 수 있듯 Jacobian $\frac{\partial z}{\partial x}$는 대각성분이 $f$의 derivative에 $x_i$를 먹인 값을 가지는 diagonal matrix입니다.

Hence, $\boxed{\cfrac{\partial z}{\partial x}=\mathrm{diag}(f^\prime(x))}$.

Diagonal matrix로 multiplication을 수행하는 것은 diagonal로 elementwise multiplication을 수행하는 것과 동일하기 때문에 chain rule을 적용할 때 이를 $\boxed{\circ f^\prime(x)}$로 적을 수 있습니다.

Matrix times column vector with respect to the matrix

$$z=Wx,\;\delta=\frac{\partial J}{\partial z}\;\text{what is }\frac{\partial J}{\partial W}=\frac{\partial J}{\partial z}\frac{\partial z}{\partial W}=\delta\frac{\partial z}{\partial W}\text{?}$$

우선 $J$라는 loss function을 가지고 있다고 가정합시다.

이 값은 scalar입니다.

loss function $J$는 행렬 $W\in \mathbb{R}^{n\,\times\,m}$의 gradients를 계산합니다. $J$를 input $nm$차원의 vector를 받고 ($W$의 entries) single output ($J$)를 출력하는 $W$의 함수라고 생각할 수 있습니다. 이는 Jacobian $\frac{\partial J}{\partial W}$이 $1\,\times\,nm$ vector임을 의미합니다. 하지만 이는 실용적으로 효과적이지 않고 아래와 같이 $n\,\times\,m$ 차원의 행렬로 derivative를 더 nice하게 쓸 수 있습니다.

$$\cfrac{\partial J}{\partial W}=\begin{bmatrix} \frac{\partial J}{\partial W_{11}} & \cdots & \frac{\partial J}{\partial W_{1m}} \\ \vdots & \ddots & \vdots \\ \frac{\partial J}{\partial W_{n1}} & \cdots & \frac{\partial J}{\partial W_{nm}} \\ \end{bmatrix}$$

위 행렬은 $W$와 같은 shape을 가지기 때문에 gradient descent를 수행할 때 $W$에서 learning rate를 곱하고 빼면 됩니다.

위 방식(gradient를 정렬하는 방법)은 $\frac{\partial z}{\partial W}$을 계산할 때 복잡해집니다. $J$는 scalar이기 때문에 matrix-form으로 jacobian을 쓸 수 있지만 $z$는 vector입니다. 이를 gradient arranging을 수행하면 $\frac{\partial z}{\partial W}$은 $n\,\times\,m\,\times\,n$ tensor가 될 것입니다. 운좋게도 이 문제를 single weight $W_{ij}$로 gradient를 계산하여 회피할 수 있습니다. $\frac{\partial z}{\partial W_{ij}}$은 vector기에 처리하기가 더 쉽습니다.

$$z_k=\sum_{t=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}$$

$$\cfrac{\partial}{\partial W_{ij}}W_{kl}=\begin{cases} 1 & \text{if}\; i=k \,\text{and}\, j=l \\ 0 & \text{otherwise} \end{cases}$$

마지막 식에서 $k\neq i$이면 $\frac{\partial}{\partial W_{ij}}W_{kl}$은 0이 되고 그 합도 0, 최종적으로 gradient $\frac{\partial z_k}{\partial W_{ij}}$도 0이 됩니다. 반대로 summation의 유일한 nonzero element는 $l=j$일 때이며 해당 값은 $x_j$입니다.

partial term이 1이 되고 나머지 term은 0이기 때문에!

고로 위에서 $\frac{\partial z_k}{\partial W_{ij}}=x_j\;\text{if}\;k=i\;\text{and}\;0\;\text{otherwise}$임을 알 수 있고 아래처럼 쓸 수도 있습니다.

$$\cfrac{\partial z}{\partial W_{ij}}=\begin{bmatrix} 0 \\ \vdots \\ 0 \\ x_j \\ 0 \\ \vdots \\ 0 \end{bmatrix} \leftarrow i\text{th element}$$

이제 $\frac{\partial J}{\partial W_{ij}}$를 계산해보면,

$$\cfrac{\partial J}{\partial W_{ij}}=\cfrac{\partial J}{\partial z}\cfrac{\partial z}{\partial W_{ij}}=\delta \cfrac{\partial z}{\partial W_{ij}}=\sum_{k=1}^{m}\delta_k \cfrac{\partial z_k}{\partial W_{ij}}=\delta_i x_j$$

The only nonzero term in the sum is $\delta_i \frac{\partial z_i}{\partial W_{ij}}$

위로부터 아래의 결론을 도출할 수 있다.

Hence, $\boxed{\cfrac{\partial J}{\partial W}=\delta x^T}$.

Row vector time matrix with respect to the matrix

$$z=xW,\;\delta=\frac{\partial J}{\partial z}\;\text{what is }\frac{\partial J}{\partial W}=\delta\frac{\partial z}{\partial W}\text{?}$$

Cross-entropy loss with respect to logits

$$\hat{y}=\mathrm{softmax}(\theta),\;J=CE(y,\hat{y}),\;\text{what is }\frac{\partial J}{\partial \theta}\text{?}$$

Gradient Layout

Example: 1-Layer Neural Network

import random
import numpy as np
import torch


def set_seed(seed: int = 42):
    """Seed fixer (random, numpy, torch)
    Args:
        seed (:obj:`int`): The seed to set.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)  # if use multi-GPU
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    
    
set_seed()

x = torch.randn(1, 2).requires_grad_()

f = torch.nn.Linear(2, 3, bias=False).requires_grad_()

y = f(x)
torch.autograd.backward(y, torch.ones_like(y), retain_graph=True)

# nn.Linear에서 Weight를 transpose로 저장 
torch.equal(y, x @ f.weight.T)

True

torch.autograd.grad(y, x, torch.ones_like(y), retain_graph=True)[0]

tensor([[0.1242, 0.0392]])

torch.autograd.grad(y, f.weight, torch.ones_like(y), retain_graph=True)[0]

tensor([[0.3367, 0.1288],
        [0.3367, 0.1288],
        [0.3367, 0.1288]])

x.T @ torch.ones_like(y) # row vector case

tensor([[0.3367, 0.3367, 0.3367],
        [0.1288, 0.1288, 0.1288]], grad_fn=<MmBackward>)

(x.T @ torch.ones_like(y)).T

tensor([[0.3367, 0.1288],
        [0.3367, 0.1288],
        [0.3367, 0.1288]], grad_fn=<PermuteBackward>)

x = torch.randn(5, 1).requires_grad_()
W = torch.randn(3, 5).requires_grad_()

y = W @ x

torch.autograd.grad(y, W, torch.ones_like(y), retain_graph=True)[0]

tensor([[-0.1731, -0.8167,  0.1210, -0.7272,  1.9706],
        [-0.1731, -0.8167,  0.1210, -0.7272,  1.9706],
        [-0.1731, -0.8167,  0.1210, -0.7272,  1.9706]])

torch.ones_like(y) @ x.T # column vector case

tensor([[-0.1731, -0.8167,  0.1210, -0.7272,  1.9706],
        [-0.1731, -0.8167,  0.1210, -0.7272,  1.9706],
        [-0.1731, -0.8167,  0.1210, -0.7272,  1.9706]], grad_fn=<MmBackward>)