Ускоренные методы оптимизации и их роль в обучении больших моделей

Ускорение

Опубликованный в 1961 году доклад Гельфанда и Цетлина — первая работа, которая мотивировала цикл исследований по ускоренным методам в оптимизации. В ней использовали овражный метод, с помощью которого авторы распознавали профиль оврага и получали ускорение.

Обычный градиентный метод — движение по нормали к множеству уровня. Если множество уровней «сосискообразное» — имеем вытянутую кривую, отвечающую f(x) = const., то есть множество, где целевая функция f(x) принимает одинаковые значения.

Если это множество сильно вытянутое, градиентный спуск будет принимать форму пилы, которая очень медленно сходится.

Идея: растянуть график.

Проблема: сложно определить, где в многомерном пространстве и особенно в пространстве больших размерностей находится центр «сосиски».

В 1964 году Б. Т. Поляк предложил добавлять в процедуру градиентного спуска моментное ускорение — использовать momentum. Ю. Е. Нестеров доказал, что при правильном добавлении моментума можно получить глобальную сходимость.

Использованная литература:

• Polyak B. T. Some methods of speeding up the convergence of iteration methods // Comput. Math. Math. Phys. — 1964. — V. 4:5. — P. 1–17.
• Nemirovski A. Orth-method for smooth convex optimization // Cybern. Soviet J. Comput. Syst. Sci. — 1982. — V. 2. — P. 937–947.
• Nesterov Y. E. A method for solving the convex programming problem with convergence rate O (1/k^2) // Dokl. Akad. nauk USSR. — 1983. — V. 269. — P. 543–547.
• Opt.mipt.ru

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

Самое важное — правильное добавление моментума. Под правильным имеется в виду, чтобы это действительно работало.

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

Также требуется гладкость целевой функции. Часть задач негладкие — например, в нейронной сети ReLU гладкость теряется из-за пороговой функции ReLU. То есть существуют нюансы, которые усложняют жизнь. Но моментум — важный тезис, который до сих пор работает. Бывает выгоднее подобрать параметры для выпуклого случая и потом пытаться подкрутить их в невыпуклом общем случае, чем специально затачиваться под невыпуклую задачу.

Где применяется моментум. Моментум включают в себя современные пакеты PyTorch и процедуры, связанные с Adam. Эти технологии присутствуют во всех сферах жизни для решения самых современных задач — например, Adam используется в GigaChat, аналоге ChatGPT от Сбера.

На данный момент не существует полностью адаптивного метода с ускорением в стохастической постановке. В конце 2023 года были предложены методы без ускорения. Исследования всё ещё ведутся, поскольку Adam — не полностью адаптивный метод, так как он не работает сам и в него входит некая константа.

Неточность в градиенте и ранняя остановка

Когда вы обучаете нейронную сеть или решаете задачу обучения, у вас есть два варианта:

1. Решать задачу оптимизации, например, методом Ньютона, LBFGS или любым другим. Если делать это очень точно и без регуляризации, есть риск переобучить модель.
2. Решать не исходную задачу, а задачу стохастической оптимизации. Ранняя остановка таких процедур — один из способов борьбы с переобучением без регуляризации.

Возникает вопрос: насколько неточный или пробатченный градиент испортит скорость сходимости. До какого-то момента она практически никак не ухудшается, как если бы шума не было. Затем достигается определённый порог, масштаб которого пропорционален шуму, а коэффициент пропорциональности характеризуется размером решения Rδ. Как только вы достигнете масштаба Rδ, надо остановиться — вы уже не решите задачу лучше, потому что оракул неточный. Если вы не остановитесь, то процедура может продолжаться, но уже не будет гарантий, что она уменьшится. Можно пойти в разнос — Поляк доказал это 40 лет назад. Многие полезные наблюдения были сделаны существенно раньше в обычной оптимизации, которая сейчас используется в сферах обычного и глубокого обучения.

Использованная литература:

• B. Poljak. Iterative algorithms for singular minimization problems, in Nonlinear Programming 4, Elsevier, 1981. Р. 147–166.
• Devolder O. Exactness, inexactness and stochasticity in first-order methods for large-scale convex optimization. — PhD thesis, 2013.
• Vasin A. et al. Accelerated gradient methods with absolute and relative noise in the gradient // Optimization Methods and Software. — 2023. — P. 1–50.

Сейчас пользуются популярностью сжатия и компрессии — процедуры, которые по ходу решения задач обучения требуют передачи градиентов от узла к серверу. Размер обучаемой модели — десятки миллиардов параметров, поэтому пересылать вектор масштаба размером 10 Гб — значит сжимать его и передавать неточно.

Масштаб неточности, который вы получаете путём сжатия, пропорционален самому вектору: чем больше вектор, который вы сжимаете, тем больше масштаб относительной неточности.

Градиентные методы. Неускоренные процедуры, то есть обычный градиентный метод, устойчивы.

Использованная литература:

• Поляк Б. Т. Введение в оптимизацию. — М.: Наука, 1983.
• De Klerk E., Glineur F., Taylor A. B. On the worst-case complexity of the gradient method with exact line search for smooth strongly convex functions // Optimization Letters. — 2017. — V. 11. — P. 1185–1199.
• Vasin A. et al. Accelerated gradient methods with absolute and relative noise in the gradient // Optimization Methods and Software. — 2023. — P. 1–50.

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

Ускоренные градиентные методы. Картина принципиально меняется, когда мы берём ускоренный метод.

Использованная литература:

• Vasin A. et al. Accelerated gradient methods with absolute and relative noise in the gradient // Optimization Methods and Software. — 2023. — P. 1–50.
• Kornilov N. et al. Intermediate Gradient Methods with Relative Inexactness // arXiv preprint arXiv:2310.00506. — 2023.

На данный момент доказано, что ускорение гарантированно будет правильным, если масштаб неточности относительно небольшой и равен числу обусловленности задач (μ/L)1/2, где μ — константа сильной выпуклости, а L — константа гладкости. Если мы не можем это гарантировать, то метод может расходиться. Этому посвящена работа Александра и его коллег. Они опирались на теорию, которая сводит поиск оценок к решению задач выпуклой оптимизации. Данный результат доказал компьютер, а 10 лет назад об этом можно было только мечтать.

Седловые задачи

Такие задачи используются, например, для моделирования атак.

Использованная литература:

• Кибардин, Владимир Михайлович. Адаптивные методы оптимизации сложных систем при аддитивных критериях: диссертация ... кандидата технических наук: 05.13.01. — Москва, 1979. — 166 с. : ил. науч. рук. Б. Т. Поляк.
• Gorbunov E. et al. Recent theoretical advances in decentralized distributed convex optimization // High-Dimensional Optimization and Probability: With a View Towards Data Science. — Cham : Springer International Publishing, 2022. — P. 253–325.

В этих постановках нет одинаковости x и y. Х — переменные, по которым берётся min, то есть, например, настоящая задача. Y — переменные, по которым берётся max, то есть то, что обеспечивает робастность найденного решения.

Отсюда появляется идея расщеплять оракульные сложности, то есть пытаться решать задачу так, чтобы отдельно считать сложность по x, по у или по разным блокам. Умение считать по блокам — очень важная опция, на которую направлены современные методы ускоренной оптимизации, которые уже решают седловые задачи.

Использованная литература:

• Borodich E. et al. Optimal Algorithm with Complexity Separation for Strongly Convex-Strongly Concave Composite Saddle Point Problems // arXiv:2307.12946.
• Kovalev D., Gasnikov A., Richtárik P. Accelerated primal-dual gradient method for smooth and convex-concave saddle-point problems with bilinear coupling // Advances in Neural Information Processing Systems. — 2022. — V. 35.
• Kovalev D., Gasnikov A. The first optimal algorithm for smooth and strongly-convex-strongly-concave minimax optimization // Advances in Neural Information Processing Systems. — 2022. — V. 35.

Седловые задачи могут возникать естественным образом. Это происходит, когда задача машинного обучения распределена на разных устройствах и функция fk хранится на k-том устройстве. Мы хотим организовать вычисление таким образом, чтобы k-тое устройство считало только градиент своей функции, а градиенты других функций пересылались бы ей.

Матрица Лапласа обладает особенностью: в случае связанной коммуникационной сети Wx = 0 тогда и только тогда, когда x 1 = х 2 и т. д., то есть когда её компоненты вектора одинаковые.

Поэтому, расщепляя в постановке задачи хk , мы можем равносильно переписать, что все х равны между собой.

Теперь делаем ключевой шаг — заменяем условие, что все х равны между собой, на Wx = 0, потому что это равносильно. W была введена, потому что её нулевые элементы соответствуют отсутствию связи между узлами. И поэтому умножение матрицы на вектор интерпретируется как коммуникация, то есть коммуницируют только те, кто связан. Таким образом в постановку задачи зашили информацию о том, кто может коммуницировать, а кто — нет.

Седловую задачу можно решать с учётом расщепления оракульных сложностей. Нюанс: ускорение не переносится на меняющиеся сети.

Использованная литература:

• Metelev D. et al. Is Consensus Acceleration Possible in Decentralized Optimization over Slowly Time-Varying Networks? // ICML, 2023.
• Metelev D. et al. Decentralized optimization over slowly time-varying graphs: Algorithms and lower bounds //Computational Management Science. — 2024. — V. 21. — № 1. — P. 1–25.

Минибатчирование

Переходим к релаксации исходной задачи, которая связывает оптимизацию с обучением в том числе нейронных сетей.

Нам неведома настоящая задача, которую мы решаем. Это задача стохастической оптимизации по ξ, где ξ, например, картинки и одновременно их разметка. Кошка и информация о том, что это кошка.

Считаем, что изображения генерируются из какого-то распределения. Берётся функция потерь от предсказаний нейронной сети и реальности. Мы стараемся минимизировать невязку между прогнозом нейросети и тем, что происходит на самом деле. Берём матожидания от невязки. Но проблема в том, что этот функционал невычислим, у нас нет закона распределения ξ.

Есть два решения этой проблемы:

1. Семплим из ξ , храним датасеты на разных устройствах и как-то решаем задачи распределённой или нераспределённой оптимизации, сформированные из этой выборки.
2. Онлайн получаем каждый раз новый элемент выборки и дообучаемся. Приходит поток данных, и каждый новый элемент выборки позволяет скорректировать веса нейронной сети. Это происходит с помощью батчирования.

Использованная литература:

• Posobie.pdf (mipt.ru)

Нам известно, что градиентный спуск неплохо решает задачу оптимизации, поэтому заменяем градиент на пробатченный градиент, то есть среднее арифметическое стохастических градиентов. Чем больше объём выборки, то есть размер батча r, тем точнее уравнение аппроксимирует градиент целевой функции. Если мы надеемся, что это действительно неплохо описывает градиент, то тогда и сама процедура должна каким-то образом аппроксимировать процедуру градиентного спуска. Но градиентный спуск сходится, и из этого рождается идея заменить в градиентном спуске градиент на пробатченный градиент. Эта идея тиражируется на моментные методы, в которых также можно заменить градиент на пробатченный и после этого выбирать размер батча.

Полезное наблюдение, работающее необязательно в идеальных выпуклых сетапах: по теории в выпуклом случае размер батча и размер шага должны быть пропорциональны. Если вы берёте большой батч, то до определённого момента можно брать и больший шаг, если изначально стартовали с маленького шага и маленького батча. Если вы берёте очень большой батч, то очень большой шаг брать нельзя, потому что даже в детерминированных процедурах размер шага лимитирован.

Как правильно выбрать размер батча. Это важно, потому что перебатчить — плохо, нужны лишние расчёты.

На рисунке выше представлены формулы, которые показывают, как сходятся процедуры такого типа. И в градиентном методе, и в ускоренном градиентном методе формулы состоят из двух слагаемых и отличаются только детерминированным слагаемым: в первом случае сходимость будет 1/T в детерминированной части, а во втором — 1/T2, то есть быстрее.

Второе слагаемое отвечает стохастической сути задачи и пропорционально (1 × r) ÷ Т, где r — размер батча. То есть за счёт выбора большого размера батча мы можем нивелировать второе слагаемое и приравнять его к первому, должным образом выбирая r. Подбираем размер батча так, чтобы второе слагаемое стало таким же маленьким, как и первое. Больше не надо.

В случае ускоренного метода за счёт большего размера батча вы можете сократить последовательное число итерации. В случае ускоренного метода оно будет равно 1 ÷ (ε)½, а в случае неускоренного метода — 1 ÷ ε, где ε — желаемая точность.

Федеративное обучение

Использованная литература:

• Woodworth B. et al. The Min-Max Complexity of Distributed Stochastic Convex Optimization with Intermittent Communication (Extended Abstract)

Когда мы говорим о батче, можно подумать, что после каждой итерации заменяем градиент на пробатченный и используем его. В районе 2003 года Нестеров предлагал делать последовательную процедуру. Вы параллельно запускаете стохастический градиентный спуск и так параллельно запускаете несколько веток. Через какое-то количество локальных шагов устройства коммуницируют и вычисляется общий батч, то есть осуществляется синхронизация, потом они снова идут локально. В момент синхронизации они обмениваются информацией и стартуют из одинаковой точки, потом расходятся, а затем снова синхронизируются.

В 2016 году постдок Нестерова Питер Рихтарик (Peter Richtarik) с коллегами написал статью The Min-Max Complexity of Distributed Stochastic Convex Optimization with Intermittent Communication. Суть статьи: мобильным устройствам невыгодно каждый раз посылать куда-то свою информацию. Надо сделать так, чтобы они сами у себя что-то обучали, а потом через какое-то количество тактов посылали куда-то, и за счёт этого можно федеративно организовать обучение.

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

1. Батчировать каждую итерацию, то есть коммуникация будет происходить каждую итерацию.
2. Не коммуницировать и просто последовательно идти.

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

Перепараметризация

Переобученные нейронные сети появляются в ситуации, когда не хватает обучающей выборки или много параметров.

Использованная литература:

• Woodworth B. E., Srebro N. An even more optimal stochastic optimization algorithm: minibatching and interpolation learning // Advances in Neural Information Processing Systems. — 2021. — V. 34. — P. 7 333–7 345.

Это соответствует тому, что в функционале, который вы обучаете, удаётся выбором параметров обнулить все слагаемые. То есть подбирайте веса нейронной сети так, чтобы конкретный выход полностью и точно описывался, потому что число подкручиваемых параметров меньше, чем число уравнений, которые надо выполнить. Этот случай означает, что у нас возникает понятие дисперсии в решении, которая будет равна нулю именно для перепараметризованных ситуаций. Перепараметризованная оптимизация — случай, когда дисперсия в решении либо 0, либо маленькая.

В этом случае мы можем точнее переписать существующие оценки, то есть заменить σ на σ✱ . σ — глобальная константа, а σ✱ — локальная. Для перепараметризованных моделей это что-то близкое к нулю. Мы можем убрать последнее слагаемое. Это означает, что методы с клиппированием, которые это дают, крутые и выгодны в таких режимах.

Условие сильного роста

Использованная литература:

• Vaswani S., Bach F., Schmidt M. Fast and faster convergence of sgd for over-parameterized models and an accelerated perceptron // The 22nd international conference on artificial intelligence and statistics. — PMLR, 2019. — P. 1 195–1 204.

Стохастический градиент бывает ограниченным не только σ, но и поправкой. Такое обобщение классических результатов крайне продуктивно для современных приложений ускоренных методов к ситуациям, когда нет полной информации, например, разных компрессий. Это позволяет сэкономить на передаче данных.

Клиппирование

Использованная литература:

• Puchkin N. et al. Breaking the Heavy-Tailed Noise Barrier in Stochastic Optimization Problems // arXiv:2311.04161
• Gorbunov E. et al. High-Probability Convergence for Composite and Distributed Stochastic Minimization and Variational Inequalities with Heavy-Tailed Noise // ICML, 2023

Спикер и его коллеги опубликовали статью, идея которой в том, что можно делать клиппирование и заменять стохастический градиент его медианой. Медиана — хорошая оценка. Медиана и клиппирование по отдельности — просто хорошо, а если их объединить, это прямо круто.

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