MIT представил метод стабильного обучения трансформеров с помощью липшицевых ограничений и оптимизатора Muon

Решение проблемы нестабильности больших трансформеров

Обучение больших трансформеров традиционно сталкивается с проблемой нестабильного роста активаций и скачков функции потерь. Исследователи из MIT предложили решение, которое направлено на корень проблемы: неограниченные нормы весов и активаций. Они вводят доказуемые липшицевы ограничения, регулируя спектр весов, при этом не используя распространённые приёмы стабилизации, такие как нормализация активаций, нормализация QK или ограничение логитов.

Что такое липшицево ограничение

Липшицево ограничение ограничивает максимальное изменение выхода нейросети при изменении входных данных или весов. Формально функция f является K-липшицевой, если для любых двух входов x1 и x2 разница выходов не превышает K умноженное на разницу входов. Меньшие липшицевы константы означают большую устойчивость и предсказуемость, что важно для стабильности, защиты от атак, приватности и обобщения.

Ограничения существующих методов стабилизации

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

Основные инновации: спектральная регуляция весов и оптимизатор Muon

Оптимизатор Muon спектрально регулирует градиенты, не позволяя спектральной норме увеличиваться сверх заданного предела. После каждого шага оптимизации применяется операция ограничивающая сингулярные значения матриц весов. Благодаря этому нормы активаций остаются очень малыми, совместимыми с низкой точностью fp8 в трансформерах масштаба GPT-2.

Результаты экспериментов без стабилизирующих приёмов

Во всех экспериментах не использовались нормализация слоёв, нормализация QK или ограничение логитов. Максимальные значения активаций в трансформерах масштаба GPT-2 не превышали ~100, в то время как в базовом варианте без ограничений они превышали 148000.

| Модель | Макс. активация | Стабилизирующие приёмы | Точность валидации | Липшицево ограничение | |-----------------------|-----------------|-----------------------|--------------------|-----------------------| | Базовый (Speedrun) | 148480 | Да | 39.4% | ∞ | | Липшицев трансформер | 160 | Нет | 39.5% | 10¹⁰²⁶⁴ |

Методы введения липшицевых ограничений

Исследованы различные методы ограничения норм веса:

Weight Decay: стандартный, но не всегда строгий по спектральной норме.
Spectral Normalization: ограничивает максимальное сингулярное значение, влияет на все сингулярные значения.
Spectral Soft Cap: новый метод, плавно ограничивающий все сингулярные значения, оптимизированный для Muon.
Spectral Hammer: ограничивает только максимальное сингулярное значение, подходит для оптимизатора AdamW.

Производительность и компромиссы

Малые модели (например, Shakespeare) достигают высокой точности с доказуемыми липшицевыми границами, превосходя базовые модели.
В больших моделях (NanoGPT, 145 млн параметров) строгие липшицевы ограничения снижают выразительность, требуя больших верхних границ для сопоставимой точности.
Сочетание Muon и спектрального ограничения лидирует в соотношении точность-ограничение.

Преимущества по стабильности и защите

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

Ограничения и перспективы

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

Влияние и ресурсы

Спектральная регуляция весов с оптимизатором Muon открывает перспективы стабильного и надёжного обучения трансформеров без традиционных нормализаций. Это способствует безопасности, эффективности и приватности ИИ. Подробнее доступно в статье, на GitHub и Hugging Face страницах.