Материал подготовлен автоматически по первоисточникам: ссылки на них — в конце статьи.
Кривизну функции потерь научились учитывать при обучении больших нейросетей без линейного поиска и ручной коррекции неудачных оценок. В препринте Joohwan Ko, Tetiana Parshakova, Diana Cai и Robert M. Gower, который не прошёл рецензирование и содержит числа, полученные самими авторами, оптимизатор SoftServe показал лучшие результаты прежде всего на плохо обусловленных задачах, где градиент сильно меняется по разным направлениям. Для языковой модели преимущество не подтвердилось, поэтому метод выглядит не общей заменой Adam, Muon или SOAP, а кандидатом для отдельных классов моделей.
Как SoftServe обращается с отрицательной кривизной
Обычный градиентный оптимизатор выбирает направление по первой производной функции потерь. Квазиньютоновский метод дополнительно приближает обратную матрицу Гессе — матрицу вторых производных, которая описывает кривизну функции. Такая оценка помогает выбирать разные масштабы шага для направлений, где поверхность потерь меняется резко или плавно.
В глубоком обучении этот подход сталкивается с отрицательной кривизной. Если приблизить её неудачно, преобразованный градиент может направить параметры к росту потерь. Классический BFGS избегает этого с помощью линейного поиска, но тот требует повторно считать функцию на полном наборе данных и плохо сочетается со случайными мини-пакетами.
SoftServe не требует, чтобы новая оценка кривизны точно удовлетворяла секущему условию — связи между изменением параметров и изменением градиента. Вместо жёсткого равенства метод штрафует расхождение и одновременно удерживает новую матрицу рядом с предыдущей. Полученная оценка остаётся положительно определённой: преобразованный градиент сохраняет направление спуска даже там, где реальная кривизна отрицательна.
Для случайного обучения SoftServe считает градиенты в начале и конце интервала с одним и тем же мини-пакетом и теми же источниками случайности. Так разница градиентов отражает изменение поверхности потерь, а не смену данных, масок прореживания или преобразований.
Как оценку кривизны уместили в память GPU
Полную матрицу для модели хранить нельзя: её размер растёт квадратично относительно числа параметров. Даже L-BFGS с десятью парами изменений параметров и градиентов потребовал бы объёма памяти, сопоставимого с двадцатью копиями модели.
SoftServe-Diag оставляет по одному положительному масштабу на параметр. Его вычислительные затраты и расход памяти растут линейно и сопоставимы с Adam, но такая версия не учитывает связи между параметрами.
SoftServe-Kron обрабатывает веса слоя как матрицу и представляет оценку кривизны двумя меньшими кронекеровскими множителями. Вместо одной огромной матрицы метод хранит две квадратные матрицы по сторонам слоя. Он обновляет их по очереди, каждый раз решая матричное квадратное уравнение с единственным положительно определённым решением.
Разложения матриц здесь заменили итерациями Newton–Schulz. Они сводят извлечение квадратного корня и обращение матриц к последовательности умножений и сложений, которые GPU выполняют эффективнее. Оценку кривизны обновляют не на каждом шаге, а через заданный интервал, чтобы распределить её стоимость между несколькими обновлениями параметров.
Где результаты меняют планы, а где пока нет
Метод проверяли на рекуррентных сетях, автоэнкодерах, моделях для решения дифференциальных уравнений и физически информированной диффузионной модели. Сравнения включали Adam, Muon, SOAP, а в малых экспериментах также квазиньютоновские варианты и связку Adam с L-BFGS.
На малых физически информированных сетях SoftServe-Kron получил наименьшую итоговую функцию потерь в пяти из шести режимов. Исключением стала детерминированная волновая задача, где ниже оказались Adam и Adam с последующим L-BFGS. На PirateNet метод был сопоставим с Muon для одного уравнения и показал наименьшую среднюю потерю для другого.
Самый крупный опыт охватывал физически информированную диффузионную модель со 136 млн параметров. После 100 тысяч вычислений градиента SoftServe-Kron достиг меньшей средней обучающей потери, чем Adam, Muon и SOAP. Это показывает, что структурированная оценка кривизны технически применима не только к небольшим сетям.
На GPT-подобной модели с 16 млн параметров картина изменилась: SoftServe-Kron обошёл AdamW по проверочной функции потерь, но уступил Muon и SOAP. Результаты поэтому не дают основания менять стандартный оптимизатор для языковых моделей только из-за этой работы.
SoftServe стоит добавить в собственное сравнение, если продукт обучает рекуррентную сеть, физически информированную модель или другую систему с плохо обусловленной функцией потерь. Сравнивать нужно не только итоговую потерю при равном числе градиентов, но и время обучения и память: матричные операции могут окупиться не на каждой архитектуре.
Переход также потребует настройки коэффициента мягкого секущего штрафа и интервала обновления кривизны. В экспериментах коэффициент оставался фиксированным в пределах запуска, а доказательства сходимости для реалистичного режима глубокого обучения работа не даёт. Поэтому практический план — испытание рядом с текущим оптимизатором, а не немедленная замена базового контура обучения.
Источники
Иллюстрация: рисунок из статьи «SoftServe: A Scalable Quasi-Newton Method for Deep Learning», Joohwan Ko, Tetiana Parshakova, Diana Cai и др., CC BY 4.0
Похоже на вашу задачу?
Расскажите, что собираете. За полчаса разложим на этапы и назовём сроки — это бесплатно и ни к чему не обязывает.
