자연 그래디언트: 좌표를 바꿔도 같은 방향

피셔 행렬을 계량으로 쓰면 가장 가파른 스텝은 그래디언트에 피셔 행렬의 역행렬을 곱한 방향이었다. 피셔 행렬의 역행렬을 곱한 스텝은 실제 학습에서 보통 경사하강과 얼마나 다르게 걸을까? 그리고 출발 문제가 물은 대로, 정말 좌표에 속지 않을까?

역사: 자주 들르는 상태만 배우는 그래디언트

강화학습에서는 상태마다 어떤 행동을 얼마의 확률로 고를지(정책)를 파라미터 $\textcolor{#1b9e77}{\theta}$ 로 적고, 오래 돌렸을 때 받는 평균 보상 $J$ 가 커지도록 $\textcolor{#1b9e77}{\theta}$ 를 그래디언트를 따라 고친다. 이 방법을 정책 그래디언트라 부른다. 런던 대학 개츠비 계산신경과학 연구단의 샴 카카데(Sham Kakade)는 2001년 NIPS 학회 논문 「A Natural Policy Gradient」에서 이 규칙 $\Delta\textcolor{#1b9e77}{\theta}_i = \partial J/\partial\textcolor{#1b9e77}{\theta}_i$ 의 단위부터 따졌다. 좌변의 단위는 $\textcolor{#1b9e77}{\theta}_i$ 인데 우변은 $\textcolor{#1b9e77}{\theta}_i$ 로 나눈 값이라 단위가 $1/\textcolor{#1b9e77}{\theta}_i$ 다. 양변의 차원이 맞지 않으니, 파라미터를 어떤 단위로 적느냐에 따라 같은 정책에서 출발해도 다른 방향으로 걷는다. 카카데는 이 규칙이 좌표에 따라 달라지는(공변이 아닌) 규칙이라고 짚었다. 그래서 정책이 정하는 행동 확률분포에서 피셔 정보로 계량을 세우고, 그 계량으로 잰 가장 가파른 방향을 따르는 그래디언트를 만들었다. 파라미터 공간의 구조를 따른 가장 가파른 방향인 자연 그래디언트는 1998년 아마리 순이치가 내놓았고, 카카데는 그것을 정책에 옮긴 것이다.

보통 그래디언트가 어떻게 막히는지는 상태가 둘뿐인 작은 문제에서 드러났다. 상태 i에서 제자리에 머물면 보상 1, 상태 j에서 머물면 보상 2, 상태를 건너가면 보상 0이다.

카카데(2001) 그림 1B의 상태 둘짜리 문제를 다시 그린 것. 상태마다 제자리에 머물거나 다른 상태로 건너가는 행동 둘이 있고, 오래 돌리면 j에 머무는 정책이 평균 보상 2로 가장 좋다. 처음 정책은 i에서 0.8, j에서 0.1의 확률로 머문다. 그림: sources/ch12/kakade_two_state.py

처음에는 i에서 머물 확률이 0.8이라 에이전트(행동을 고르는 프로그램)는 시간의 80% 남짓을 i에서 보낸다. 보통 그래디언트는 상태마다 그 상태에 머무는 비율 $p(s)$ 만큼 무게를 실어 배우므로, 자주 들르는 i에서 머물기를 먼저 빨리 배운다. i에 오래 머물수록 j에 들르는 비율은 더 줄고, j에서 배우는 속도는 더 느려진다. 논문의 계산에서 j에 머무는 비율은 0.2에서 $10^{-7}$ 까지 떨어졌고, 평균 보상은 1에서 극도로 평평한 고원에 갇혔다. 자연 그래디언트는 상태마다 고르게 배워 가장 좋은 정책(j에 머물기)까지 걸리는 시간을 $10^7$ 분의 1로 줄였다. 기둥 높이, 이웃 기둥의 높이 차, 가장 높은 기둥, 구멍 수를 특징으로 쓴 테트리스에서도 보통 그래디언트는 아주 느렸고 자연 그래디언트는 점수를 빠르게 끌어올렸다.

상태 둘짜리 문제(i에서 머물면 보상 1, j에서 머물면 보상 2, 건너가면 0)를 다시 계산한 것. 위: 평균 보상. 보통 그래디언트(실선)는 평균 보상 1의 고원에 오래 머물고, 자연 그래디언트(점선)는 일찍 2에 닿는다. 아래: 상태 j에 머무는 비율. 보통 그래디언트에서는 가까이 떨어진다. 가로축은 로그 눈금의 시간. 카카데(2001) 그림 1의 설정(머물 확률을 시그모이드로 적고 i 0.8, j 0.1에서 출발)을 따라 정확한 그래디언트를 적분했다. 세부 설정이 논문과 다를 수 있어, 이 계산에서는 걸린 시간의 비가 약 배, j의 비율 최솟값이 약 으로 논문의 값보다 한 자리 작다. 계산: sources/ch12/kakade_two_state.py — 상태 둘짜리 문제(i에서 머물면 보상 1, j에서 머물면 보상 2, 건너가면 0)를 다시 계산한 것. 위: 평균 보상. 보통 그래디언트(실선)는 평균 보상 1의 고원에 오래 머물고, 자연 그래디언트(점선)는 일찍 2에 닿는다. 아래: 상태 j에 머무는 비율. 보통 그래디언트에서는 $10^{-6}$ 가까이 떨어진다. 가로축은 로그 눈금의 시간. 카카데(2001) 그림 1의 설정(머물 확률을 시그모이드로 적고 i 0.8, j 0.1에서 출발)을 따라 정확한 그래디언트를 적분했다. 세부 설정이 논문과 다를 수 있어, 이 계산에서는 걸린 시간의 비가 약 $10^6$ 배, j의 비율 최솟값이 약 $10^{-6}$ 으로 논문의 값보다 한 자리 작다. 계산: sources/ch12/kakade_two_state.py

논문은 스스로 한계도 적었다. 카카데의 피셔 행렬 $\textcolor{#d62728}{F}$ 는 상태에 머무는 비율 $p(s)$ 로 가중해 평균한 것이라, 파라미터의 단위를 바꾸면 학습 곡선이 완전히 같지는 않았다. 이 $\textcolor{#d62728}{F}$ 는 완전한 불변 계량이 아니라는 것이다. 그러면 계량이 어떤 조건을 갖출 때 걷는 방향이 좌표에 휘둘리지 않는가? 이 절은 그것을 확인한다.

단위가 안 맞는다는 건 물리 시간에 제일 먼저 확인하는 건데요. 그걸 2001년에야 짚었다니 의외예요.

물리량에는 미터나 초가 붙어 있어서 어긋나면 바로 보여요. 정책 파라미터 $\textcolor{#1b9e77}{\theta}$ 에는 단위가 적혀 있나요?

아니요, 시그모이드에 넣는 그냥 실수예요. 아, 단위가 안 보이니 어긋난 것도 안 보였겠네요.

그래도 고원은 그래프에 보였을 것 같아요. 평균 보상이 1에서 오래 멈춰 있으면 이상하다고 느끼지 않았을까요?

상태 둘짜리 문제는 가장 좋은 평균 보상이 2라는 걸 알아서 1이 고원인 줄 알아요. 가장 좋은 값을 모르는 테트리스 같은 문제에서는 그 평평한 구간이 어떻게 보일까요?

다 배워서 멈춘 것처럼 보이겠어요. 게다가 j에 머무는 비율이 $10^{-7}$ 이면 에이전트를 돌려 모은 기록에 j가 거의 나오지 않으니까, j에 더 큰 보상이 있다는 단서도 데이터에 없고요.

그래서 카카데는 $\textcolor{#1b9e77}{\theta}$ 의 숫자가 아니라 행동 확률분포가 얼마나 바뀌는지로 스텝을 쟀어요. 다만 상태마다 분포가 따로 있어서 머무는 비율 $p(s)$ 로 평균을 냈고, 그 평균 때문에 완전히 불변은 아니었죠.

그럼 $\textcolor{#1b9e77}{\theta}$ 를 어떤 단위로 적든 같은 방향으로 걸으려면, 계량은 좌표를 바꿀 때 어떻게 바뀌어야 해요?

같은 출발점, 다른 길

가장 단순한 경우로 손실이 $\textcolor{#6b6b6b}{L}(\textcolor{#1b9e77}{\theta}) = \frac12 \textcolor{#1b9e77}{\theta}^T \textcolor{#d62728}{F}\, \textcolor{#1b9e77}{\theta}$ 인 경우를 보자. 공분산을 아는 2차원 정규분포의 평균을 맞출 때 음의 로그우도가 이 꼴이고, 그때 피셔 행렬이 바로 이 $\textcolor{#d62728}{F}$ 다. 아래 그림의 $\textcolor{#d62728}{F}$ 는 한 방향이 다른 방향보다 열 배 가파르게 골랐다. 보통 경사는 학습률을 0.2보다 키우면 가파른 방향으로 튕겨 나가므로 $\textcolor{#6b672e}{\eta} = 0.17$ 로 걸어야 하고, 그러면 골짜기의 두 벽 사이를 지그재그로 오가며 네 스텝 뒤에 손실이 처음의 약 1/6로 준다. 그래디언트에 $\textcolor{#d62728}{F}^{-1}$ 을 곱한 스텝은 $\textcolor{#6b672e}{\eta} = 0.5$ 에서 최솟값을 향해 곧장 가서, 네 스텝 뒤 손실이 처음의 1/256이다.

같은 손실 등고선(회색 타원) 위에서 같은 점을 출발한 두 길. 위: 보통 경사하강()은 가파른 벽 사이를 지그재그로 오간다. 아래: 그래디언트에 피셔 행렬의 역행렬을 곱한 자연 그래디언트()는 최솟값으로 곧장 간다. 손실은 이고, 는 한 방향이 열 배 가파른 행렬로 골랐다 — 같은 손실 등고선(회색 타원) 위에서 같은 점을 출발한 두 길. 위: 보통 경사하강( $\eta = 0.17$ )은 가파른 벽 사이를 지그재그로 오간다. 아래: 그래디언트에 피셔 행렬의 역행렬을 곱한 자연 그래디언트( $\eta = 0.5$ )는 최솟값으로 곧장 간다. 손실은 $L(\theta) = \frac12\theta^T F\theta$ 이고, $F$ 는 한 방향이 열 배 가파른 행렬로 골랐다

직관적인 비유를 하자면, 유클리드 그래디언트는 "지도 위에서의 가장 가파른 방향"이고, 피셔 계량으로 보정한 그래디언트인 자연 그래디언트는 "실제 지형 위에서의 가장 가파른 방향"이다. 메르카토르 도법으로 그린 지도에서는 극지방이 과도하게 크게 보이므로, 지도 위에서 가장 짧아 보이는 이동이 실제 지형 위에서도 가장 짧은 이동은 아니다. 피셔 계량은 이 "지도의 왜곡"을 보정한다. 유클리드 공간처럼 피셔 행렬이 단위행렬이면 두 스텝이 같지만, 확률 모델의 파라미터 공간처럼 계량의 눈금이 장소와 방향마다 다르면 두 스텝은 크게 다를 수 있다.

자연 그래디언트와 아마리의 정리 (1998)

자연 그래디언트(natural gradient)는 다음과 같이 정의된다.

\tilde{\nabla}\textcolor{#6b6b6b}{L} = \textcolor{#d62728}{F}^{-1}\, \nabla \textcolor{#6b6b6b}{L}

\begin{array}{ll} \tilde{\nabla}\textcolor{#6b6b6b}{L} & \text{자연 그래디언트 (걸어갈 방향의 반대, 벡터)} \\ \textcolor{#d62728}{F} & \text{피셔 정보행렬} \\ \nabla \textcolor{#6b6b6b}{L} & \text{보통의 그래디언트 } (\partial \textcolor{#6b6b6b}{L}/\partial\textcolor{#1b9e77}{\theta}^j) \text{. } \nabla \text{는 여기서 접속이 아니라 그래디언트라 기본색} \end{array}

자연 그래디언트 (좌표 독립 경사 / Natural Gradient, $\tilde{\nabla}\textcolor{#6b6b6b}{L} = \textcolor{#d62728}{F}^{-1}\nabla \textcolor{#6b6b6b}{L}$ ): 파라미터 공간을 피셔 계량을 가진 리만 매니폴드로 보았을 때의 “진정한” 최급강하 방향. 유클리드 그래디언트에 피셔 행렬의 역행렬을 곱한 것이다. $\textcolor{#d62728}{F} = I$ 이면 보통 그래디언트와 같지만, 파라미터 공간이 휘어져 있거나 눈금이 방향마다 다르면 방향이 크게 달라질 수 있다. 자연 그래디언트의 핵심 성질은 재파라미터화 불변성이며, 이 불변성은 매니폴드 위의 벡터가 좌표에 의존하지 않는 것과 같은 원리다.

아마리의 핵심 정리는 자연 그래디언트는 재파라미터화에 대해 불변이다라는 것이다. 파라미터를 $\textcolor{#1b9e77}{\theta}$ 에서 $\textcolor{#1b9e77}{\varphi} = f(\textcolor{#1b9e77}{\theta})$ 로 바꾸고 새 좌표에서 자연 그래디언트를 계산하면, 새 좌표의 자연 그래디언트는 원래 좌표의 자연 그래디언트를 $f$ 로 옮겨 놓은 것과 정확히 일치한다.

ML에서: 자연 그래디언트의 후예들

실용적인 관점에서 자연 그래디언트의 가장 큰 장벽은 계산 비용이다. 파라미터가 $\textcolor{#c4724f}{n}$ 개면 $\textcolor{#c4724f}{n} \times \textcolor{#c4724f}{n}$ 행렬의 역행렬을 매 단계 구해야 하는데, 현대 신경망에서 $\textcolor{#c4724f}{n}$ 은 수억에 달하므로 이 역행렬 계산을 그대로 하기는 불가능하다. 하지만 자연 그래디언트의 정신은 여러 실용적 알고리즘에 살아 있다. K-FAC(Kronecker-Factored Approximate Curvature, 마튼스·그로스 2015)은 피셔 행렬을 층마다 두 작은 행렬의 크로네커 곱(한 행렬의 칸마다 다른 행렬 전체를 곱해 넣어 만든 큰 행렬)으로 근사한다. 널리 쓰는 옵티마이저 Adam은 그래디언트 제곱의 이동평균(모델 분포 대신 데이터 샘플로 스코어 제곱을 평균한 경험적 피셔 행렬의 대각 성분에 해당)의 제곱근으로 나눈다는 점에서 대각 근사의 먼 친척으로 볼 수 있다. 강화학습의 TRPO(Trust Region Policy Optimization)는 정책 공간에서의 자연 그래디언트를 KL 발산 제약(신뢰 영역)과 결합한 것이며, PPO(Proximal Policy Optimization, 옛 정책 가까이에서만 고치는 정책 최적화)는 TRPO의 실용적 근사다. 이들 모두 "파라미터 공간의 기하학을 존중하라"는 아마리의 통찰에서 뻗어 나온 가지들이다.

문제 4: 샤워 꼭지 두 개

확인 샤워기에 온수(60°C) 꼭지와 냉수(20°C) 꼭지가 따로 있고, 지금은 둘 다 1분에 5 L씩 나와 40°C 물이 1분에 10 L 나온다. 물 양은 그대로 두고 온도만 1°C 올리고 싶다. (1) 온수 꼭지만 1분에 0.25 L 더 열면 온도와 물 양은 어떻게 되는가? (2) 두 꼭지를 어떻게 돌려야 하는가?

60°C 온수 꼭지와 20°C 냉수 꼭지가 1분에 5 L씩 내보내 40°C 물 10 L를 만든다

같이 풀기 세션 4

온수를 더 열면 되죠. 섞인 온도는 $(60 \times 5.25 + 20 \times 5)/10.25 \approx 40.49$ °C. 1도가 안 오르네요. 그럼 더 많이 열면 되겠다.

그때 물 양은 어떻게 되죠? 문제는 물 양을 그대로 두라고 했어요.

10.25 L요. 더 열면 더 늘고요. 온수만 돌렸는데 물 양까지 같이 바뀌네요.

온수를 연 만큼 냉수를 잠그면 물 양은 10 L로 유지돼. 온수 5.25, 냉수 4.75면 $(315 + 95)/10 = 41$ °C. 딱 1도.

「온수 꼭지만」은 꼭지 쪽에서 보면 똑바른 한 방향이에요. 온도와 물 양 쪽에서 보면 어느 방향이었죠?

온도와 물 양이 함께 움직이는 비스듬한 방향이요. 결과 쪽에서 「온도만」 똑바로 가려면 꼭지 쪽에서는 $(+0.25,\ -0.25)$ , 두 꼭지를 함께 돌려야 하네요.

꼭지 하나가 결과 두 개에 얽혀 있어서 그렇구나. 하나만 돌리면 원하지 않는 쪽까지 끌려가네요.

문제 5: 학습률만 키우면 같을까

킬러 두 파라미터가 서로 얽힌 모델의 한 점에서(계산하기 쉽게 고른 값) 피셔 행렬이 $\textcolor{#d62728}{F} = \begin{pmatrix} 2 & 1 \\ 1 & 1 \end{pmatrix}$ 이고 그래디언트가 $\nabla \textcolor{#6b6b6b}{L} = (1, 0)$ 이다. 자연 그래디언트 $\textcolor{#d62728}{F}^{-1}\nabla \textcolor{#6b6b6b}{L}$ 을 구하고, 보통 그래디언트와 방향이 어떻게 다른지 설명하라. 방향이 달라도 여전히 손실이 줄어드는 방향인 이유는 무엇인가?

그래디언트 을 스텝 평면 에 그린 것. 피셔 행렬에는 비대각 성분이 있다 — 그래디언트 $\nabla L = (1, 0)$ 을 스텝 평면 $(\delta\theta^1, \delta\theta^2)$ 에 그린 것. 피셔 행렬에는 비대각 성분이 있다

같이 풀기 세션 5

PyTorch로 작은 모델에 자연 그래디언트를 돌려 봤는데요, 보통 SGD(확률적 경사하강법)랑 궤적이 거의 비슷하고 더 빨리 수렴하기만 하더라고요. 결국 자연 그래디언트는 학습률을 방향마다 조금씩 키운 거 아니에요? 학습률만 잘 고르면 같을 것 같아요.

이 문제는 계산이 금방이야. 피셔 행렬 대각이 2랑 1이니까 역행렬은 $\text{diag}(1/2,\ 1)$ , 그러면 $\textcolor{#d62728}{F}^{-1}\nabla \textcolor{#6b6b6b}{L} = (1/2,\ 0)$ . 방향은 그래디언트랑 똑같고 길이만 반이네. 민준이 말이 맞는 것 같은데?

서연 학생, 역행렬을 구할 때 비대각 성분 1은 어디로 갔어요?

아… 대각 성분만 뒤집었어요. 대각행렬일 때만 되는 계산인데. 제대로 하면 행렬식이 $2 - 1 = 1$ 이니까 $\textcolor{#d62728}{F}^{-1} = \begin{pmatrix} 1 & -1 \\ -1 & 2 \end{pmatrix}$ , $\textcolor{#d62728}{F}^{-1}\nabla \textcolor{#6b6b6b}{L} = (1,\ -1)$ . 그래디언트는 $x$ 축 방향인데 자연 그래디언트는 45도 기울어졌어요.

방향이 바뀌네요. 그럼 제 PyTorch 실험은 뭐였죠?

그 모델의 피셔 행렬이 대각에 가까웠을 거예요. 파라미터들이 서로 얽혀 있지 않으면 자연 그래디언트는 정말로 "방향마다 다른 학습률"이 돼요. Adam이 대각 근사의 친척인 것도 그래서고요. 하지만 비대각 성분이 있으면, 한 파라미터를 움직일 때 다른 파라미터도 같이 보정해야 분포 공간에서 똑바로 가요. 이건 성분별 학습률로는 흉내 낼 수 없어요.

그럼 45도나 틀어졌는데도 손실이 줄어드는 게 보장돼요? 그래디언트 반대 방향이 가장 가파른 방향이라면서요.

내적을 보면 될 것 같아. 손실의 1차 변화는 $\nabla \textcolor{#6b6b6b}{L} \cdot \delta\textcolor{#1b9e77}{\theta} = -\textcolor{#6b672e}{\eta}\,\nabla \textcolor{#6b6b6b}{L}^T \textcolor{#d62728}{F}^{-1}\nabla \textcolor{#6b6b6b}{L}$ 이야.

그 값의 부호는 무엇이 정하죠?

가운데 낀 $\textcolor{#d62728}{F}^{-1}$ 이요. $\textcolor{#d62728}{F}$ 가 양의 정부호(0이 아닌 벡터 $v$ 마다 $v^T \textcolor{#d62728}{F} v > 0$ )면 $\textcolor{#d62728}{F}^{-1}$ 도 양의 정부호라서, 손실의 1차 변화는 항상 음수예요.

숫자로도 볼게. $(1, 0)\cdot(1, -1) = 1 > 0$ 이니까 반대로 걸으면 줄어드네.

응. 가장 가파른 방향은 "어떤 계량으로 재느냐"에 따라 달라질 뿐이야.

정확해요. 그래디언트 반대 방향은 유클리드 계량으로 잰 가장 가파른 방향이고, 자연 그래디언트는 피셔 계량으로 잰 가장 가파른 방향이에요. 둘 다 내려가는 방향이지만, 좌표를 바꿔도 살아남는 건 자연 그래디언트뿐이죠.

문제 4(샤워 꼭지 두 개)랑 같네요. 온수만 열면 물 양까지 바뀌어서 냉수를 같이 잠가야 했던 것처럼, 비대각 성분이 그 "같이 맞춰 주기"네요. 방향도 똑같이 $(1,\ -1)$ 쪽이고요.

선형대수 시험에서도 "대각행렬이면"이라는 조건 없이 역행렬을 성분별로 뒤집었다가 감점당한 적이 있어요. 그 뒤로 조건부터 확인하는데도 또 놓쳤네요.