Риманова геометрия нейросетей

GPT-4 имеет 1.7 триллиона параметров. Loss landscape - в 1.7-триллионо-мерном пространстве. Никто никогда не визуализировал это пространство. Но информационная геометрия предсказывает: миллионы минимумов, все эквивалентны по generalization, все соединены путями низкой loss. Это не философия - это практика: model souping, K-FAC, NTK theory - все используют этот геометрический факт.

• muP parameterization (Microsoft 2021): hyperparameters от small models переносятся на 100B+ через NTK theory
• SAM optimizer: sharpness-aware training даёт +1-2% top-1 accuracy на ImageNet
• Distributed Shampoo (DeepMind): natural gradient на 100B-parameter scale, AlphaFold/Gemini training

Neural Tangent Kernel: линеаризация больших сетей

**Neural Tangent Kernel** (Jacot, Gabriel, Hongler 2018): в пределе бесконечной ширины нейросеть обучается как kernel regression с фиксированным kernel $\Theta(x, x') = \langle \nabla_\theta f(x), \nabla_\theta f(x') \rangle$. Динамика обучения линеаризуется вокруг initialization, и сеть становится математически эквивалентна kernel machine.

**Геометрический смысл:** $\Theta$ - это риманова метрика на пространстве функций. Градиентный спуск по параметрам $\theta$ соответствует градиентному спуску в функциональном пространстве с метрикой NTK: $\partial_t f_t(x) = -\eta \sum_{x'} \Theta(x, x') (f_t(x') - y(x'))$. Обучение - это flow по этой метрике.

**Lazy training vs feature learning.** В NTK-режиме веса почти не двигаются от initialization - линеаризуется только output. Для small width и больших step sizes происходит feature learning: представления реально перестраиваются. Современные LLM находятся между двумя режимами, ближе к lazy при достаточной ширине.

**muP parameterization** (Yang & Hu 2021) используется в Microsoft для hyperparameter transfer от small к large models. Идея: правильное масштабирование learning rate по ширине сохраняет NTK-инварианты. Тюнинг hyperparameters на 1B-модели → перенос на 100B без переподбора. Экономия: тысячи GPU-часов.

NTK не учитывает feature learning, который критичен для transfer learning и downstream-задач. Чистая NTK-теория предсказывает generalization, но не объясняет, почему pre-trained features универсальны. Современные работы (Tensor Programs IV, Yang) расширяют NTK на feature-learning режим.

Loss landscape: геометрия минимумов

**Loss landscape** - функция потерь $L(\theta)$ на пространстве параметров. Невыпуклая, теоретически имеет миллионы локальных минимумов. Эмпирически (Choromanska et al. 2015, Goodfellow 2014): большинство минимумов 'хорошие' и неотличимы по generalization. Нейросети не страдают от bad minima на практике.

**Hessian spectrum** обученных сетей low-rank: большинство собственных значений $\lambda_i \approx 0$ (плоские направления), несколько dominant направлений задают effective dimension. Информационно-геометрически: касательное пространство в минимуме - почти вырождено, истинная размерность задачи $\ll \dim \theta$.

**Sharpness vs flatness debate** (Keskar 2017, Dinh 2017): большие batches приводят к sharper минимумам $\to$ хуже generalization. Sharp минимум - высокая кривизна Hessian, flat - низкая. **SAM** (Sharpness-Aware Minimization, Foret 2021) явно оптимизирует $\max_{\|\varepsilon\| \leq \rho} L(\theta + \varepsilon)$, выбирая плоские минимумы и достигая SOTA на ImageNet.

SAM используется в Vision Transformer training (Chen et al. 2021): +2.7% top-1 на ImageNet для ViT-S/16. Lion optimizer (Google 2023) явно учитывает кривизну landscape через sign-based updates - на 100B-моделях даёт 2-15% economy compute. Geometry of landscape переходит из теории в production.

Sharpness не инвариантна к репараметризациям (Dinh 2017): можно сделать любой минимум sharp или flat подменой координат на $\theta' = \alpha\theta$. Корректная мера - sharpness в Fisher metric или PAC-Bayes bounds. Это reconciles с практическим успехом SAM: SAM работает в 'natural' parameterization SGD.

Mode connectivity: пути между минимумами

**Mode connectivity** (Garipov 2018, Draxler 2018) - эмпирическое открытие: ВСЕ минимумы нейросетей соединены кривыми низкой loss. Никаких изолированных долин - есть единое связное многообразие хороших решений. Это переворачивает классическое представление о невыпуклой оптимизации.

**Permutation symmetry** (Entezari, Sedghi 2022): нейросеть инвариантна к перестановке нейронов в скрытых слоях - это даёт $\sim n!$ эквивалентных минимумов в $n$-мерных слоях. Если permute neurons одного решения, любые два минимума становятся соединимы прямой линией - **linear mode connectivity**.

**LMC после warmup** (Frankle et al. 2020): после нескольких эпох обучения, два независимых runs с одной точки становятся linearly connected. Геометрически: SGD первые шаги выбирает 'basin of attraction', потом всё дрейфует внутри одного аттрактора. Этот переход - lottery ticket hypothesis в пространстве оптимизации.

**Model souping** (Wortsman 2022) использует LMC: усреднение весов fine-tuned моделей даёт +1-3% accuracy без extra compute. Применяется в Anthropic, Google для merging mixture-of-experts и model fusion. Geometric fact (LMC) $\to$ production technique с миллионными savings.

**Теоретическая импликация:** SGD не находит уникальный минимум - он находит точку на богатом многообразии эквивалентных решений. Симметрия group action permutations $S_n$ редуцирует пространство до quotient manifold. Информационная геометрия предсказывает: реальная размерность задачи определяется дуальной плоской структурой.

Natural gradient в production: K-FAC и Shampoo

**K-FAC** (Kronecker-Factored Approximate Curvature, Martens & Grosse 2015) - блочно-диагональная аппроксимация Fisher information matrix через Kronecker factorization: $F \approx A \otimes B$, где $A, B$ - layer-wise activation и gradient covariance статистики. Превращает $O(n^2)$ inverse в два $O(\sqrt{n}^2)$ inverse.

**Shampoo** (Anil, Gupta 2020) - full-matrix preconditioner на уровне каждого parameter tensor: $G \to L^{-1/4} G R^{-1/4}$, где $L, R$ - left/right preconditioners. **Distributed Shampoo** (2023) масштабируется на 100B+ параметров через распределённое вычисление матричных корней. Used by DeepMind for AlphaFold training.

**Sophia** (Liu et al. 2023) - diagonal Hessian preconditioner specifically для LLMs: $\theta \leftarrow \theta - \eta \cdot \text{clip}(g / \hat{H}, \rho)$. Использует stochastic Hutchinson estimator для $\hat{H}$. На GPT-2 scale даёт 2x faster convergence чем Adam. Геометрически: diagonal $H$ - локальная аппроксимация Fisher.

Преимущество над Adam: K-FAC использует кросс-параметрическую кривизну (Kronecker structure), Adam - только diagonal. Разница 2-3x в скорости training больших моделей. Стоимость GPT-4 training: USD 100M+; reduction в шагах = USD millions saved. Distributed Shampoo training partially used in Gemini and Claude development.

K-FAC требует $O(n^2)$ памяти на блок (где $n$ - dim слоя). Для трансформеров с $d_{model} = 12288$ это уже терабайты. Shampoo решает через distributed sharding факторов. Production trade-off: натуральный градиент даёт 2-3x convergence speedup, но удваивает memory footprint - окупается на длинных обучениях.

Итоги

• NTK $\Theta(x, x') = \langle \nabla_\theta f(x), \nabla_\theta f(x') \rangle$ - метрика на пространстве функций; в пределе бесконечной ширины обучение становится kernel regression
• Loss landscape больших сетей имеет low-rank Hessian; sharpness vs flatness определяет generalization, и SAM явно оптимизирует плоские минимумы
• Mode connectivity: все минимумы соединены кривыми низкой loss; permutation symmetry даёт линейную связность - основа model souping
• K-FAC и Shampoo - блочные аппроксимации Fisher information matrix; Sophia использует диагональ гессиана и ускоряет LLM training в 2x

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

Куда ведёт риманова геометрия нейросетей:

• Natural gradient — K-FAC и Shampoo - production-аппроксимации natural gradient
• IG в deep learning — Базовый язык для NTK, K-FAC, mode connectivity
• Wasserstein gradient flows — Альтернативная геометрия для loss landscape

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

• NTK предполагает infinite width и lazy training. В каких условиях современные LLM (Llama, GPT) выходят за рамки NTK-режима, и какие геометрические эффекты это создаёт?
• Если все минимумы соединены путями низкой loss и эквивалентны по generalization, имеет ли смысл говорить о 'правильном' решении задачи обучения? Что определяет выбор конкретной точки на этом многообразии?
• K-FAC даёт 2-3x ускорение, но требует $O(n^2)$ памяти на блок. При каких размерах модели trade-off с Adam перестаёт быть выгодным, и почему Shampoo решает эту проблему через distributed factorization?

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

• ig-07-natural-gradient — Modern approximations like K-FAC realize natural gradient at scale
• ig-10-deep-learning — DL geometry foundation
• ig-13-generative — Both apply IG to modern ML
• ot-14-gradient-flows — Wasserstein flows offer alternative landscape view
• lt-01-pac-intro — NTK predicts generalization bounds for wide networks