Orth-Dion ускоряет распределенное обучение нейросетей, устраняя геометрическое искажение

Исследователи представили новый алгоритм Orth-Dion, предназначенный для оптимизации распределенного обучения нейронных сетей. Метод решает проблему геометрического несоответствия в низкоранговом спектральном сжатии градиентов, которая замедляла сходимость предыдущих подходов.

Алгоритм основан на оптимизаторе Dion, который использует низкоранговое сжатие градиентов для уменьшения объёмов передаваемых данных. Однако Dion применяет нормализацию столбцов правого фактора, что не соответствует геометрии низкорангового спектрального метода Muon. Это приводило к замедлению сходимости — скорость падала в корень из r раз, где r — ранг аппроксимации.

Orth-Dion заменяет нормализацию столбцов на QR-ортогонализацию правого фактора. Это позволяет получить точный полярный фактор ранга r, устраняя геометрическое несоответствие. В результате достигается скорость сходимости O(?(L_r/T)), где L_r — константа кривизны, что соответствует показателям полных спектральных методов.

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

Эксперименты на предобучении больших языковых моделей подтвердили предсказанное масштабирование ?r и показали, что Orth-Dion устраняет разрыв в сходимости с Muon при тех же коммуникационных затратах, что и Dion. Это делает новый метод перспективным для крупномасштабного распределённого обучения.