Тот или иной механизм предобуславливания заложен во всех современных оптимизаторах. Адам использует диагональный предобуславливатель, не учитывающий взаимные корреляции между параметрами. Шампунь же опирается на внедиагональное приближение, имея при этом посильную сложность
O(m^2 + n^2) для матрицы весов размера (m, n).В Shampoo поддерживаются две матрицы L и R из произведений градиентов. От них считается степень -¼.
В оригинальной реализации считают разложение матриц на собственные числа. Данная процедура представляет собой итеративный процесс плохо параллелизующийся на несколько GPU.
Взамен этого авторы предлагают два варианта итеративных алгоритмов для подсчета матричных корней:
1️⃣ CN - Coupled Newton Iteration
2️⃣ NDB: The Newton-Denman-Beavers 🦫 Iteration
Оба варианта требуют только операций матричного умножения, поэтому куда более hardware friendly. Кроме того, они более численно устойчивы, и позволяют проводить операции в fp16 вместо fp32.
Далее авторы анализируют сходимость. В оригинальном распределенном Shampoo для того, чтобы спектр матрицы вписывалась в единичный круг делили на норму Фробениуса, которая может быть заметно больше спектральной, что замедляет сходимость. Здесь же используют оценку на спектральную норму из частного Рэлея.
Чтобы понизить асимптотическую сложность, предобуславливатель заменяют на блочно-диагональный с размером блока ~1k. Для батчевой обработки предобуславливатели для L и R стакаются в один тезнор и считаются одновременно.
🧪 Эксперименты
Для валидации метода обучают семейство моделей архитектуры Llama-2 в Шиншилла-оптимальном сетапе.
При использовании разложения на собственные значения DASH дает ускорение порядка 20%. Для CN и NDB итераций ускорение от DASH в сравнении с распределенным Shampoo доходит до 3-4 раз. При этом валидационная перплексия на уровне бейзлайна в пределах шума.
💡 Выводы
Выглядит как серьезное ускорение Shampoo. Интересно 🤔, насколько оно окажется конкурентно способно Мюону в более боевой постановке - на большой модели и большом количестве данных.