Новый оптимизатор SignMuon сокращает трафик при распределенном обучении нейросетей в 32 раза
Разработчики из сообщества arXiv предложили SignMuon — новый метод оптимизации для распределенного обучения нейросетей. Он передаёт только знаки градиентов (1 бит) вместо полных 32-битных чисел, что радикально снижает объём сетевого трафика.
SignMuon сочетает два подхода: большинственное голосование из signSGD и матричную структуру Muon. Каждый узел вычисляет направление по формуле Muon, используя итерацию Ньютона–Шульца для полярного разложения импульса, а затем передаёт только знаки. Агрегация происходит по принципу большинства, при этом локальный полярный шаг дополнительно ортогонализирует веса без дополнительных затрат на связь.
Эксперименты показали, что SignMuon обеспечивает 32-кратное снижение пропускной способности по сравнению с float32 (4? по сравнению с int8). В тестах на CIFAR-10 с ResNet-50 (330 конфигураций) метод достиг лучшей точности валидации — 92.15%. Версия с большинственным голосованием на 4 GPU показала точность 92.02% при сокращении времени обучения на 37% при одинаковом эффективном размере батча.
На модели nanoGPT SignMuon продемонстрировал более низкую перплексию и лучшую производительность в любой момент времени по сравнению с другими знаковыми методами. Слабая масштабируемость сохраняется вплоть до 16 GPU.
С теоретической стороны авторы доказали, что при условии гладкости спектральной нормы и ограниченной дисперсии стохастических градиентов метод сходится со скоростью O(1/?T) по невыпуклой метрике стационарности L1. При симметричном шуме большинственное голосование уменьшает стохастическую ошибку в ?M раз, где M — число рабочих узлов, что аналогично signSGD.
SignMuon требует всего одну целочисленную all-reduce операцию на итерацию, вся ортогонализация выполняется локально. Это делает метод перспективным для обучения больших языковых моделей и других архитектур, где коммуникация является узким местом.
