Семинар MTML Lab «Динамическая релаксация: от Signum и Muon к мягким правилам обновления»

О чем пойдет речь?

В докладе я хочу обсудить проблему, которая возникает у оптимизаторов вроде Signum и Muon. Такие методы часто выбирают не точный градиентный шаг, а его более грубую версию: например, берут только знак координат или ортогонализуют обновление. На ранних стадиях обучения это часто полезно. Но ближе к концу оказывается, что такой шаг слишком грубый: когда одним параметрам еще нужно двигаться далеко, другим уже нужна очень точная и маленькая подстройка. Из-за этого метод может начать колебаться там, где хотелось бы уже спокойно сходиться.

Идея работы состоит в том, чтобы сделать этот переход от "жестких" обновлений к более мягким не вручную и не эвристически, а на уровне общей конструкции. В координатном случае sign можно заменить его сглаженной версией, например через tanh (ту самую странную функцию из калькулятора), чтобы шаг постепенно переставал быть чисто знаковым и начинал учитывать масштаб. В спектральном случае, если хочется так же смягчить Muon, возникает уже другая нелинейность: там появляется алгебраический аналог tanh, связанный с итерациями Newton-Schulz (NS).

Мотивация понятна, а как вообще такие вещи доказывать? Я попытаюсь ответить на этот вопрос. Речь пойдет скорее об общей идее доказательств: как построить язык, в котором можно аккуратно описывать переход от грубых, но устойчивых шагов к более точным. По дороге придется вспомнить и сопряженные функции по Фенхелю (да, не удивляйтесь, они действительно могут пригодиться, и, возможно, учили вы их не совсем зря).

Если останется время, в конце я коротко расскажу еще об одном любопытном сюжете: похоже, что у итераций NS в Muon-подобных методах может быть свое оптимальное число шагов, так что делать их бесконечно много, возможно, не только не нужно, но и не является лучшей стратегией.