Natural Gradient는 왜 좌표를 묻지 않는가

유클리드 gradient의 parameterization 의존성 문제부터 Fisher 계량 하의 steepest descent 유도, K-FAC·Shampoo의 실전 근사까지, Natural Gradient의 철학과 구조를 추적한다.

torch.optim.SGD를 쓸 때 우리는 ∇L을 “방향”이라고 부른다. 그런데 같은 목적함수라도 $\sigma$ 좌표로 최적화하느냐 $\log\sigma$ 좌표로 하느냐에 따라 gradient 경로 전체가 달라진다. 좌표 선택이 수렴 경로를 바꾼다면, 우리가 “gradient”라고 부르는 것은 과연 무엇인가?

∇L은 방향이 아니다

편미분 벡터 $\nabla L = (\partial_1 L, \dots, \partial_n L)$ 은 “steepest ascent 방향”처럼 보인다. 하지만 이것은 좌표계가 직교정규이고 계량이 항등 행렬 $g = I$ 일 때만 사실이다.

미분기하 언어로 말하면, $dL = \partial_i L \, d\theta^i$ 는 **1-form(covector)**이다. 방향을 입력으로 받아 실수를 돌려주는 함수. 반면 “gradient”는 방향을 가리키는 벡터다. 두 대상을 연결하는 것이 계량 $g$ 이다.

g(\operatorname{grad} L, V) = dL(V) \quad \forall V \in T_p M

좌표로 쓰면 $(\operatorname{grad} L)^i = g^{ij} \partial_j L$ . 유클리드에서는 $g^{ij} = \delta^{ij}$ 이므로 두 개가 일치하지만, 이것은 특수한 일치일 뿐이다.

재매개변수화 $\phi = \phi(\theta)$ , Jacobian $J = \partial\phi/\partial\theta$ 아래 gradient는 역전치 Jacobian으로 변환한다.

\nabla^\phi L = J^{-T} \nabla^\theta L

이것이 covector 변환 규칙이다. $\mathcal{N}(\mu, \sigma^2)$ 에서 $\text{KL}(N(\mu,\sigma^2) \| N(0,1))$ 을 최소화할 때 이 차이가 선명해진다. $\sigma$ 좌표에서 gradient의 $\sigma$ 성분은 $\sigma - 1/\sigma \approx -1/\sigma$ 로 $\sigma \to 0$ 에서 발산한다. $\log\sigma$ 좌표에서는 $-1 + e^{2\tau} \approx -1$ 로 유계다. 같은 분포, 같은 목적함수, 다른 안정성 — 이것이 parameterization이 gradient에 새기는 편향이다.

Fisher 계량과 Natural Gradient

분포 공간에서 자연스러운 거리는 KL divergence다. 무한소에서 KL은 다음과 같이 근사된다.

\operatorname{KL}(p_\theta \| p_{\theta+d\theta}) = \frac{1}{2} d\theta^T F(\theta) d\theta + O(\|d\theta\|^3)

여기서 $F(\theta) = \mathbb{E}[\nabla \log p \cdot \nabla \log p^T] = -\mathbb{E}[\nabla^2 \log p]$ 가 Fisher 정보 행렬이다. 0차 항은 $\operatorname{KL}(p\|p)=0$ 이고, 1차 항은 score의 기댓값이 0이므로 소멸한다. KL의 로컬 구조는 Fisher 계량으로 완전히 결정된다.

이로부터 Natural Gradient는 제약 최적화의 해로 등장한다. “KL 반경 $\varepsilon$ 안에서 $L$ 을 가장 많이 줄이는 방향”은 다음 문제의 해다.

\min_{d\theta} \nabla L(\theta)^T d\theta \quad \text{s.t.} \quad d\theta^T F(\theta) d\theta \leq \varepsilon^2

정리 1 · Natural Gradient = KL-ball Steepest Descent

위 제약 최적화 문제의 유일한 해는

d\theta^* = -\frac{\varepsilon}{\sqrt{\nabla L^T F^{-1} \nabla L}} \underbrace{F^{-1} \nabla L}_{=:\,\tilde{\nabla} L}

이다. 즉 Fisher 타원체 내에서 $L$ 을 가장 많이 줄이는 방향은 항상 $-F^{-1}\nabla L$ 방향이다.

▷ 증명

Lagrangian $\mathcal{L} = \nabla L^T d\theta + \lambda(d\theta^T F d\theta - \varepsilon^2)$ 의 정류 조건은 $\nabla L + 2\lambda F d\theta = 0$ , 즉 $d\theta = -\frac{1}{2\lambda} F^{-1} \nabla L$ . $\nabla L \neq 0$ 이면 최적점은 경계 $d\theta^T F d\theta = \varepsilon^2$ 에 있다. 이를 대입해 $\lambda^* = \frac{1}{2\varepsilon}\sqrt{\nabla L^T F^{-1} \nabla L}$ 을 얻고, 다시 대입하면 위 공식이 나온다. $F \succ 0$ 이므로 해는 유일하다. $\square$

∎

지수족 MLE에서 $L(\theta) = -\log p(x|\theta)$ 의 Hessian은 $\nabla^2 L = F$ 다. 따라서 $F^{-1}\nabla L = H^{-1}\nabla L$ — Natural Gradient는 지수족 MLE에서 Newton’s method와 정확히 일치하며, 이론적으로 한 스텝에 수렴한다.

좌표 불변성 — Natural Gradient가 “옳은” 이유

Fisher는 $(0,2)$ 텐서 법칙으로 변환한다.

F_\phi = J^{-T} F_\theta J^{-1}, \quad F_\phi^{-1} = J F_\theta^{-1} J^T

이 두 변환을 결합하면 Natural Gradient의 핵심 성질이 나온다.

\tilde{\nabla}_\phi L = F_\phi^{-1} \nabla_\phi L = (J F_\theta^{-1} J^T)(J^{-T} \nabla_\theta L) = J \tilde{\nabla}_\theta L

Natural gradient는 **반변 벡터 변환 규칙 $X_\phi = J X_\theta$ **를 따른다. 이는 접벡터 $d\theta$ 와 같은 변환이므로, 방향으로 쓰는 것이 기하학적으로 정당하다.

연속 시간 flow $\dot\theta = -F_\theta^{-1} \nabla_\theta L$ 에서 $\phi(t) = \phi(\theta(t))$ 로 변환하면 $\dot\phi = J\dot\theta = -F_\phi^{-1}\nabla_\phi L$ . 두 좌표계에서 Natural Gradient flow를 풀면 완전히 같은 분포 궤적을 그린다. 이산 업데이트에서는 $O(\eta^2)$ 오차가 남지만, $\eta \to 0$ 극한에서 정확히 불변이다.

반면 유클리드 GD $\phi_{t+1} = \phi_t - \eta \nabla_\phi L$ 이 $\theta$ 좌표의 경로와 일치하려면 $JJ^T = I$ , 즉 Jacobian이 직교행렬이어야 한다. 일반적인 재매개변수화에서는 성립하지 않는다.

실전 근사 — K-FAC과 Shampoo

$n = 10^9$ 파라미터 신경망에서 $F$ 는 저장 자체가 불가능하다( $O(n^2)$ 메모리, $O(n^3)$ 역행렬). 핵심 관찰은 Fisher가 층 구조와 Kronecker 구조를 갖는다는 것이다.

층 $\ell$ 의 출력 $z = Wa$ 에서 $\nabla_W L = g a^T$ 이므로:

F_\ell = \mathbb{E}[(a \otimes g)(a \otimes g)^T] = \mathbb{E}[aa^T \otimes gg^T]

K-FAC(Martens & Grosse 2015)은 $a$ 와 $g$ 가 독립이라 근사하여

F_\ell \approx \underbrace{\mathbb{E}[aa^T]}_{A_\ell} \otimes \underbrace{\mathbb{E}[gg^T]}_{G_\ell}

으로 분해한다. Kronecker 역행렬 $(A \otimes G)^{-1} = A^{-1} \otimes G^{-1}$ 과 $\operatorname{vec}$ 공식을 결합하면 업데이트가 닫힌 형태로 나온다.

\Delta W_\ell = -\eta \, G_\ell^{-1} (\nabla_{W_\ell} L) A_\ell^{-1}

저장은 $O(p^2 + q^2)$ , 계산은 $O(p^3 + q^3)$ 으로 원래 $O((pq)^2)$ , $O((pq)^3)$ 에서 극적으로 감소한다. Shampoo(Gupta+ 2018)는 이를 일반 텐서로 확장해 각 축마다 preconditioner를 누적하며, Anil+ (2020)은 Shampoo가 full-matrix AdaGrad의 최적 Kronecker 근사임을 보였다.

TRPO(Schulman+ 2015)는 다른 접근을 택한다. $F^{-1}v$ 를 conjugate gradient로 푸는데, 각 반복에서 Fisher-vector product $Fv = \nabla_\theta(\nabla_\theta L \cdot v)$ 만 계산한다. 두 번의 backward pass, $O(n)$ 비용으로 10~20 CG 반복이면 충분한 근사를 얻는다.

✎ 트레이드오프

Natural Gradient는 parameterization 불변이고 지수족 MLE에서 Newton과 동치다. 그 대가는 Fisher 역행렬의 계산 비용이다. K-FAC은 Kronecker 구조로 이를 줄이지만 $a$ 와 $g$ 의 독립 가정에 의존하고, Empirical Fisher( $y_i$ 를 데이터에서 취함)는 $\theta$ 가 최적에 수렴할수록 추정치가 0으로 붕괴해 업데이트가 폭주할 수 있다(Kunstner+ 2019). True Fisher는 $y \sim p_\theta$ 에서 샘플링해야 한다.

정리

$\nabla L$ 은 covector다. 유클리드에서만 벡터(방향)와 일치하며, 이 일치는 $g = I$ 라는 특수 가정의 귀결이다.
Natural Gradient $\tilde{\nabla} L = F^{-1}\nabla L$ 은 KL ball 내 steepest descent이며, 지수족 MLE에서 Newton’s method와 동치다.
Fisher의 공변 변환 법칙 덕분에 Natural Gradient는 반변 벡터처럼 변환되고, 연속 시간 flow는 좌표 선택에 무관한 분포 궤적을 그린다.
K-FAC, Shampoo, TRPO-CG는 모두 이 불변 방향을 현실적 비용으로 근사하는 전략이다.

Adam의 per-parameter scaling이 Fisher 대각 근사이고, BatchNorm이 암묵적으로 Fisher를 등방화하는 재매개변수화라는 사실은 우연이 아니다 — 실용적 딥러닝 최적화의 많은 “tricks”는 Natural Gradient 철학의 공학적 근사다.