Информационная геометрия

Natural Gradient: ковариантная производная distributions

В 1998 Amari задал странный вопрос: что значит шагнуть на $\epsilon$ в пространстве распределений? Если параметризация может быть $\sigma$, $\log\sigma$ или $\sigma^2$ - получаются три разных SGD с тремя разными траекториями. Natural gradient убирает этот произвол: одна геометрия, одна direction, инвариант к репараметризации. Это та же разница, что между евклидовой геометрией и римановой - и для нейросетей она настоящая.

  • **TRPO и PPO в RL.** OpenAI Five (Dota 2) и DeepMind AlphaStar тренировались на natural policy gradient в trust region формулировке. Без KL-constraint policy collapsing случается за десятки итераций.
  • **K-FAC в DeepMind.** AlphaGo Zero policy network обучалась через K-FAC. Sophia (Liu 2023) ускорил GPT-2-уровня pretraining в 2x по сравнению с Adam на тех же flops.
  • **RLHF в Claude/ChatGPT.** PPO clip - это суррогат natural gradient KL-constraint: clip ratio $\frac{\pi_\theta}{\pi_{\text{old}}}$ - грубое приближение KL trust region из TRPO.

Предварительные знания

  • Fisher information metric
  • KL и Bregman divergences

Определение: ковариантная производная distributions

Vanilla gradient зависит от параметризации

Обычный градиент $\nabla L(\theta)$ - это вектор частных производных по координатам $\theta$. Проблема: смена параметризации меняет градиент. Если для гауссиана взять $\theta=(\mu, \sigma)$ или $\theta=(\mu, \log\sigma)$ - получатся разные шаги SGD, разные траектории, разная скорость сходимости. Геометрически бессмысленно: distribution одна и та же.

Натуральный градиент $\tilde{\nabla} L$ - это vanilla градиент, домноженный на обратную Fisher information matrix $F(\theta)$. Ключевое свойство: при репараметризации $F$ преобразуется как метрический тензор, а $\nabla L$ - как ковектор. Их произведение остаётся инвариантным. Direction шага в пространстве distributions один и тот же независимо от того, в каких координатах считаем.

**Fisher information matrix:** $F_{ij}(\theta) = \mathbb{E}_{x\sim p_\theta}\big[\partial_i \log p_\theta(x)\, \partial_j \log p_\theta(x)\big]$. Это метрика на манифольде распределений; её также называют Fisher-Rao metric.

Геометрический смысл: natural gradient - direction steepest descent в KL-divergence, а не в евклидовой норме параметров. Шаг $\theta \leftarrow \theta - \eta \tilde{\nabla} L$ минимизирует loss среди всех шагов фиксированного KL-радиуса. Vanilla SGD минимизирует среди шагов фиксированного евклидова радиуса - что метрически произвольно.

Откуда взялся natural gradient

Amari (1998) формализовал natural gradient как ковариантную производную на статистическом манифольде. Идея вызревала десятилетиями: Akaike (1981) применял Fisher для model selection (AIC), Кульбак (1959) определил KL-divergence как меру информационного расхождения. Amari соединил это в один framework: information geometry. Прорыв случился, когда Peters & Schaal (2008) показали на робототехнике, что natural policy gradient сходится в 5-10 раз быстрее vanilla policy gradient.

Чем natural gradient $\tilde{\nabla} L = F^{-1} \nabla L$ отличается от vanilla градиента $\nabla L$?

Риманова геометрия и обратная Fisher matrix

Manifold распределений и Fisher metric

Параметрическое семейство $\{p(\cdot|\theta) : \theta \in \Theta\}$ - это многомерный манифольд: каждой точке $\theta$ соответствует distribution. Fisher metric $F(\theta)$ задаёт скалярное произведение в касательном пространстве: расстояние между близкими distributions измеряется через KL, и в локальном пределе KL квадратична через $F$.

Это уравнение - сердце natural gradient. Минимизация loss $L(\theta + \epsilon) - L(\theta) \approx \nabla L^T \epsilon$ при ограничении $D_{KL} \le \delta$ даёт направление $\epsilon^* \propto -F^{-1} \nabla L$. То есть natural gradient - это решение задачи trust region в KL-метрике.

**Геодезики и parallel transport.** SGD по евклидовой метрике движется по прямым в $\theta$-пространстве - но прямая в $\theta$ - это не прямая на манифольде distributions. Natural gradient следует геодезикам Fisher-Rao: кратчайшим путям в смысле KL. Christoffel symbols описывают, как касательные векторы переносятся вдоль кривых; для exponential family они имеют замкнутую форму.

**Связь с TRPO в RL.** Trust Region Policy Optimization (Schulman 2015) явно решает: $\max_\theta \mathbb{E}[A(s,a)] $ при $\mathbb{E}[D_{KL}(\pi_{\theta_{\text{old}}} \| \pi_\theta)] \le \delta$. Линеаризация целевой и квадратичная аппроксимация ограничения - и появляется natural gradient. PPO - дешёвая эвристическая замена этого constraint через clipping.

Что измеряет Fisher metric $F(\theta)$ в локальном смысле?

K-FAC и Shampoo: практические аппроксимации

Почему прямое $F^{-1}$ невозможно

Для модели с $d$ параметрами Fisher matrix имеет размер $d \times d$. Хранение - $O(d^2)$, инверсия - $O(d^3)$. Для GPT-2 small ($d \approx 10^8$) это $10^{16}$ FLOPS на инверсию каждый шаг - неприемлемо. Все production-методы natural gradient - это аппроксимации $F^{-1}$, торгующие точность на cost.

  • Diagonal Fisher — Только диагональ $F$. Cost: $O(d)$. Используется в AdaGrad, RMSProp, Adam (как эвристика). Точность плохая - игнорирует корреляции между параметрами, но cheap.
  • Empirical Fisher — Заменяем $\mathbb{E}_{p_\theta}$ на среднее по batch: $F \approx \tfrac{1}{B}\sum \nabla \log p \cdot \nabla \log p^T$. Дешевле истинной Fisher, но смещает оценку при mismatch model/data.
  • K-FAC (Kronecker-Factored) — Для слоя сети $F \approx A \otimes G$, где $A$ - covariance активаций, $G$ - covariance градиентов. Тогда $F^{-1} \approx A^{-1} \otimes G^{-1}$. Cost: $O(\text{layer\_dim}^3)$ вместо $O(d^3)$. Martens & Grosse 2015.
  • Shampoo / Sophia — Block-diagonal preconditioners второго порядка. Shampoo (Gupta 2018) использует Kronecker-factored statistics похоже K-FAC, но как general-purpose optimizer. Sophia (Liu 2023) - diagonal Hessian estimator для LLM.
  • Conjugate Gradient — Решаем линейную систему $F x = g$ итеративно через matrix-vector products $F v$. $F v$ можно вычислить за один backward pass без материализации $F$. Hessian-Free optimization (Martens 2010).

**Ответ на вопрос 'какая аппроксимация в Adam?'.** Adam использует diagonal empirical Fisher (через $v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2$). Это natural gradient в самом грубом приближении. Именно поэтому Adam часто называют 'cheap natural gradient' - он эмпирически близок, но игнорирует off-diagonal Fisher.

**Где K-FAC реально применялся.** DeepMind использовал K-FAC для policy training в AlphaGo и AlphaZero. Sophia (Liu et al. 2023) показал ускорение GPT-2-уровня pretraining в 2x по сравнению с Adam. Shampoo был использован Google для PaLM. Trade-off: компилятор и memory сложнее, но wall-clock convergence лучше для больших batch.

Natural gradient всегда лучше Adam, просто слишком дорог.

Natural gradient и Adam решают разные задачи. Natural gradient - принципиальный second-order метод; Adam - diagonal-Fisher эвристика. Когда compute позволяет (RL с маленькими сетями, fine-tuning) и градиенты сильно скоррелированы - natural gradient выигрывает 5-10x. Для огромных DNN с миллиардами параметров и стохастичным batch Adam часто достаточен, потому что off-diagonal Fisher шумит и дорого считается.

Это не вопрос 'лучше/хуже', а trade-off между принципиальностью и cost. K-FAC и Shampoo занимают середину - block-diagonal aproximation позволяет использовать natural gradient на больших моделях.

Ключевые идеи

  • Natural gradient $\tilde{\nabla}L = F^{-1}\nabla L$ инвариантен к репараметризации - direction шага одинаков в любых координатах
  • Геометрически это steepest descent в KL-divergence: $D_{KL}(p_\theta \| p_{\theta+\epsilon}) \approx \tfrac12 \epsilon^T F \epsilon$, и natural gradient - решение trust-region задачи
  • Прямое $F^{-1}$ стоит $O(d^3)$ - для нейросетей нужны аппроксимации: diagonal Fisher (Adam), K-FAC, Shampoo, Sophia, conjugate gradient
  • Natural gradient - принципиальный фундамент TRPO, PPO, natural policy gradient в RL и K-FAC/Shampoo в deep learning

Связи с другими темами

Natural gradient - центральный узел information geometry: соединяет Fisher metric с практикой оптимизации.

  • Fisher information metric — Метрика, в которой natural gradient является steepest descent. Без Fisher формула $F^{-1}\nabla L$ не имеет смысла.
  • RL и bandits — TRPO, PPO, natural Q-learning - всё построено на natural gradient. Trust region KL constraint - это и есть Fisher метрика.
  • Neural networks training — K-FAC, Shampoo, Sophia - production аппроксимации natural gradient для глубоких сетей. AlphaGo и LLM-pretraining используют их в большом масштабе.

Вопросы для размышления

  • Если параметризация модели не влияет на natural gradient, то почему vanilla SGD до сих пор стандарт - и в каких сценариях этот произвол на практике безвреден?
  • Adam = diagonal empirical Fisher - это обоснованное приближение или совпадение? Что теряется и что сохраняется по сравнению с full Fisher?
  • TRPO явно решает trust-region задачу, PPO заменяет её на clipping. Какая интуиция geometric стоит за clip ratio - и где это приближение ломается?

Связанные уроки

  • ig-02-fisher-metric — Натуральный градиент - это градиент, измеренный в Fisher-метрике; без неё формула $F^{-1}\nabla L$ не имеет смысла
  • ig-04-kl-bregman — Локальная KL квадратична в Fisher: $D_{KL}(p_\theta \| p_{\theta+\epsilon}) \approx \tfrac12 \epsilon^T F \epsilon$ - это и обосновывает natural gradient
  • ig-05-dual-flat — Двойственно-плоская структура объясняет, почему natural gradient инвариантен относительно репараметризации
  • ig-12-rl-bandits — TRPO, PPO и natural policy gradient - прямое применение natural gradient в RL
  • ig-14-neural — K-FAC, Shampoo и Sophia - аппроксимации natural gradient для глубоких сетей
  • stat-01-sampling
Natural Gradient: ковариантная производная distributions

0

1

Войти

Какая аппроксимация Fisher matrix используется в Adam?