Натуральный градиент в глубоком обучении

Adam - самый популярный оптимизатор. Но Adam - диагональная аппроксимация натурального градиента без учёта кросс-корреляций параметров. K-FAC использует Кронекер-факторизованную матрицу Фишера и ускоряет сходимость в 2-4 раза. Distributed Shampoo масштабируется на 100B+ параметров и используется при обучении Gemini. Это не академические упражнения - это миллионы долларов экономии вычислений.

• K-FAC (Martens & Grosse 2015): обучение ResNet-50 на ImageNet в 2.5x меньше шагов чем Adam
• Distributed Shampoo (Google 2023): используется при обучении AlphaFold 2 и частично Gemini
• Sophia optimizer (Stanford 2023): diagonal Hessian приближение, 2x faster LLM training на GPT-2 scale

K-FAC: Кронекер-факторизация матрицы Фишера

Для нейронной сети с L слоями матрица Фишера блочно-диагональна по слоям (приближение). В каждом слое l: входная активация a_{l-1} ∈ R^n, выходной градиент g_l ∈ R^m, вес W ∈ R^{m x n}. Блок Фишера для слоя l: F_l = E[g_l g_l^T] ⊗ E[a_{l-1} a_{l-1}^T] = G_l ⊗ A_l.

Damping parameter в K-FAC играет роль learning rate scale. Слишком малый damping - нестабильность (матрица Фишера вырождена). Слишком большой - приближаемся к обычному SGD. Adaptive damping (Levenberg-Marquardt rule) автоматически подстраивает его по изменению loss.

Shampoo и Distributed Shampoo: масштабирование

Shampoo (Gupta et al. 2018): для матрицы параметров W ∈ R^{m x n} поддерживает два накопленных фактора L = sum G_t G_t^T и R = sum G_t^T G_t. Обновление: W = W - lr * L^{-1/4} G R^{-1/4}. Это full-matrix preconditioner для каждой матрицы весов.

Distributed Shampoo (2023): Kronecker факторы L и R для большого слоя (например, embedding с dim=32000) требуют O(32000^2) = 1 GB памяти каждый. Решение: шардировать строки/столбцы L и R между GPU через ZeRO-style partitioning. Каждый GPU хранит свой кусок и вычисляет частичный корень. Это позволяет применять Shampoo к transformer-моделям на 100B параметров.

Riemannian SGD и оптимизация на многообразиях

Некоторые параметры нейросетей живут не в евклидовом пространстве. Параметры Stiefel-многообразия (ортогональные матрицы): Q^T Q = I. Гиперболические эмбеддинги: норма строго меньше 1. Симплекс вероятностей: sum p_i = 1. Для таких параметров натуральный градиентный спуск требует Retraction - проекции обратно на многообразие.

Adam как приближение натурального градиента

Adam: theta = theta - lr * m_t / (sqrt(v_t) + eps). Здесь v_t = E[g^2] - диагональная оценка второго момента градиента. Диагональ матрицы Фишера: F_{ii} = E[(d log p / d theta_i)^2] ≈ E[g_i^2] = v_t. Таким образом, Adam - диагональное приближение натурального градиента.

Что Adam упускает: 1) Кросс-корреляции между параметрами (K-FAC учитывает). 2) Структуру матрицы Фишера (не просто второй момент градиента). 3) Геометрию параметрического многообразия для structured parameters. Поэтому K-FAC и Shampoo превосходят Adam на задачах с strong inter-parameter correlations.

Несмотря на теоретическое превосходство K-FAC, на практике Adam часто выигрывает из-за: 1) Простоты тюнинга (один lr вместо lr + damping + update_freq). 2) Меньших затрат памяти. 3) Хорошей работы с adaptivity на разных задачах. K-FAC выгоден только при очень длинном обучении, где экономия шагов перевешивает стоимость вычислений.

Итоги

• K-FAC: $F \approx A \otimes G$ для каждого слоя; natural gradient $= G^{-1} W_{\text{grad}} A^{-1}$; $O(\sqrt{n}^2)$ вместо $O(n^2)$
• Shampoo: full-matrix preconditioner $W \to L^{-1/4} G R^{-1/4}$; Distributed Shampoo шардирует факторы по GPU
• Riemannian SGD: обновление $\theta \leftarrow \exp_{\theta}(-\eta \cdot F^{-1}\nabla L)$ с retraction на многообразие
• Adam как приближение: $g/(\sqrt{v}+\epsilon)$ - diagonal rescaling без кросс-корреляций; K-FAC учитывает блочную структуру

Связанные темы

Натуральный градиент в production:

• Natural gradient basics — Теоретическая основа
• IG в deep learning — Landscape и Fisher metric
• Риманова геометрия нейросетей — K-FAC, NTK, mode connectivity

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

• K-FAC требует O(n^2) памяти на блок. При каких размерах слоёв это становится непрактичным, и как Distributed Shampoo решает эту проблему через шардирование?
• Adam на практике часто работает лучше K-FAC несмотря на теоретическое превосходство натурального градиента. Какие свойства реальных задач объясняют этот парадокс?
• Sophia использует диагональный гессиан вместо матрицы Фишера. В каком смысле это «лучше» Adam? Что теряется по сравнению с K-FAC?

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

• ig-07-natural-gradient — ig-17 - production-реализации натурального градиента
• ig-14-neural — K-FAC и Shampoo - из ig-14
• ig-16-alpha-divergences — Альфа-геометрия объясняет выбор метрики