Информационная геометрия
K-FAC: кронекеровская аппроксимация матрицы Фишера
Матрица Фишера для GPT-3 с 175 миллиардами параметров имеет размер 3*10^22 - хранить её физически невозможно. K-FAC (2015) решает это через кронекеровскую факторизацию, снижая память с O(n^4) до O(n^2). Google Brain использовал K-FAC для ускорения обучения AlphaGo Zero.
- Distributed K-FAC от DeepMind (2021) обучил ResNet-50 на ImageNet за 55 эпох вместо 90 у SGD при том же качестве. Meta использует K-FAC-вдохновлённые методы в оптимизаторах для LLaMA.
Кронекеровская факторизация блоков Фишера
Матрица Фишера для нейросети размером 100M параметров имеет размер 10^16 элементов - хранить её невозможно. K-FAC (Martens, Grosse 2015) аппроксимирует каждый блок F_l для слоя l как кронекеровское произведение A_l x G_l, где A_l - ковариация входов, G_l - ковариация градиентов выходов. Это снижает память с O(n^4) до O(n^2) на слой. Google Brain применял K-FAC для обучения AlphaGo Zero и трансформеров.
Сложность обращения K-FAC блока для слоя с n входами и m выходами...
Сходимость и практические настройки K-FAC
K-FAC требует периодического обновления статистик A и G (каждые 10-100 шагов) и инвертирования (каждые 100-500 шагов). Martens & Grosse показали, что при правильном демпинге K-FAC сходится за 10-20x меньше итераций, чем SGD, на задачах обучения рекуррентных сетей. На ResNet-50 (ImageNet) K-FAC достигает 75% top-1 за 55 эпох против 90 для SGD.
Почему K-FAC не обновляет статистики A и G на каждом шаге?
Распределённый K-FAC и расширения
Распределённый K-FAC (Osawa et al., 2019; Pauloski et al., 2021) масштабирует алгоритм на сотни GPU: статистики A_l и G_l вычисляются параллельно по слоям, инвертирование распределяется между процессами. На ImageNet с 512 GPU K-FAC достигает точности SGD за 35 эпох вместо 90.
Расширения K-FAC: EKFAC (George et al., 2018) дополнительно rescale по собственным значениям для лучшего приближения; KFRA учитывает нелинейности активационных функций; Shampoo (Gupta et al.) - обобщение K-FAC на тензорные параметры.
Какое преимущество даёт eigen-decomposition факторов A_l и G_l в K-FAC?
В eigen-базисе кронекерово произведение диагонально: (Λ_A ⊗ Λ_G)_{ii} = λ_i^A · λ_i^G. Инвертирование - просто деление каждого элемента на λ_i^A · λ_i^G + λ_damp, что занимает O(nm) вместо O((nm)^3).
Ключевые результаты
- K-FAC аппроксимирует блок Фишера слоя как A_l x G_l.
- Обращение занимает O(n^3 + m^3) вместо O((nm)^3).
- Статистики обновляются с momentum каждые T шагов, инвертирование - каждые T' шагов.
- На практике K-FAC даёт 2-10x ускорение по числу итераций при сопоставимом времени на шаг.
- Eigen-decomposition позволяет инвертировать (A⊗G+λI) за O(nm) в eigen-базисе.