Гессиан
Вспоминаем теорию
Гессиан играет большую роль в численной оптимизации. Недавно появились новые алгоритмы оптимизации (например, Sophia). Они используют Гессиан или его приближения для достижения лучших локальных минимумов и более быстрой сходимости. В этой статье мы как раз и познакомимся с Гессианом — ценным инструментом для нахождения минимума функции.
Для начала рассмотрим общую задачу оптимизации:
- \( \theta \) представляет собой переменные (это обычно параметры нейронной сети), которые мы стремимся оптимизировать;
- \( y \) соответствует истинным лейблам;
- \( \hat y \) обозначает предсказания, сделанные моделью с параметрами \( \theta \).
Для наглядности мы можем представить практический сценарий, когда мы решаем задачу регрессии для прогнозирования цен на дома и используем среднеквадратичную ошибку (MSE) в качестве нашей функции потерь.
В численной оптимизации основополагающим принципом считается Условие Оптимальности Первого Порядка: если \( \theta^\star \) является локальным минимумом \( \mathcal L \), то
\( \nabla \mathcal L(\theta^\star, \cdot)= 0. \)
Простыми словами, когда мы найдем решение, градиент нашей функции потерь будет равен нулю. Если \( \theta \) — скаляр, то производная функции в точке решения — нуль (привет, матан). Точка, которая удовлетворяет условию \( \nabla_\theta \mathcal L(\cdot)= 0 \), называется стационарной точкой \( \mathcal L \). Она может быть и локальным минимумом, и локальным максимумом, и седловой точкой, как показано на рисунке 1.
Поскольку нас интересуют только локальные минимумы, нам понадобится еще одно условие для гарантии того, что наше решение \( \theta^\star \) не является максимумом. Это условие называется Необходимое Условие Оптимальности Второго Порядка: если \( \theta^\star \) является локальным минимумом \( \mathcal L \), то
\( \nabla^2_\theta \mathcal L(\theta^\star,\cdot) \succeq 0. \)
Здесь \( \nabla^2_\theta \mathcal L(\theta^\star,\cdot) \)представляет собой вторую производную \( \mathcal{L} \) по отношению к \( \theta \) и называется Гессианом.
Важно отметить: к сожалению, седловые точки также удовлетворяют необходимому условию оптимальности второго порядка, как показано на рисунке 1. Для гарантии нашего решения как локального минимума нам нужно еще более сильное условие, известное как Достаточное Условие Оптимальности Второго Порядка: \( \nabla^2_\theta \mathcal L(\theta^\star,\cdot) \succ 0 \).
Теперь, когда мы рассмотрели основные теоретические аспекты, давайте перейдем к алгоритмам, которые используют информацию второго порядка для оптимизации наших параметров \( \theta \).
Метод Ньютона с точным вычислением
Метод Ньютона (метод Ньютона-Рафсона) — метод, в первую очередь используемый для решения задач нахождения корней нелинейной функции. Фактически, Условие Оптимальности Первого Порядка может быть сформулировано как задача нахождения корней: нам нужно найти параметры \( \theta^\star \), которые делают градиент функции потерь равным нулю. Для достижения этой цели мы можем применить метод Ньютона с разложением \( \nabla \mathcal{L} \) в ряд Тейлора:
\( \nabla \mathcal L(\theta_k) + \frac{\partial }{\partial \theta} (\nabla \mathcal L(\theta_k))(\theta — \theta_k) = 0, \\ \theta_{k+1} = \theta_k — \nabla^2 \mathcal L(\theta_k)^{-1} \nabla \mathcal L(\theta_k). \)
Если мы начинаем с функции потерь \( \mathcal{L} \), мы разлагаем ее в ряд Тейлора и учитываем члены до второго порядка. Иными словами, метод Ньютона локально аппроксимирует функцию потерь квадратичной функцией, вычисляет минимум этой квадратичной аппроксимации и итеративно уточняет его до сходимости, как показано на анимации ниже:
Методы типа метода Ньютона
Мы увидели, что метод Ньютона полагается на Гессиан для обновления параметров. Однако вычисление Гессиана может быть затратным. Поэтому ученые разработали различные методы для приближенного вычисления Гессиана. Обычно они следуют общей структуре:
\( \theta_{k+1} = \theta_k — B_k^{-1}\nabla \mathcal L(\theta_k) \)
где \( B_k \) — приближение Гессиана.
Есть несколько популярных методов типа метода Ньютона:
- Метод Левенберга-Марквардта (ЛМ)
- Градиентный метод (ГМ)
- Квазиньютоновский метод
Давайте рассмотрим их по порядку.
Метод Левенберга-Марквардта (ЛМ)
Метод Левенберга-Марквардта — широко используемый метод для решения нелинейной задачи наименьших квадратов. Он часто применяется в задачах одновременной локализации и построения карты (SLAM) на основе камер и LiDAR. Основная идея метода ЛМ заключается в локальном приближении функции потерь квадратичной моделью.
Рассмотрим такую функцию потерь наименьших квадратов:
\( \mathcal{L} = \frac{1}{2}||l( \theta, y, \hat y)||_2^2 + \frac{1}{2} \lambda ||\theta||_2^2. \)
Далее разложим \( l(\theta, y, \hat y) \) в ряды Тейлора и оставим только первые два члена:
\( l \approx
l(\theta_k) + \nabla l(\theta_k)^T (\theta — \theta_k) + \mathrm{h.o.t.} \)
Затем подставим это обратно в функцию потерь и получим:
\( \begin{align} \mathcal{L} &= \frac{1}{2}\big(l — \nabla l (\theta — \theta_k)\big)^T\big(l — \nabla l (\theta — \theta_k)\big) + \frac{1}{2}\lambda\theta^T \mathbb{I} \theta \\ &=\frac{1}{2}\big(l^Tl + 2l^T \nabla l (\theta — \theta_k) + (\theta — \theta_k)^T \nabla l^T \nabla l (\theta — \theta_k)\big) + \frac{1}{2}\lambda\theta^T \mathbb{I} \theta \\ & = \frac{1}{2} l^T l + l^T \nabla l(\theta — \theta_k) + \frac{1}{2}(\theta — \theta_k)\underbrace{(\nabla l^T \nabla l + \lambda \mathbb{I})}_{B} (\theta — \theta_k) \end{align} \)
Из уравнения (3) видно, что приближение Гессиана методом ЛМ задается формулой \( B_k = \nabla l^T \nabla l + \lambda \mathbb{I} \). Регуляризационный член \( \lambda \mathbb{I} \) гарантирует, что \( B \) всегда обратима, даже когда \( \nabla l^T \nabla l \) является вырожденной. Следовательно, одна итерация метода ЛМ выражается как:
\( \theta_{k+1} = \theta_k — \left(\nabla l^T \nabla l + \lambda \mathbb{I}\right)^{-1} \nabla \mathcal L(\theta_k) \)
Градиентный метод (ГМ)
Градиентный метод, или метод наискорейшего спуска — один из наиболее широко используемых методов оптимизации в машинном обучении. Он приближает матрицу Гессе матрицей, состоящей из единиц и взвешенной обратным значением скорости обучения \( \alpha: B = \frac{1}{\alpha}\mathbb{I} \). Следовательно, шаги оптимизации будут выглядеть таким образом:
\( \theta_{k+1} = \theta_k — \alpha \nabla \mathcal L(\theta_k) \)
Квазиньютоновский метод
Квазиньютоновский методы приближают матрицу Гессе \( B_{k+1} \) на основе информации о \( B_k \), \( \nabla \mathcal L(\theta_k) \) и \( \nabla \mathcal L(\theta_{k+1}) \). Среди различных предложенных методов для определения матрицы Гессе \( B_{k+1} \) одним из широко используемых и успешных является алгоритм Бройдена-Флетчера-Гольдфарба-Шанно (BFGS), где кривизна функции приближается с применением конечных разностей. Проще говоря, мы выбираем матрицу \( B_{k+1} \) так, чтобы удовлетворить условию секущей:
\( B_{k+1} \left(\theta_{k+1} — \theta_k \right) = \nabla \mathcal L(\theta_{k+1}) — \nabla \mathcal L(\theta_{k}). \)
Логика выбора этого условия: матрица Гессе \( \nabla^2 \mathcal L(\theta_k) \) удовлетворяет ему при приближении \( \theta_{k+1} \) к \( \theta_{k} \). Поскольку несколько матриц могут удовлетворять условию секущей, мы предполагаем, что \( B_{k+1} \) и \( B_k \) должны быть близки друг другу. Поэтому мы выбираем \( B_{k+1} \) как матрицу, наиболее близкую предыдущему приближению Гессиана \( B_k \) из всех матриц, удовлетворяющих условию секущей. Для измерения близости между \( B_k \) и \( B_{k+1} \) мы можем использовать концепцию дифференциальной энтропии между случайными переменными с нулевыми средним значением и гауссовскими распределениями \( \mathcal{N}(0, B_k) \) и \( \mathcal{N}(0, Z) \). Такую концепцию мы формулируем как полуопределенную программу (задачу оптимизации). Она имеет аналитическое решение в виде:
\( \begin{align} B_{k+1} &= B_k — \frac{B_k s s^T B_k}{s^T B_k s} + \frac{y y^T}{s^T y} \\ s &= \theta_{k+1} — \theta_k \\ y &=\nabla \mathcal L(\theta_{k+1}) — \nabla \mathcal L(\theta_{k}) \end{align} \)
Напомним, что Гессиан часто является плотной матрицей и имеет размеры \( n_\theta \times n_\theta \), где \( n_\theta \) представляет собой количество параметров для оптимизации. Следовательно, хранение этой матрицы в памяти может потреблять много ресурсов. Для решения этой проблемы была разработана вариация метода BFGS, известная как BFGS с ограниченной памятью (L-BFGS).
Как L-BFGS преодолевает проблему памяти? Она сохраняет историю предыдущих \( m \) обновлений как параметров \( \theta \), так и градиента \( \nabla \mathcal L(\theta) \). Этот подход позволяет L-BFGS неявно представлять приближение Гессиана, используя только несколько векторов.
Масштабируемый метод второго порядка для глубокого обучения
Sophia (Second-order Clipped Stochastic Optimization) идет еще дальше в снижении вычислительных затрат, приближая Гессиан с помощью диагональной матрицы. Кроме того, она включает механизм обрезки для контроля максимального размера обновления параметров. Эта функция важна, потому что быстрые изменения в ландшафте функции потерь могут сделать информацию второго порядка ненадежной.
Если говорить точнее, Sophia оценивает диагональные элементы Гессиана функции потерь, используя mini-batch входных данных каждые \( k \) шагов. На каждом шаге Sophia обновляет параметры с применением экспоненциального скользящего среднего (EMA) градиента, разделенного на EMA оценки диагонального гессиана. Он затем обрезается скаляром. Эксперименты по обучению LLM на датасете GPT-2 показывают: Sophia в два раза быстрее AdamW!
Пример
Наш пример взят из исследования оптимизаторов, которое использовало четырехслойную сеть из 20 нейронов в каждом слое + все нейроны с функцией активации гиперболического тангенса (tanh). Эта сеть обучалась на данных \( \mathcal{D} = \{x_i, \sin(x_i) \}_{i=1}^{100} \), где \( x_i \) выбирались из равномерного распределения \( \mathcal{U}(0,1) \). Во время обучения в качестве функции потерь использовалась среднеквадратичная ошибка (MSE). Давайте посмотрим на результаты тренировки:
Результаты этого простого примера и выводы статьи о методе Sophia показывают нам, что методы оптимизации второго порядка (вроде метода Ньютона) существенно быстрее. Плюс они дают более низкие значения функции потерь. Это свидетельствует о потенциальной пользе методов второго порядка для обучения нейронных сетей.
