딥러닝 최적화의 숨은 기하학 — Sharpness에서 Fisher까지

1차 방법은 gradient만 본다. 2차 방법은 곡률까지 본다. 그런데 실전 딥러닝은 이 둘의 경계를 끊임없이 오간다 — 이론적으로 불안정해야 할 learning rate에서 수렴하고, 파라미터 공간의 기하를 무시해야 할 방식으로 업데이트하면서도 좋은 해에 도달한다. 왜 이런 일이 가능한가?

이론적 경계 너머에서 학습이 일어난다

고전적인 수렴 분석은 $\eta \leq 2/L$ 조건을 요구한다. $L$-smooth 함수에서 이 경계를 넘으면 descent lemma가 깨지고 발산이 예측된다.

하지만 Cohen et al. (2021)이 full-batch GD 실험에서 관찰한 것은 달랐다. Hessian의 최대 고유값, 즉 sharpness $\lambda_{\max}(H)$는 훈련 초반에 빠르게 증가하다가 정확히 $2/\eta$ 경계에 도달한 뒤 그 근방에서 진동한다. 그리고 loss는 여전히 감소한다.

$\lambda_{\max}(H(x_t)) \;\longrightarrow\; 2/\eta \quad \text{with } f(x_t) \text{ still decreasing}$

이것이 Edge-of-Stability 현상이다. 두 국면으로 나뉜다. Phase I에서는 $\lambda_{\max} < 2/\eta$이고 loss가 빠르게 떨어진다. Phase II에서는 $\lambda_{\max} \approx 2/\eta$ 근방에서 sharpness가 진동하지만 loss는 완만하게 계속 감소한다.

이 관찰의 핵심 함의는 sharpness 증가는 loss 감소의 필연적 대가라는 것이다. 실전 딥러닝은 이 경계에서 최대 속도로 달리면서 동시에 안정성을 확보한다.

배치 크기에는 임계점이 있다

Edge-of-Stability는 full-batch GD에서 가장 명확하다. Mini-batch로 넘어오면 gradient noise가 개입한다. 이 noise의 크기를 정량화한 것이 **Gradient Noise Scale(GNS)**이다.

Mini-batch gradient의 분산은 다음과 같다.

$\text{Var}[\hat{g}] = \frac{\text{tr}(\Sigma(x))}{B}$

여기서 $\Sigma(x)$는 개별 샘플 gradient의 공분산이다. Signal-to-noise ratio가 1이 되는 배치 크기를 critical batch size라 한다.

$B_{\text{crit}} = \frac{\text{tr}(\Sigma(x))}{\|\nabla f(x)\|^2}$

$B < B_{\text{crit}}$이면 noise가 signal보다 작다. 이 영역에서는 $\eta \propto B$ 선형 스케일링이 성립한다 — 배치를 2배로 늘리면 learning rate도 2배 올릴 수 있다. $B > B_{\text{crit}}$이면 배치를 더 늘려도 수익이 감소한다. Wall-clock time은 $\sqrt{B}$에 비례해 늘어나지만 수렴 속도는 제자리다.

2차 정보를 쓰면 무엇이 달라지는가

Edge-of-Stability와 GNS는 모두 loss landscape의 곡률, 즉 Hessian 정보와 관련된다. Newton method는 이 정보를 직접 활용한다.

$x_{t+1} = x_t - H(x_t)^{-1} \nabla f(x_t)$

Quadratic 함수에서는 한 번의 step으로 optimum에 도달한다. Condition number에 대한 의존성이 사라진다. 하지만 ResNet-50(26M 파라미터)의 full Hessian은 약 2.7 페타바이트의 메모리를 요구한다. 현실에서는 불가능하다.

파라미터 공간이 아닌 분포 공간에서의 최적화

K-FAC은 단순한 공학적 근사가 아니다. 그 배경에는 Natural Gradient라는 더 근본적인 아이디어가 있다.

일반 gradient descent는 Euclidean 거리 기준으로 내리막을 찾는다. 그러나 확률 모델에서 파라미터 공간의 Euclidean 거리는 분포 간 실제 차이를 반영하지 못한다. Fisher Information Matrix는 KL divergence의 local quadratic 근사다.

$\text{KL}(p_\theta \| p_{\theta + d\theta}) \approx \frac{1}{2} (d\theta)^\top F(\theta) (d\theta)$

이 metric 하에서 최속하강 방향이 Natural Gradient다.

$\theta_{t+1} = \theta_t - \eta F(\theta_t)^{-1} \nabla L(\theta_t)$

두 가지 결정적인 성질이 있다. 첫째, reparameterization invariance — 파라미터를 $\theta \to \phi = g(\theta)$로 재정의해도 기하학적 방향이 변하지 않는다. Euclidean gradient는 재정의에 의존하지만 Natural Gradient는 분포 공간에서의 방향을 그대로 유지한다. 둘째, K-FAC은 이 Fisher를 Kronecker 구조로 근사하므로, K-FAC은 실용적인 Natural Gradient 구현이다.

정리

네 개의 챕터를 관통하는 하나의 원리가 있다 — loss landscape의 곡률을 얼마나, 어떤 방식으로 활용하는가.

• Edge-of-Stability: GD는 $\lambda_{\max} = 2/\eta$ 경계에서 sharpness를 자동으로 제어하며 최대 속도로 학습한다.
• Gradient Noise Scale: 배치 크기 $B_{\text{crit}}$를 기준으로 선형 스케일링이 유효한 영역과 수익 감소 영역이 나뉜다.
• K-FAC: Kronecker 구조로 Fisher를 근사해 Newton의 수렴 속도를 $O(p)$ 메모리로 달성한다.
• Natural Gradient: Euclidean이 아닌 정보 기하의 metric에서 최적화하면 reparameterization에 불변이 된다.

Adam이 여전히 실전의 기본값인 이유는 이 모든 것을 모르고 써도 어느 정도 작동하기 때문이다. 하지만 학습이 불안정할 때, 배치 크기 스케일링이 기대한 대로 되지 않을 때, 더 빠른 수렴이 필요할 때 — 이 기하학이 진단과 처방의 언어가 된다.