Предварительно обусловленный неточный стохастический ADMM для глубоких моделей

Более умная подготовка для более умного ИИ

Современные системы искусственного интеллекта — от чат-ботов до генераторов изображений — основаны на массивных нейросетях, обучение которых известно своей сложностью и высокой стоимостью. По мере того как компании и исследователи распределяют данные по множеству устройств и серверов, стандартные методы обучения часто замедляются, становятся нестабильными или просто не справляются с непредсказуемостью реальных данных. В этой статье представлена новая семья алгоритмов обучения, основанная на методе под названием PISA, которая обещает более быстрое и надежное обучение для широкого круга глубоких моделей при меньшем числе математических допущений о данных.

Почему существующие методы обучения испытывают трудности

Большинство моделей глубокого обучения обучаются с помощью вариантов стохастического градиентного спуска — подхода, который многократно сдвигает параметры модели в направлении, уменьшающем ошибку. С годами появилось множество усовершенствований — таких как Adam, RMSProp и другие — которые делают эти шаги «умнее», адаптируя размеры шагов или добавляя моментум. Однако эти методы обычно предполагают, что данные для обучения аккуратно перемешаны и статистически схожи на разных машинах, а также что некоторые математические величины остаются ограниченными. На практике, особенно в сценариях вроде федеративного обучения, где телефоны или устройства периферии содержат сильно различающиеся данные, эти допущения часто нарушаются, что приводит к медленной сходимости или плохой работе.

Новый способ согласовывать многих обучающихся

Авторы опираются на другой оптимизационный фреймворк, известный как метод переменных направлений множителей (ADMM), который хорош тем, что разбивает большую задачу на множество меньших, решаемых параллельно. Их основное достижение — PISA (preconditioned inexact stochastic ADMM) — сохраняет сильные стороны ADMM, избегая его обычных недостатков, таких как необходимость вычислять полные градиенты по всем данным или выполнять дорогие обращения матриц. Вместо этого PISA позволяет каждому клиенту или рабочему узлу обновлять свою копию модели, используя только мини-батч данных, а затем координировать эти обновления через центральную переменную. Тщательно разработанные матрицы «предварительного обусловливания» (preconditioning) изменяют направления обновлений так, чтобы обучение шло более плавно и эффективно.

Более строгие гарантии при более мягких предположениях

Отличительная черта PISA — её теоретическое обоснование. Авторы доказывают, что их алгоритм сходится при одном, относительно мягком предположении: градиент функции потерь является липшицевым в пределах ограниченной области — условие, выполняемое для многих стандартных функций потерь нейронных сетей. В отличие от большинства стохастических методов, PISA не требует, чтобы градиенты были несмещёнными, имели ограниченную дисперсию или поступали из идеально перемешанных данных. Несмотря на эту ослабленную постановку, метод достигает линейной скорости сходимости по тому, как быстро стабилизируются значения функции и обновления параметров, что ставит его в ряд лучших алгоритмов по сравнению с таблицей результатов. Это делает PISA особенно привлекательной для гетерогенных, неравномерных распределений данных, характерных для реальных условий эксплуатации.

Практические варианты для реальных глубоких сетей

Чтобы сделать фреймворк применимым к большим нейросетям, авторы предлагают два эффективных варианта: SISA и NSISA. SISA использует информацию о втором моменте — по существу отслеживая, насколько большими были прежние обновления в каждом направлении параметров — чтобы формировать простые диагональные предобуславливатели, подобные идеям Adam и RMSProp, но встроенные в структуру ADMM. NSISA идёт дальше, включая технику, известную как ортогонализация Ньютона–Шульца, вдохновлённую оптимизатором Muon, чтобы лучше согласовать моментум с полезными направлениями в пространстве параметров. Оба варианта сохраняют гарантии сходимости PISA, оставаясь при этом достаточно лёгкими по вычислительной нагрузке для современных GPU и крупных моделей.

Производительность в задачах зрения, языка и генерации

Авторы тестируют SISA и NSISA на широком наборе задач глубокого обучения. В экспериментах федеративного обучения с намеренно скошенными распределениями меток — жёсткая постановка, где каждый клиент видит только подмножество классов — SISA значительно превосходит популярные методы вроде FedAvg, FedProx, FedNova и Scaffold, достигая значительно более высокой тестовой точности на бенчмарках, таких как MNIST и CIFAR-10. Для стандартной классификации изображений с моделями вроде ResNet и DenseNet на CIFAR-10 и ImageNet SISA сопоставима или превосходит сильные оптимизаторы, включая SGD с моментумом, AdaBelief и AdamW. При дообучении языковых моделей GPT2 разных размеров NSISA показывает более низкую валидационную потерю за меньшее реальное время по сравнению со специализированными оптимизаторами, такими как Shampoo, SOAP, Adam-mini и Muon, и преимущество становится заметнее на самых больших моделях. Он также стабилизирует обучение генеративных состязательных сетей, достигая более низких значений метрики Fréchet inception distance, которая измеряет визуальное качество и разнообразие сгенерированных изображений.

Что это значит для повседневного ИИ

Проще говоря, эта работа показывает, что можно обучать мощные модели ИИ быстрее и надёжнее, даже когда данные шумные, несбалансированные или распределены по множеству устройств. Переосмыслив базовый процесс оптимизации, а не просто подбирая темп обучения, PISA и её варианты предлагают единый инструмент, который хорошо работает для задач зрения, языка, обучения с подкреплением и генеративных задач. Для конечных пользователей это может означать более умную персонализацию на телефонах, более способные языковые и визуальные модели и более эффективное использование вычислительных ресурсов в крупных дата-центрах — всё это становится возможным благодаря алгоритму обучения, лучше соответствующему реалиям современных систем ИИ.

Цитирование: Zhou, S., Wang, O., Luo, Z. et al. Preconditioned inexact stochastic ADMM for deep models. Nat Mach Intell 8, 234–245 (2026). https://doi.org/10.1038/s42256-026-01182-3

Ключевые слова: оптимизация глубокого обучения, федеративное обучение, стохастический ADMM, большие языковые модели, гетерогенные данные