야코비안: 벡터함수의 미분은 행렬이다

신경망에서 가장 흔한 한 줄은 nn.Linear(784, 128)입니다. 784개짜리 벡터를 받아 128개짜리 벡터를 내놓는 층이고, MNIST 분류기라면 첫 줄에 그대로 등장합니다. 이 층의 미분은 무엇일까요.

지난 글에서 계산 그래프를 손으로 따라갈 때, 각 자리의 「국소 비율」은 전부 수 하나였습니다. 입력도 수, 출력도 수였기 때문입니다. 그런데 이 층은 784개를 받아 128개를 내놓으므로, 물어야 할 것이 784 × 128 = 100,352가지입니다 — 「$j$ 번째 입력을 흔들면 $i$ 번째 출력이 얼마나 흔들리는가」를 모든 짝에 대해 물어야 합니다.

답을 담을 그릇은 수도 벡터도 아니고 표입니다. 그 표에 이름이 있고, 이 글은 그 이름과 그것이 하는 일을 세웁니다.

모든 짝을 표 하나에 담는다

$f: \mathbb{R}^n \to \mathbb{R}^m$ 인 함수를 봅니다. 입력이 $n$ 개, 출력이 $m$ 개이므로 편도함수는 $m \times n$ 개입니다.

정의. $f$ 의 야코비안(Jacobian)은 $(i, j)$ 성분이 $\dfrac{\partial f_i}{\partial x_j}$ 인 $m \times n$ 행렬이다. $J_f(x) = \begin{bmatrix} \dfrac{\partial f_1}{\partial x_1} & \cdots & \dfrac{\partial f_1}{\partial x_n} \\[6pt] \vdots & \ddots & \vdots \\[6pt] \dfrac{\partial f_m}{\partial x_1} & \cdots & \dfrac{\partial f_m}{\partial x_n} \end{bmatrix}$

배치를 외우는 대신 한 문장으로 기억하면 됩니다 — 행은 출력, 열은 입력입니다. 그래서 모양이 (출력 개수) × (입력 개수)이고, nn.Linear(784, 128)의 야코비안은 $128 \times 784$ 입니다.

한 줄과 한 칸의 뜻이 서로 다릅니다.

• $i$ 번째 행은 출력 하나에 대한 그래디언트입니다. 지난 글에서 본 $\nabla f_i$ 를 눕혀 놓은 것입니다.
• $j$ 번째 열은 입력 하나를 흔들었을 때 출력 전체가 흔들리는 방향입니다.

작은 예로 손을 풀어 봅니다. $f(x_1, x_2) = (x_1^2 x_2,\; x_1 + 3x_2)$ 는 $\mathbb{R}^2 \to \mathbb{R}^2$ 이므로 야코비안이 $2 \times 2$ 입니다. 네 칸을 하나씩 채웁니다.

$J_f = \begin{bmatrix} \dfrac{\partial (x_1^2x_2)}{\partial x_1} & \dfrac{\partial (x_1^2x_2)}{\partial x_2} \\[8pt] \dfrac{\partial (x_1 + 3x_2)}{\partial x_1} & \dfrac{\partial (x_1 + 3x_2)}{\partial x_2} \end{bmatrix} = \begin{bmatrix} 2x_1x_2 & x_1^2 \\ 1 & 3 \end{bmatrix}$

점 $(2, 3)$ 에서는

$J_f(2,3) = \begin{bmatrix} 12 & 4 \\ 1 & 3 \end{bmatrix}$

입니다. 야코비안은 상수가 아니라 점마다 다른 행렬이라는 점을 눈여겨봅니다. 아래 행은 마침 상수지만 위 행은 $x$ 에 따라 변합니다.

야코비안은 그 점에서 함수를 대신하는 선형사상

지난 글에서 도함수를 「가장 좋은 선형근사의 기울기」로 읽었습니다. 벡터함수에서도 문장이 그대로입니다. 달라진 것은 곱셈이 행렬곱이 된 것뿐입니다.

$f(a + h) \approx f(a) + J_f(a)\, h$

오른쪽 둘째 항은 $h$ 에 대한 선형사상입니다. $h$ 는 $n$ 차원, $J_f$ 는 $m \times n$ 이므로 $J_f h$ 는 $m$ 차원 — 모양이 맞습니다. 행렬을 선형사상으로 읽는 글에서 「행렬은 벡터를 벡터로 보내는 규칙」이라고 했는데, 야코비안이 하는 일이 정확히 그것입니다. 휘어진 함수를 한 점 근처에서 곧게 편 것이 야코비안입니다.

수로 확인합니다. $a = (2,3)$, $h = (0.01,\, -0.02)$ 로 잡습니다.

$J_f(a)\,h = \begin{bmatrix} 12 & 4 \\ 1 & 3 \end{bmatrix}\begin{bmatrix} 0.01 \\ -0.02 \end{bmatrix} = \begin{bmatrix} 0.12 - 0.08 \\ 0.01 - 0.06 \end{bmatrix} = \begin{bmatrix} 0.04 \\ -0.05 \end{bmatrix}$

실제 변화는 $f(2.01,\, 2.98) - f(2,3) = (12.039498,\ 10.95) - (12,\ 11) = (0.039498,\ -0.05)$ 입니다. 첫 성분의 오차가 0.0005, 둘째는 0입니다.

그림으로 읽으면 이렇습니다. 점 $a$ 를 중심으로 아주 작은 사각형을 그리면, $f$ 는 그것을 평행사변형으로 보냅니다. 그 평행사변형의 두 변이 야코비안의 두 열이고, 사각형이 작을수록 근사가 정확해집니다.

합성하면 야코비안이 곱해진다

지난 글의 연쇄법칙이 행렬 한 줄로 다시 적힙니다. $f: \mathbb{R}^n \to \mathbb{R}^m$ 과 $g: \mathbb{R}^m \to \mathbb{R}^p$ 를 이어 붙이면

연쇄법칙(행렬 형태). $J_{g \circ f}(x) = J_g\big(f(x)\big)\; J_f(x)$

입니다. 오른쪽은 $(p \times m)(m \times n)$ 이라 결과가 $p \times n$ — 왼쪽과 모양이 맞습니다. 가운데 $m$ 이 사라지는 것이 요점입니다. 중간층의 폭은 최종 야코비안의 모양에 나타나지 않습니다.

여기서 지난 글의 「경로마다 곱하고 경로끼리 더한다」가 어디로 갔는지 확인해 둡니다. 행렬곱의 정의가

$[J_g J_f]_{ij} = \sum_{k=1}^{m} \frac{\partial g_i}{\partial u_k}\frac{\partial u_k}{\partial x_j}$

이므로, 곱은 경로 하나이고 합 기호가 경로끼리 더하는 자리입니다. 지난 글의 문장이 행렬곱 안에 통째로 들어가 있는 셈입니다.

앞의 $f$ 에 $g(u_1, u_2) = (u_1 + u_2,\; u_1 u_2)$ 를 이어 봅니다. $J_g = \begin{bmatrix} 1 & 1 \\ u_2 & u_1 \end{bmatrix}$ 이고 $u = f(2,3) = (12, 11)$ 이므로

$J_g(u)\,J_f(2,3) = \begin{bmatrix} 1 & 1 \\ 11 & 12 \end{bmatrix}\begin{bmatrix} 12 & 4 \\ 1 & 3 \end{bmatrix} = \begin{bmatrix} 13 & 7 \\ 144 & 80 \end{bmatrix}$

입니다. 합성함수를 직접 써서 미분해도 같은지 봅니다. $g \circ f = (x_1^2x_2 + x_1 + 3x_2,\; x_1^3x_2 + 3x_1^2x_2^2)$ 이고

$\frac{\partial}{\partial x_1}(x_1^2x_2 + x_1 + 3x_2) = 2x_1x_2 + 1 = 13, \qquad \frac{\partial}{\partial x_1}(x_1^3x_2 + 3x_1^2x_2^2) = 3x_1^2x_2 + 6x_1x_2^2 = 144$

로 첫 열이 맞습니다. 나머지 둘도 $x_1^2 + 3 = 7$, $x_1^3 + 6x_1^2x_2 = 80$ 으로 맞습니다.

층이 $L$ 개인 신경망이면 이 곱이 그대로 늘어납니다.

$J_{\text{전체}} = J_L \, J_{L-1} \cdots J_2 \, J_1$

지난 글에서 「비율이 곱해지니 $r^L$ 로 벌어진다」고 했던 것의 정확한 형태가 이것입니다. 곱해지는 것이 수가 아니라 행렬일 뿐이고, 크기가 벌어지는 이야기는 그대로입니다 — 다만 「크기」를 재려면 행렬의 크기를 재는 자가 필요하고, 그것이 행렬 노름 또는 특잇값입니다.

자주 쓰는 층들의 야코비안을 정리해 둡니다. 대부분 따로 계산할 것이 없습니다.

층	식	야코비안	모양
선형층	$y = Wx$	$W$	$m \times n$
편향 덧셈	$y = x + b$	$I$	$n \times n$
원소별 활성함수	$y_i = \phi(x_i)$	$\operatorname{diag}(\phi'(x))$	$n \times n$
잔차 블록	$y = x + F(x)$	$I + J_F$	$n \times n$

선형층의 야코비안이 $W$ 그 자체라는 첫 줄이 특히 중요합니다. $y_i = \sum_j W_{ij}x_j$ 이므로 $\partial y_i/\partial x_j = W_{ij}$ 이고, 그러니 야코비안을 「구할」 것이 없습니다 — 이미 메모리에 있습니다. 잔차 블록 줄의 $I$ 도 지난 글에서 「곱해지는 값에 1을 심는다」고 했던 것의 행렬판입니다.

행렬식은 부피가 몇 배가 되는가

야코비안이 정사각행렬($m = n$)이면 행렬식을 잴 수 있고, 그 값에 뜻이 있습니다.

초급의 행렬식 글에서 $2 \times 2$ 행렬식이 두 열벡터가 만드는 평행사변형의 넓이라고 했습니다. 앞 절에서 작은 사각형이 야코비안의 두 열이 만드는 평행사변형으로 간다는 것을 봤으니, 둘을 붙이면 곧바로

$|\det J_f(a)|$ 는 점 $a$ 근처에서 $f$ 가 넓이(부피)를 몇 배로 늘리는가다.

가 됩니다. 우리 예에서 $\det J_f(2,3) = 12 \times 3 - 4 \times 1 = 32$ 이므로, $(2,3)$ 근처의 작은 사각형은 32배 넓은 평행사변형이 됩니다. 뒤의 코드에서 한 변이 0.001인 사각형을 실제로 보내 32.009배를 얻습니다.

부호도 뜻이 있습니다. 음수면 방향이 뒤집혔다는 뜻이고, 0이면 그 점에서 차원이 무너졌다는 뜻입니다 — 평행사변형이 선분으로 납작해진 것이고, 그러면 그 근처에서 $f$ 는 되돌릴 수 없습니다.

밀도가 바뀌는 이유

이 사실이 확률에서 곧바로 쓰입니다. 확률변수 $X$ 를 $Y = f(X)$ 로 옮기면 밀도가 어떻게 되는가.

확률의 총량은 보존되므로 대응하는 영역의 확률이 같아야 합니다. 작은 영역 $dV$ 가 $|\det J_f|\,dV$ 로 늘어났으니, 같은 확률을 더 넓은 자리에 펴 담느라 밀도는 그만큼 묽어집니다.

$p_Y(y)\;\big|\det J_f(x)\big|\,dV = p_X(x)\,dV \quad \Longrightarrow \quad p_Y(y) = \frac{p_X(x)}{\big|\det J_f(x)\big|}$

이것을 변수변환 공식(change of variables)이라고 합니다. 한 줄로 읽으면 넓힌 만큼 묽어지고 좁힌 만큼 진해진다입니다.

가장 간단한 확인은 $X \sim \text{Uniform}(0,1)$ 에 $Y = 2X$ 를 먹이는 것입니다. $f'(x) = 2$ 이므로 밀도가 절반이 되어야 하고, 실제로 $Y$ 는 $(0,2)$ 위의 균등분포라 밀도가 $1/2$ 입니다.

극좌표도 같은 식입니다. $x = r\cos\theta$, $y = r\sin\theta$ 의 야코비안은

$J = \begin{bmatrix} \cos\theta & -r\sin\theta \\ \sin\theta & r\cos\theta \end{bmatrix}, \qquad \det J = r\cos^2\theta + r\sin^2\theta = r$

이고, 적분에서 $dx\,dy$ 를 $r\,dr\,d\theta$ 로 바꿔 적는 그 $r$ 이 바로 이 행렬식입니다. 반지름이 큰 쪽의 칸이 더 넓기 때문입니다.

가역 변환을 여러 개 이어 붙여 단순한 분포를 복잡한 분포로 만드는 생성 모델이 이 공식 위에 서 있습니다. 로그가능도가 $\log p_X(x) - \sum_k \log|\det J_k|$ 로 적히므로, 그런 모델의 설계는 대부분 행렬식을 싸게 계산할 수 있는 층을 고르는 일입니다. 삼각행렬이나 대각행렬이 되도록 층을 짜면 행렬식이 대각 성분의 곱이라 $O(n)$ 에 끝납니다.

층 하나의 야코비안은 얼마나 큰가

이제 처음의 물음으로 돌아갑니다. nn.Linear(784, 128)의 야코비안은 $128 \times 784$ 이고 성분이 100,352개인데, 앞 표에서 봤듯 그것은 $W$ 그 자체라 새로 만들 것이 없습니다.

문제는 합성입니다. 폭이 4096인 층이 100개인 모델을 생각해 봅니다. 층 하나의 야코비안은 $4096 \times 4096$ 이고 float32로 67MB입니다. 곱을 앞에서부터 차례로 쌓아 나가면

$J_{100}\,(J_{99}\,(\cdots(J_2\,J_1)))$

중간 결과가 매번 $4096 \times 4096$ 이고, 곱 한 번이 $4096^3 \approx 6.9 \times 10^{10}$ 번의 곱셈입니다. 99번 하면 $6.8 \times 10^{12}$ 번 — 순전파 한 번의 비용과 비교하면 터무니없습니다.

그런데 실제 autograd는 이 행렬을 하나도 만들지 않습니다. 손실은 수 하나이므로 우리가 최종적으로 원하는 것은 $1 \times n$ 짜리 행 하나뿐인데, 위의 곱셈은 그 행 하나를 얻으려고 $4096 \times 4096$ 짜리 표를 100번 만들었다 버린 셈입니다.

곱셈의 순서를 바꾸는 것만으로 이 비용이 사라집니다. 그 순서에 이름이 붙어 있고, 다음 글이 그것을 셉니다.

코드로 확인하기

from math import cos, sin, pi

f = lambda x1, x2: (x1 ** 2 * x2, x1 + 3 * x2)
g = lambda u1, u2: (u1 + u2, u1 * u2)
a = (2.0, 3.0)

Jf = [[2 * a[0] * a[1], a[0] ** 2], [1.0, 3.0]]      # [[12, 4], [1, 3]]
print(Jf)

# ① 선형근사 — J h 가 실제 변화를 대신한다
h = (0.01, -0.02)
approx = [Jf[i][0] * h[0] + Jf[i][1] * h[1] for i in range(2)]
fa, fb = f(*a), f(a[0] + h[0], a[1] + h[1])
print([round(v, 6) for v in approx])                  # [0.04, -0.05]
print([round(fb[i] - fa[i], 6) for i in range(2)])    # [0.039498, -0.05]

# ② 합성의 야코비안 = 야코비안의 곱
u = f(*a)
Jg = [[1.0, 1.0], [u[1], u[0]]]
prod = [[sum(Jg[i][k] * Jf[k][j] for k in range(2)) for j in range(2)] for i in range(2)]
print(prod)                                           # [[13, 7], [144, 80]]

eps = 1e-6
gf = lambda x1, x2: g(*f(x1, x2))
num = [[(gf(a[0] + eps * (j == 0), a[1] + eps * (j == 1))[i]
       - gf(a[0] - eps * (j == 0), a[1] - eps * (j == 1))[i]) / (2 * eps)
        for j in range(2)] for i in range(2)]
print([[round(v, 3) for v in row] for row in num])    # [[13.0, 7.0], [144.0, 80.0]]

# ③ 행렬식은 넓이 배율 — 작은 사각형을 실제로 보내 본다
s = 1e-3
o, p, q = f(*a), f(a[0] + s, a[1]), f(a[0], a[1] + s)
v1 = (p[0] - o[0], p[1] - o[1])
v2 = (q[0] - o[0], q[1] - o[1])
print(round(abs(v1[0] * v2[1] - v1[1] * v2[0]) / s ** 2, 3))   # 32.009
print(Jf[0][0] * Jf[1][1] - Jf[0][1] * Jf[1][0])               # 32.0

# ④ 극좌표의 행렬식은 r 이다
for r, th in [(1.0, 0.0), (2.0, pi / 3), (0.5, pi / 4)]:
    J = [[cos(th), -r * sin(th)], [sin(th), r * cos(th)]]
    print(r, round(J[0][0] * J[1][1] - J[0][1] * J[1][0], 12))
# 1.0 1.0 / 2.0 2.0 / 0.5 0.5

③이 이 글에서 가장 볼 만한 줄입니다. 함수를 직접 실행해 얻은 넓이 배율 32.009와, 편도함수 넷으로 계산한 행렬식 32가 같습니다. 사각형을 더 작게 잡으면 소수점 아래도 맞아 들어갑니다.

정리

• 야코비안은 모든 입력·출력 짝의 편도함수를 담은 $m \times n$ 행렬이다. 행은 출력, 열은 입력이라 모양이 (출력, 입력)이다.
• 한 행은 출력 하나의 그래디언트, 한 열은 입력 하나를 흔들었을 때의 출력 변화 방향이다.
• 야코비안은 그 점에서 함수를 대신하는 선형사상이다. $f(a+h) \approx f(a) + J_f(a)h$ 이고, 작은 사각형이 야코비안의 두 열이 만드는 평행사변형으로 간다.
• 합성하면 야코비안이 곱해진다 — $J_{g\circ f} = J_g J_f$. 행렬곱의 합 기호가 지난 글의 「경로끼리 더한다」이고, 가운데 차원은 사라진다.
• 선형층의 야코비안은 $W$ 그 자체, 원소별 활성함수는 $\operatorname{diag}(\phi')$, 잔차 블록은 $I + J_F$ 다. 대부분 따로 구할 것이 없다.
• $|\det J|$ 는 국소 부피 배율이다. 0이면 그 점에서 차원이 무너져 되돌릴 수 없다.
• 부피가 늘어난 만큼 밀도는 묽어진다 — $p_Y(y) = p_X(x)/|\det J_f(x)|$ 이고, 적분의 $r\,dr\,d\theta$ 도 이 공식의 한 경우다.
• 야코비안을 명시적으로 만들면 폭 4096짜리 100층에서 곱셈이 $10^{12}$ 번대가 된다. 실제 autograd는 만들지 않는다.

층 하나의 미분이 행렬이라는 사실이 이 글의 결론이고, 그 사실이 곧바로 비용 문제를 낳았습니다. 표 하나가 $4096 \times 4096$ 인데 층마다 그것을 곱해 나가면 감당이 되지 않고, 정작 우리가 원하는 것은 손실 하나에 대한 행 하나뿐입니다.

행렬곱은 결합법칙이 성립하므로 어느 쪽부터 곱해도 답은 같습니다. 그런데 비용은 같지 않습니다. 다음 글에서 그 차이를 곱셈 횟수로 세면, 딥러닝이 왜 반드시 뒤에서부터 미분하는지가 계산 한 번으로 나옵니다.

---

읽어주셔서 감사합니다. 😊