Обучение крупномасштабных трансформеров является сложной задачей в области глубокого обучения, особенно когда модели становятся больше и выразительнее. Исследователи из MIT решают эту проблему на корню: нестабильный рост активаций и всплески потерь, вызванные неограниченными нормами весов и активаций. Их решение — обеспечить доказуемые границы Липшица для трансформера путём спектральной регуляции весов — без использования нормализации активаций, нормы QK или методов ограничения логаритов.
Что такое граница Липшица и зачем её устанавливать?
Граница Липшица для нейронной сети определяет максимальное изменение выхода в ответ на возмущения входа (или веса). Математически функция $f$ является $K$-Липшицевой, если:
$∥f(x1)−f(x2)∥≤K∥x1−x2∥ ∀x1,x2$
Меньшая граница Липшица ⇒ большая устойчивость и предсказуемость. Это важно для стабильности, устойчивости к враждебным воздействиям, конфиденциальности и обобщения. Более низкие границы означают, что сеть менее чувствительна к изменениям или враждебному шуму.
Мотивация и постановка задачи
Традиционно для обучения стабильных трансформеров в масштабе использовались различные методы стабилизации:
* нормализация слоёв;
* нормализация QK;
* ограничение логаритов tanh.
Но эти методы не решают основную проблему роста спектральной нормы (наибольшего сингулярного значения) в весах, что является причиной взрывных активаций и нестабильности обучения, особенно в больших моделях.
Центральная гипотеза: если мы будем регулировать веса спектрально — помимо оптимизатора или активаций — мы сможем поддерживать жёсткий контроль над липшицевостью, потенциально решая проблему нестабильности в её источнике.
Ключевые инновации
Спектральная регуляция весов и мюонный оптимизатор
Мюонный оптимизатор спектрально регулирует градиенты, гарантируя, что каждый шаг градиента не увеличивает спектральную норму сверх установленного предела.
Исследователи распространяют регуляцию на веса: после каждого шага они применяют операции для ограничения сингулярных значений каждой весовой матрицы. В результате нормы активаций остаются на удивление маленькими — редко превышают значения, совместимые с точностью fp8 в их трансформерах масштаба GPT-2.
Удаление трюков со стабильностью
Во всех экспериментах не использовались нормализация слоёв, нормализация QK, ограничение логаритов tanh. Тем не менее максимальные значения активаций в их трансформерах масштаба GPT-2 никогда не превышали ~100, в то время как неограниченная базовая линия превышала 148 000.
Методы обеспечения ограничений Липшица
Исследователи изучили и сравнили различные методы ограничения норм весов по их способности:
* поддерживать высокую производительность;
* гарантировать границу Липшица;
* оптимизировать компромисс между производительностью и липшицевостью.
Методы
* Распад веса: стандартный метод, но не всегда строгий в отношении спектральной нормы.
* Спектральная нормализация: гарантирует, что верхнее сингулярное значение ограничено, но может повлиять на все сингулярные значения глобально.
* Спектральное мягкое ограничение: новый метод, который плавно и эффективно применяет $σ→min(σmax,σ)$ ко всем сингулярным значениям параллельно (используя нечётные полиномиальные аппроксимации). Это разработано для стабильных обновлений высокого ранга в мюоне для получения жёстких границ.
* Спектральный молоток: устанавливает только наибольшее сингулярное значение равным $σmax$, лучше всего подходит для оптимизатора AdamW.
Экспериментальные результаты и выводы
Оценка модели в различных масштабах
* Shakespeare (маленький трансформер, <2-Липшиц): достигает 60% точности проверки с доказуемой границей Липшица ниже. Превосходит неограниченную базовую линию по потерям при проверке.
* NanoGPT (145 млн параметров): с границей Липшица <10, точность проверки: 21,2%. Чтобы соответствовать сильному неограниченному базовому уровню (39,4% точности), потребовалась большая верхняя граница 10²⁶⁴. Это подчёркивает, как строгие ограничения Липшица часто идут на компромисс с выразительностью в больших масштабах.
Эффективность метода ограничения веса
Мюон + спектральное ограничение: лидирует на границе компромисса — более низкие константы Липшица при совпадающих или лучших потерях при проверке по сравнению с AdamW + распад веса.
Спектральная мягкая нормализация (в мюоне) последовательно обеспечивает наилучший фронт на компромиссе между потерями и липшицевостью.
Стабильность и надёжность
Устойчивость к враждебным воздействиям резко возрастает при более низких границах Липшица. В экспериментах модели с ограниченной константой Липшица испытывали гораздо меньшее падение точности при враждебных атаках по сравнению с неограниченными базовыми линиями.
Величины активаций
При спектральной регуляции весов максимальные активации остаются крошечными (совместимыми с fp8), по сравнению с неограниченными базовыми линиями, даже в масштабе. Это открывает возможности для обучения и вывода с низкой точностью в аппаратных средствах, где меньшие активации снижают вычислительные, запоминающие и энергетические затраты.
Ограничения и открытые вопросы
* Выбор «самого жёсткого» компромисса для норм веса, масштабирования логаритов и масштабирования внимания всё ещё основан на пробах и ошибках, а не на принципе.
* Текущие верхние границы являются неточными: рассчитанные глобальные границы могут быть астрономически большими (например, 10²⁶⁴), в то время как реальные нормы активации остаются небольшими.
* Неясно, возможно ли сопоставление производительности с неограниченной базовой линией при строго малых границах Липшица при увеличении масштаба — необходимы дополнительные исследования.
Заключение
Спектральная регуляция весов, особенно в сочетании с мюонным оптимизатором, может стабильно обучать большие трансформеры с установленными границами Липшица без нормализации активаций или других вспомогательных методов. Это устраняет нестабильность на более глубоком уровне и поддерживает активации в компактном, предсказуемом диапазоне, значительно повышая устойчивость к враждебным воздействиям и потенциально повышая эффективность аппаратного обеспечения.
Это направление работы указывает на новые эффективные вычислительные примитивы для регулирования нейронных сетей с широким применением в области конфиденциальности, безопасности и развёртывания ИИ с низкой точностью.
1. Какие проблемы решают исследователи из MIT при обучении крупномасштабных трансформеров?
Исследователи из MIT решают проблему нестабильного роста активаций и всплесков потерь, вызванных неограниченными нормами весов и активаций при обучении крупномасштабных трансформеров.
2. Что такое граница Липшица и почему она важна для трансформеров?
Граница Липшица для нейронной сети определяет максимальное изменение выхода в ответ на возмущения входа (или веса). Меньшая граница Липшица означает большую устойчивость и предсказуемость, что важно для стабильности, устойчивости к враждебным воздействиям, конфиденциальности и обобщения.
3. Какие методы использовали исследователи для обеспечения границ Липшица у трансформеров?
Исследователи использовали спектральную регуляцию весов и мюонный оптимизатор для обеспечения границ Липшица у трансформеров. Они распространили регуляцию на веса, применяя операции для ограничения сингулярных значений каждой весовой матрицы после каждого шага.
4. Какие экспериментальные результаты были получены при использовании спектральной регуляции весов?
При использовании спектральной регуляции весов исследователи достигли высокой точности проверки с доказуемой границей Липшица ниже, чем у неограниченной базовой линии. Они также показали, что устойчивость к враждебным воздействиям резко возрастает при более низких границах Липшица.
5. Какие ограничения и открытые вопросы существуют в области регулирования трансформеров с использованием границ Липшица?
В области регулирования трансформеров с использованием границ Липшица существуют следующие ограничения и открытые вопросы:
* Выбор «самого жёсткого» компромисса для норм веса, масштабирования логаритов и масштабирования внимания всё ещё основан на пробах и ошибках, а не на принципе.
* Текущие верхние границы являются неточными: рассчитанные глобальные границы могут быть астрономически большими, в то время как реальные нормы активации остаются небольшими.
* Неясно, возможно ли сопоставление производительности с неограниченной базовой линией при строго малых границах Липшица при увеличении масштаба — необходимы дополнительные исследования.