МАРС: Развязывание силы снижения дисперсии для обучения больших моделей

Обучение глубоких нейронных сетей, а в последнее время и крупных моделей, требует эффективных и масштабируемых оптимизаторов. Адаптивные алгоритмы градиентного спуска, такие как Adam, AdamW и их вариации, играют ключевую роль в этой задаче. Несмотря на разработку множества алгоритмов снижения дисперсии за последнее десятилетие, направленных на ускорение стохастической оптимизации как в выпуклых, так и в невыпуклых условиях, снижение дисперсии не получило широкого распространения при обучении глубоких нейронных сетей или крупных языковых моделей. В результате этот подход остается менее популярным в современном ИИ. В данной статье, чтобы раскрыть потенциал снижения дисперсии для эффективного обучения крупных моделей, мы предлагаем единую оптимизационную структуру MARS (Make vAriance Reduction Shine), которая объединяет методы предобусловленного градиента и снижение дисперсии с помощью техники масштабированного стохастического рекурсивного импульса. В рамках нашего подхода мы представляем три варианта MARS, которые используют обновления градиентов, предобусловленные на основе AdamW, Lion и Shampoo соответственно. Мы также проводим связь между нашими алгоритмами и существующими оптимизаторами. Экспериментальные результаты по обучению моделей GPT-2 показывают, что MARS значительно превосходит AdamW.