Математический взгляд на градиентный спуск: концепция модульной двойственности от ученых MIT

Корректен ли привычный шаг оптимизации?

Градиентный спуск служит базовым инструментом обучения нейронных сетей, однако с точки зрения строгой математики стандартный алгоритм обновления весов имеет скрытый изъян. В работе лаборатории MIT CSAIL под названием «Modular Duality in Deep Learning» авторы обращают внимание на то, что параметры сети и градиенты принадлежат принципиально разным пространствам: веса находятся в векторном пространстве, а градиенты — в сопряженном (выступая в роли ковекторов).

Прямое вычитание градиента из весов без преобразования игнорирует тот факт, что кривизна функции потерь может быть крайне неоднородной по разным направлениям. Для математически выверенного шага необходимо использовать двойственное отображение, которое согласует пространства с помощью заданной нормы.

Концепция модульной двойственности

Авторы статьи предлагают формализовать компоненты нейросети через понятие изолированного модуля. Каждый такой модуль характеризуется четырьмя параметрами:

• Функция прямого прохода (forward): стандартное преобразование входного вектора и весов в выходные значения;
• Масса (mass): гиперпараметр, определяющий вклад модуля в обучение всей системы и регулирующий скорость сходимости;
• Чувствительность (sensitivity): мера стабильности выхода при небольших возмущениях на входе;
• Норма весов (norm): показатель сложности параметров, отвечающий за контроль емкости и обобщающую способность.

Если модуль удовлетворяет условию липшицевости (так называемый «well-normed module»), оптимизация становится устойчивой: шаг по весам не приводит к взрыву активаций, а архитектура в целом приобретает повышенную стойкость к шумам и состязательным атакам (adversarial attacks).

Связь с оптимизатором Shampoo и ускорение вычислений

На примере классического линейного слоя исследователи продемонстрировали, что переход через двойственное отображение заменяет матрицу градиентов ее полярной проекцией, где сингулярные числа приравниваются к единице. Данное преобразование математически родственно известному матричному оптимизатору Shampoo.

Кроме того, возникающий коэффициент спектрального масштабирования точно совпадает с принципами параметризации µP (muP). Это позволяет свободно масштабировать ширину слоев сети без необходимости заново подбирать скорость обучения.

Поскольку классическое сингулярное разложение (SVD) вычислительно затратно, авторы оптимизировали расчеты с помощью прямоугольных итераций Ньютона — Шульца, что позволило добиться высокой скорости без потери точности.

Практическая значимость

Предложенный фундаментальный аппарат решает сразу несколько важных задач в машинном обучении:

• Формирует единый математический базис и строгую систему типов для пространств активаций;
• Обеспечивает быстрое и стабильное обучение моделей (проверено экспериментами на базе NanoGPT);
• Помогает отслеживать согласованность обновлений весов с входящими активациями для защиты от вычислительных шумов.