자연 그래디언트

학습률은 한 스텝에 파라미터를 얼마나 옮길지 정하는 숫자인데, 이 숫자는 파라미터 좌표의 칸 수로 적혀 있다. 좌표의 한 칸을 움직였을 때 모델이 내는 분포가 얼마나 달라지는지는 모델마다, 자리마다 다르다. 그래서 같은 학습률이 어떤 모델에서는 알맞고 어떤 모델에서는 지나치다. GPT-3 논문(브라운(Tom Brown) 외 2020)의 표에서 폭(한 층의 뉴런 수) 768 모델은 학습률 6.0 × 10⁻⁴ 로, 폭 12288 모델은 0.6 × 10⁻⁴ 로 학습했다.

물음은 이것이다. 한 스텝의 크기를 좌표의 칸 수 대신, 분포가 달라지는 정도로 처음부터 잴 수는 없는가?

이 장에서 학습률(learning rate)은 lr 로 쓴다(흔히 쓰는 η 는 이 책에서 다른 양의 이름이다).

자연 그래디언트: 피셔 계량을 따라 걷기

같은 학습률이 폭에 따라 왜 다른 뜻을 갖는지, 손으로 끝까지 셀 수 있는 가장 작은 예부터 보자. 폭 n 인 층 하나가 출력 y = w·x 하나를 내고, 입력 x 의 좌표는 ±1 이며, 손실 ½(y − t)² 의 그래디언트는 (y − t)x 이다. 경사하강 한 스텝 뒤 출력의 변화는

\Delta \textcolor{#506890}{y} = -\,\textcolor{#b88068}{\mathrm{lr}}\,(\textcolor{#506890}{y} - \textcolor{#885868}{t})\,\|\textcolor{#6f8fa6}{x}\|^2 = -\,\textcolor{#b88068}{\mathrm{lr}}\,(\textcolor{#506890}{y} - \textcolor{#885868}{t})\, \textcolor{#a05000}{n}

\begin{array}{ll} \Delta \textcolor{#506890}{y} & \text{한 스텝 뒤 출력의 변화} \\ \textcolor{#b88068}{\mathrm{lr}} & \text{학습률} \\ \textcolor{#506890}{y},\ \textcolor{#885868}{t} & \text{출력과 목표값} \\ \textcolor{#6f8fa6}{x} & \text{입력 벡터 (좌표마다 크기 1)} \\ \textcolor{#a05000}{n} & \text{폭 (입력 차원)} \end{array}

그래디언트의 각 좌표는 작아도, 그 변화가 입력과 같은 방향으로 늘어서서 n 개가 한꺼번에 더해지기 때문에 출력 변화가 n 에 비례한다. 씨앗을 고정해 네 폭에서 한 스텝씩 걸어 보면 이렇다.

import numpy as np
rng = np.random.default_rng(0)
for n in [64, 256, 1024, 4096]:
    x = rng.choice([-1.0, 1.0], n)              # 입력: 좌표마다 크기 1
    w = rng.standard_normal(n) / np.sqrt(n)     # 표준 초기화
    y, target, lr = w @ x, 1.0, 0.01
    g = (y - target) * x                        # ½(y−t)² 의 그래디언트
    y_new = (w - lr * g) @ x
    print(f"n={n:5d}  출력 변화 {y_new - y:+.3f}  = −lr·(y−t)·n = {-lr * (y - target) * n:+.3f}")
# n=   64  출력 변화 -0.707  = −lr·(y−t)·n = -0.707
# n=  256  출력 변화 +2.189  = −lr·(y−t)·n = +2.189
# n= 1024  출력 변화 +9.973  = −lr·(y−t)·n = +9.973
# n= 4096  출력 변화 +38.057  = −lr·(y−t)·n = +38.057

같은 lr = 0.01 이 n = 64 에서는 목표 쪽으로 알맞게 가고, n = 4096 에서는 38만큼 튀어 나간다. 폭을 바꿀 때마다 학습률을 1/n 꼴로 손으로 다시 맞춰야 한다는 뜻이다. 실제 신경망은 층이 수십 개이고 층마다 입력 차원이 다르니, 이 맞춤을 층마다 따로 해야 한다.

이 층을 피셔 계량(곳마다 작은 스텝의 길이를 분포가 달라지는 정도로 재는 규칙)으로 읽으면 출력 변화가 왜 폭 n 을 따라 커지는지 보인다. 층이 평균 y, 분산 σ² 인 정규분포를 낸다고 보자. 가중치 공간의 피셔 정보와, 피셔 정보로 잰 스텝 dw 의 길이 제곱은

\begin{gathered} \textcolor{#0e8f8f}{I}(\textcolor{#2e9e6e}{w}) = \frac{1}{\textcolor{#7888e8}{\sigma}^2}\,\mathbb{E}\big[\textcolor{#6f8fa6}{x}\textcolor{#6f8fa6}{x}^{\top}\big] \\ d\textcolor{#2e9e6e}{w}^{\top}\,\textcolor{#0e8f8f}{I}\,d\textcolor{#2e9e6e}{w} = \frac{\mathbb{E}\big[(\textcolor{#6f8fa6}{x}\cdot d\textcolor{#2e9e6e}{w})^2\big]}{\textcolor{#7888e8}{\sigma}^2} \end{gathered}

\begin{array}{ll} \textcolor{#0e8f8f}{I} & \text{가중치 공간의 피셔 정보} \\ \textcolor{#2e9e6e}{w} & \text{가중치} \\ \textcolor{#6f8fa6}{x} & \text{입력} \\ \textcolor{#7888e8}{\sigma}^2 & \text{출력 정규분포의 분산 (상수)} \end{array}

피셔로 잰 스텝의 길이는 「출력이 얼마나 움직였나」를 출력의 눈금으로 잰 값이다. 경사하강 한 스텝에서 x·dw 가 n 에 비례하니 피셔 길이도 n 에 비례하지만, 유클리드 길이 ‖dw‖ 는 √n 에 비례할 뿐이다. 학습률은 스텝의 유클리드 길이를 정하는데, 출력 분포가 실제로 달라지는 정도는 피셔 계량으로 잰 길이를 따른다.

“폭마다 학습률을 손으로 다시 맞추는 대신, 한 스텝의 크기를 처음부터 분포의 눈금으로 잴 수는 없는가?”

정의

보통의 경사하강법은 파라미터 공간의 유클리드 거리를 쓴다. 아마리(Shun’ichi Amari)가 1998년 논문 「자연 그래디언트는 학습에서 효율적으로 작동한다(Natural Gradient Works Efficiently in Learning)」에서 다듬은 자연 그래디언트(natural gradient)는 피셔 계량 아래에서 가장 가파른 방향으로 간다. 같은 논문은 파라미터 공간에 이런 구조가 있으면 보통의 그래디언트는 가장 가파른 방향이 아니라는 데서 출발한다.

\textcolor{#2e9e6e}{w} \leftarrow \textcolor{#2e9e6e}{w} - \textcolor{#b88068}{\mathrm{lr}}\cdot \textcolor{#0e8f8f}{I}(\textcolor{#2e9e6e}{w})^{-1}\, \nabla \textcolor{#c05080}{\mathcal{L}}(\textcolor{#2e9e6e}{w})

\begin{array}{ll} \textcolor{#2e9e6e}{w} & \text{가중치} \\ \textcolor{#b88068}{\mathrm{lr}} & \text{학습률} \\ \textcolor{#0e8f8f}{I} & \text{피셔 정보 행렬} \\ \textcolor{#c05080}{\mathcal{L}} & \text{손실} \end{array}

그래디언트를 피셔 정보 행렬로 나누면, 분포가 크게 바뀌는 방향의 스텝은 줄고 조금만 바뀌는 방향의 스텝은 는다. 아래 그림은 손실의 등고선이 한쪽으로 길게 늘어난 경우에 두 방법을 비교한다.

손실 ½(0.5w₁² + 4w₂²), 피셔 정보 diag(0.5, 4) 에서 그래디언트 그대로 걷기(위, 학습률 0.42)와 자연 그래디언트로 걷기(아래, 학습률 0.5). 위는 가파른 w₂ 방향으로 오가며 튀고, 아래는 손실이 가장 작은 점 (0, 0)으로 곧장 간다

그림의 손실은 ½(0.5w₁² + 4w₂²) 이고 피셔 정보는 I = diag(0.5, 4) 이다(출력이 정규분포인 선형 모델에서는 둘이 같다). 그래디언트 그대로 걸으면 한 스텝마다 w₂ 는 −0.68 배, w₁ 은 0.79 배가 된다. 가파른 w₂ 방향의 그래디언트가 스텝을 좌우해서 골짜기를 오가며 튀고, 완만한 w₁ 방향으로는 조금씩밖에 못 나아간다. 그래디언트를 I 로 나누면 두 방향 모두 한 스텝마다 0.5 배가 되어 손실이 가장 작은 점 (0, 0)으로 곧장 간다. 열 스텝 뒤 (0, 0)까지 남은 거리는 그래디언트 그대로가 0.2847, 자연 그래디언트가 0.0031 이다.

폭 n 인 층(y = w·x, 입력 좌표 ±1)에서는 I 가 x 방향으로 n 에 비례하는 값을 가지므로, I⁻¹ 이 그 방향의 스텝을 1/n 로 줄인다. 손으로 하던 폭 맞춤이 저절로 들어 있다.

하나뿐인 계량, 그러나 파라미터 위에서는

첸초프(Chentsov)의 정리에 따르면, 분포의 공간에서 충분통계량 변환(데이터를 요약해도 정보를 잃지 않는 변환) 아래에서 바뀌지 않는 계량은 (상수배를 빼면) 피셔 하나뿐이다. 하지만 신경망의 파라미터 공간은 분포의 공간이 아니다. 뉴런의 순서를 바꾸거나 ReLU 앞뒤로 크기를 옮기는 것처럼, 서로 다른 가중치가 똑같은 분포를 내는 일이 흔하다(비식별성). 그래서 파라미터 위의 피셔 행렬은 대개 특이하고(역행렬이 없고), 거기에 첸초프의 유일성은 그대로 적용되지 않는다. 정확한 말은 이렇다. 파라미터 위의 피셔는 분포 공간의 피셔 계량을 파라미터로 끌어온(pullback) 것이다. 실제 구현은 I 의 대각선에 작은 값 ε을 더해 역을 구한다.

파이썬

폭 n 인 층에 자연 그래디언트를 적용해 본다. 입력 하나로 만든 피셔 I = xxᵀ 는 x 방향 하나에만 값이 있어 역행렬이 없으므로, 대각선에 ε = 10⁻⁶ 을 더해 푼다. 층과 씨앗은 위의 경사하강 코드와 같다.

import numpy as np
rng = np.random.default_rng(0)
lr, eps = 0.5, 1e-6
for n in [64, 256, 1024, 4096]:
    x = rng.choice([-1.0, 1.0], n)
    w = rng.standard_normal(n) / np.sqrt(n)
    y, target = w @ x, 1.0
    g = (y - target) * x                        # 경사하강 코드와 같은 층과 그래디언트
    I = np.outer(x, x)                          # 피셔 정보 (출력 분산 1, 입력 하나)
    dw = -lr * np.linalg.solve(I + eps * np.eye(n), g)
    print(f"n={n:5d}  오차 y−t {y - target:+.3f}  출력 변화 {x @ dw:+.3f}  스텝 길이 ‖dw‖ {np.linalg.norm(dw):.4f}")
# n=   64  오차 y−t +1.104  출력 변화 -0.552  스텝 길이 ‖dw‖ 0.0690
# n=  256  오차 y−t -0.855  출력 변화 +0.428  스텝 길이 ‖dw‖ 0.0267
# n= 1024  오차 y−t -0.974  출력 변화 +0.487  스텝 길이 ‖dw‖ 0.0152
# n= 4096  오차 y−t -0.929  출력 변화 +0.465  스텝 길이 ‖dw‖ 0.0073

lr = 0.5 를 모든 폭에 그대로 썼는데, 출력 변화는 늘 오차의 절반만큼 목표 쪽이다. 폭에 따라 줄어드는 것은 파라미터 공간의 스텝 길이 ‖dw‖ 이다(오차가 같다면 1/√n 꼴). 폭마다 손으로 하던 1/n 맞춤을 I⁻¹ 이 대신한다.

ML에서: 큰 모델의 옵티마이저가 흉내 내는 것

폭 4096 층 하나의 피셔만 해도 가중치 수의 제곱만 한 행렬이라, 실제 신경망에서 I⁻¹ 을 스텝마다 그대로 구할 수는 없다. 그래서 오늘날 큰 모델의 옵티마이저들은 자연 그래디언트의 성질 가운데 일부를 싸게 흉내 낸다. 층마다 피셔를 작은 두 행렬로 쪼개 근사하거나(K-FAC, Shampoo), 그래디언트 제곱의 평균으로 좌표마다 스텝을 나누거나(Adam), 폭에 따라 초기화와 학습률을 미리 정해 두거나(μP: 폭이 달라져도 학습률을 그대로 옮겨 쓰게 하는 설정), 스텝의 크기를 피셔 대신 다른 노름으로 재고 모든 방향을 고르게 맞춘다(Muon). 어떻게 계산하고 왜 그런 꼴이 나왔는지는 최적화 책에서 다룬다.

이 절의 결론

“자연 그래디언트는 피셔 계량으로 길이를 잰 한 스텝이어서, 학습률 하나로 폭이 다른 층을 같은 출력 변화만큼 움직이게 한다. 파라미터 위의 피셔 행렬은 분포 공간의 피셔 계량을 끌어온 것이라 역행렬이 없을 수 있고, 그때는 대각선에 작은 값을 더해 푼다.”

문제 1. 다이얼 눈금이 다른 두 오븐

오븐 A 는 다이얼을 한 칸 돌리면 10°C 오르고, 오븐 B 는 한 칸에 25°C 오른다. 두 오븐 모두 지금 160°C 이고 목표는 200°C 다. (가) 레시피가 「남은 온도 차(°C)에 0.1 을 곱한 만큼 다이얼을 칸으로 돌려라」라고 적었다. 이 지시를 세 번 따르면 두 오븐의 온도는 각각 어떻게 되는가? (나) 지시를 「남은 온도 차의 절반만큼 온도를 올려라」로 바꾸면, 세 번 뒤 온도와 한 번마다 돌린 칸 수는? (다) 두 지시 가운데 어느 쪽이 오븐을 가리지 않는가? 그 지시를 따르려면 무엇을 알아야 하는가?

함께 풀기

같은 레시피니까 두 오븐 다 비슷하게 가겠죠. A 는 0.1 × 40 = 4칸, 40°C 올라서 한 번에 200°C 예요. B 도 4칸이니까…

B 의 4칸은 몇 도예요?

100°C 요. 260°C 가 됐어요. 다음엔 남은 차가 −60 이라 6칸 내리면 −150°C, 110°C. 그다음엔 335°C. 점점 더 크게 넘어가요.

한 번마다 남은 차가 (1 − 0.1 × 25) = −1.5 배가 되니까 그래. A 는 1 − 0.1 × 10 = 0 배라 한 번에 끝났고. 칸으로 적은 지시는 칸 하나가 몇 도인지에 따라 뜻이 달라져.

(나)는요?

남은 차가 40, 20, 10, 5 로 줄어서 두 오븐 모두 180, 190, 195°C 예요. 돌린 칸은 A 가 2, 1, 0.5칸이고 B 가 0.8, 0.4, 0.2칸이에요. 칸 수는 달라도 온도는 똑같이 가요.

대신 (나)를 따르려면 우리 오븐이 한 칸에 몇 도인지 알아야 하네요. 올릴 온도를 그 값으로 나눠서 칸으로 바꾸니까요.

칸 수가 좌표고, 한 칸당 온도가 피셔 계량이네. 자연 그래디언트가 그래디언트를 피셔 정보로 나누는 게 올릴 온도를 칸 수로 바꾸는 그 환산이야.

문제 2. 좌표 이름표를 바꾸면

데이터는 정규분포 N(0, 1) 에서 오고, 모델은 평균을 0 으로 고정한 N(0, σ²) 이다. 지금 σ = 0.5 다. 손실은 L(σ) = log σ + 1/(2σ²)(데이터와 모델 사이 KL 에서 σ 와 상관없는 상수를 뺀 것)이다. 피셔 정보는 σ 를 좌표로 쓰면 2/σ², s = log σ 를 좌표로 쓰면 2 다. (가) 보통 경사하강 lr = 0.01 로 한 스텝 걸은 뒤의 σ 를, σ 좌표로 걸을 때와 s 좌표로 걸을 때 각각 구하라. (나) 자연 그래디언트 lr = 0.01 로 두 좌표에서 한 스텝씩 걸은 뒤의 σ 를 구하라. (다) 어느 쪽 걷기가 「좌표를 무엇으로 적느냐」에 덜 휘둘리는가? lr 을 0.001 로 줄이면 어떻게 되는가?

함께 풀기

손실도 같고 lr 도 같으니까 좌표만 바꾼다고 다른 곳에 가진 않겠죠. 그래도 계산해 볼게요. σ 좌표에서는 dL/dσ = 1/σ − 1/σ³ = 2 − 8 = −6 이니까 σ = 0.5 + 0.01 × 6 = 0.56.

s 좌표에서는 dL/ds = σ·dL/dσ = −3 이라 s 가 0.03 늘고, σ = 0.5·e^0.03 = 0.515. 같은 손실을 같은 lr 로 내려왔는데 0.56 과 0.515 예요.

둘 중 어느 쪽이 맞는 스텝이에요?

음… 맞고 틀리고를 가를 기준이 없는 것 같아요. 두 좌표의 한 칸이 같은 길이인지부터 모르겠어요.

σ = 0.5 근처에서 s 가 0.01 늘면 σ 는 얼마쯤 늘어요?

s = log σ 니까 dσ = σ ds… 0.005 쯤이요. 두 좌표의 한 칸은 길이가 달라요.

그래디언트도 다르게 나왔잖아. −6 이랑 −3.

그래디언트도 같은 비율로 달라지니까, s 좌표의 스텝은 σ 로 옮기면 σ² 배가 돼요. 0.06 대 0.015, 1/σ² = 4 배 차이예요.

자연 그래디언트로 걸으면요?

σ 좌표부터요. 그래디언트 −6 을 피셔 2/σ² = 8 로 나누면 −0.75, σ = 0.5 + 0.0075 = 0.5075.

s 좌표에서는요?

−3 을 피셔 2 로 나누면 −1.5, s 가 0.015 늘어 σ = 0.5·e^0.015 = 0.50756. 거의 같아요.

그런데 완전히 같지는 않네요. 0.00006 차이가 나요.

lr 을 0.001 로 줄이면 두 좌표에서 얻은 σ 의 차이는 어떻게 돼요?

스텝은 0.00075 로 10분의 1이 되는데, 차이는 0.00000056 으로 100분의 1이 돼요.

스텝에 견주면 차이는 어떻게 돼 가요?

점점 사라져요. 자연 그래디언트는 스텝을 분포가 달라지는 정도로 재니까, 아주 작은 스텝에서는 좌표를 무엇으로 적든 같은 분포로 가요.

그럼 남는 0.00006 같은 차이는 뭐야?

휘어진 좌표 위를 곧게 걸어서 생기는 차이야. 스텝이 10분의 1일 때 100분의 1이 됐으니까 스텝 제곱 크기고.

문제 1(눈금이 다른 두 오븐)이랑 같네요. 칸 수로 적은 지시는 오븐마다 결과가 달랐는데, 온도로 적은 지시는 같았잖아요. 여기선 σ 와 s 가 한 칸의 크기가 다른 두 좌표고, 경사하강이 칸 수로 적은 지시, 자연 그래디언트가 분포의 변화로 적은 지시예요.

적분에서 변수를 바꾸면 dx 앞에 자코비안이 붙어서 값이 그대로인 것처럼, 피셔가 좌표를 바꿀 때 생기는 배율의 제곱을 알아서 곱해 주는 거네요.

문제 3. 그래디언트 클리핑을 피셔 계량으로 재면

학습이 튀는 것을 막으려고 흔히 그래디언트 클리핑(스텝의 유클리드 길이가 정한 값을 넘으면 그 값으로 줄이는 장치)을 건다. 폭 n 인 층(y = w·x, 입력 좌표 ±1, 출력 분산 σ² = 1)에서 한 스텝 dw 가 늘 유클리드 길이 0.05 로 잘린다고 하자(방향은 그래디언트 방향 그대로, 곧 x 방향). (가) n = 1024 와 n = 4096 에서 잘린 스텝은 출력을 얼마나 움직이는가? (나) 잘린 스텝의 피셔 길이를 구하라. (다) 두 폭에서 출력이 같은 만큼 움직이게 하려면 n = 4096 의 자르는 길이를 얼마로 잡아야 하는가? 자르는 기준을 피셔 길이로 바꾸면 어떻게 되는가?

함께 풀기

스텝 길이를 0.05 로 묶었으니까 폭이 어떻든 같은 크기로 움직이는 거잖아요. 클리핑만 걸면 폭 걱정은 끝 아니에요?

길이 0.05 로 잘린 스텝이 출력을 얼마나 바꾸는지 n = 1024 에 넣어 봐요.

스텝이 x 방향이고 길이가 0.05 니까 x·dw = ‖x‖ × 0.05 = √1024 × 0.05 = 1.6. n = 4096 이면 √4096 = 64 라서 3.2 요. 길이는 같은데 출력은 두 배를 움직이네요.

스텝이 x 와 같은 방향이면 출력은 유클리드 길이의 √n 배로 움직여. 피셔 길이로 재면 σ = 1 이라 |x·dw| 그대로, 1.6 과 3.2 야. 피셔 길이는 두 스텝이 출력을 다르게 움직인다는 걸 그대로 보여 주고, 유클리드 길이는 못 보여 줘.

그럼 n = 4096 에서도 출력이 1.6 만 움직이게 하려면요?

√n 에 반비례로, 0.05 × 32/64 = 0.025 요. 결국 폭마다 자르는 값을 다시 맞춰야 하네요.

자르는 기준을 처음부터 피셔 길이 1.6 으로 잡으면요?

아까 피셔 길이가 |x·dw| 그대로였으니까, 피셔 길이가 출력 변화예요. 그러면 두 폭 모두 출력이 1.6 씩 움직여요.

유클리드 길이로는 얼마씩 잘리는 건데?

알아서 0.05 와 0.025 가 되지. 문제 1(두 오븐)에서 칸 대신 온도로 지시한 것처럼, 자르는 기준을 피셔 길이로 바꾸니 폭마다 다시 맞출 일이 없어졌어.

유클리드 길이로만 자르는 클리핑은, 보고서를 「10쪽 이내」로만 묶으면 글자를 줄여서 내용은 얼마든지 늘릴 수 있는 것과 같네요. 쪽수가 아니라 내용으로 재야 분량이 맞춰지죠.