Hessian-Free оптимизация с помощью TensorFlow

в 17:11, , рубрики: machine learning, python, TensorFlow, Алгоритмы, математика, машинное обучение, Программирование

Добрый день! Я хочу рассказать про метод оптимизации известный под названием Hessian-Free или Truncated Newton (Усеченный Метод Ньютона) и про его реализацию с помощью библиотеки глубокого обучения — TensorFlow. Он использует преимущества методов оптимизации второго порядка и при этом нет необходимости считать матрицу вторых производных. В данной статье описан сам алгоритм HF, а так же представлена его работа для обучения сети прямого распространения на MNIST и XOR датасетах.
Hessian-Free оптимизация с помощью TensorFlow - 1


Немного о методах оптимизации

Обучение нейросети включает в себя минимизацию функции ошибки (loss function) по отношению к ее параметрам (weights), которых может быть очень много. Поэтому существуют множество методов оптимизации для решения данной проблемы.

Градиентный спуск (Gradient Descent)

Градиентный спуск — это простейший метод для последовательного нахождения минимума дифференцируемой функции $f$ (в случае нейронных сетей это функция стоимости). Имея несколько параметров $x$ (веса сети) и дифференцируя по ним функцию получаем вектор частичных производных или вектор градиента:

$nabla f(x)=langlefrac{delta f}{delta x_1}, frac{delta f}{delta x_2}, ...,frac{delta f}{delta x_n}rangle$

Градиент всегда указывает по направлению максимального роста функции. Если же мы будем двигаться в противоположную сторону (т.е. $-nabla f(x)$) то со временем придем к минимуму, что нам и требовалось. Простейший алгоритм градиентного спуска:

  1. Инициализация: случайно выбираем параметры $x_0$
  2. Вычисляем градиент: $nabla f(x_i)$
  3. Изменяем параметры в сторону отрицательного градиента: $x_{i+1}=x_i - alpha nabla f(x_i)$, где $alpha$ — некоторый параметр скорости обучения (learning rate)
  4. Повторяем предыдущие шаги пока градиент не станет достаточно близок к нулю

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

Метод Ньютона (Newton's Method)

А что если мы возьмем и будем использовать информацию которую нам дает вторые производные по функции стоимости? Самый известный метод оптимизации с использованием вторых производных — это метод Ньютона. Основная идея этого метода заключается в минимизации квадратичной аппроксимации функции стоимости. Что же это значит? Давайте разберемся.

Возьмем одномерный случай. Предположим у нас есть функция: $f: mathbb{R} tomathbb{R}$. Что бы найти точку минимума, надо найти ноль её производной, потому-что мы знаем: $f'(x)=0$ находится в минимуме $f(x)$. Аппроксимируем функцию $f$ разложением в ряд Тейлора второго порядка:

$f(x+delta) approx f(x)+f'(x)delta +frac{1}{2}delta f''(x)delta$

Мы хотим найти такую $delta$, что $f(x+delta)$ будет минимумом. Для этого возьмем производную по $x$ и приравняем к нулю:

$frac{d}{dx}bigg(f(x)+f'(x)delta+frac{1}{2}delta f''(x)deltabigg)=f'(x) + f''(x)delta=0impliesdelta=-frac{f'(x)}{f''(x)}$

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

$x_{n+1}=x_n-frac{f'(x_n)}{f''(x_n)}=x_n-(f''(x_n))^{-1}f'(x_n)$

Со временем решение сойдется к минимуму.

Рассмотрим многомерный случай. Положим у нас есть многомерная функция $f: mathbb{R^n}$ тогда:

$f'(x)tonabla f(x)$

$f''(x)to H(f)(x)$

Где $H(f)(x)$ — гессиан (Hessian) или матрица вторых производных. Исходя из этого для обновления параметров имеем такую формулу:

$x_{n+1}=x_n-(H(f)(x))^{-1}nabla f(x_n)$

Проблемы с Методом Ньютона

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

Но у данного метода есть один большой минус. Для оптимизации функции стоимости надо находить матрицу Гессе или гессиан $H$. Положим $mathbb{x}$ — вектор параметров, тогда:

$H(f)(mathbf{x})=begin{pmatrix} frac{delta f}{delta x_1 delta x_1} & frac{delta f}{delta x_1 delta x_2} & cdots & frac{delta f}{delta x_1 delta x_n} \ frac{delta f}{delta x_2 delta x_1} &frac{delta f}{delta x_2 delta x_2} & cdots & frac{delta f}{delta x_2 delta x_n}\ vdots & vdots & ddots & vdots \ frac{delta f}{delta x_m delta x_1} & frac{delta f}{delta x_m delta x_2} & cdots & frac{delta f}{delta x_m delta x_n} end{pmatrix}$

Как можно видеть гессиан это матрица вторых производных размера $ntimes n$ и что бы ее посчитать потребуется $O(n^2)$ вычислительных операций, что может быть очень критично для сетей у которых сотни или тысячи параметров. По мимо этого для решения задачи оптимизации с помощью метода Ньютона необходимо найти обратную матрицу Гессе $H^{-1}(f)(x)$, для этого она должна быть положительно определенной для всех $n$.

Положительно определенная матрица.

Матрица $mathbf A$ размерности $ntimes n$ называется неотрицательно определенной если она удовлетворяет условию: $mathbf v^Tmathbf Amathbf vgeq 0;;; forall ;;; mathbf v in mathbb{R^n}$. Если при этом выполняется строгое неравенство, то матрица называется положительно определенной. Важным свойством таких матриц является их несингулярность, т.е. существование обратной матрицы $mathbf A^{-1}$. Тут нам и приходит на помощь метод под названием Hessian Free оптимизация.


Hessian-Free оптимизация

Основная идея HF оптимизации заключается в том, что за основу мы берем метод Ньютона, но используем более подходящий способ минимизации квадратичной функции. Но для начала введем основные понятия которые понадобятся в дальнейшем.
Пусть $theta=(mathbf W, mathbf b)$ — параметры сети, где $mathbf W$ — матрица весов (weights), $mathbf b$ вектор смещений (biases), тогда выходом сети назовем: $f(x, theta)$, где $x$ — входной вектор. $L(t, f(x,theta))$ — функция потерь (loss function), $t$ — целевое значение. А функцию которую будем минимизировать определим как среднюю от потерь по всем тренировочным примерам (training batch) $S$:

$h(theta)=frac{1}{|S|}sum_{(x, y)in S}L(t, f(x, theta))$

Далее согласно методу Ньютона определим квадратичную функцию, полученную путем разложения в ряд Тейлора второго порядка:

$h(theta+delta)equiv M(delta)=h(theta)+nabla h(theta)^Tdelta+frac{1}{2}delta^THdeltaqquadqquad(1)$

Далее взяв производную по $delta$ от формулы выше и приравнивая её к нулю получаем:

$nabla M(delta)=nabla h(theta) + Hdelta=0qquadqquad(2)$

Что бы найти $delta$ будем использовать метод сопряженных градиентов (conjugate gradient method).

Метод сопряженных градиентов (conjugate gradient method)

Метод сопряженных градиентов (CG) — это итерационный методя решения систем линейных уравнений типа: $mathbf Amathbf x=mathbf b$.

Краткий алгоритм CG:
Входные данные: $mathbf b$, $mathbf A$, $mathbf x_0$, $i=0$ — шаг алгоритма CG
Инициализация:

  1. $mathbf r_0=mathbf b - mathbf Amathbf x_0$ — вектор ошибки (residual)
  2. $mathbf d_0=mathbf r_0$ — вектор направления поиска (search direction)

Повторяем пока не выполнится условие остановки:

  1. $alpha_i=frac{mathbf r_i^Tmathbf r_i}{mathbf d_i^Tmathbf Amathbf d}$
  2. $mathbf x_{i+1}=mathbf x_i + alpha_imathbf d_i$
  3. $mathbf r_{i+1}=mathbf r_i - alpha_imathbf Amathbf d_i$
  4. $beta_i=frac{mathbf r_{i+1}^Tmathbf r_{i+1}}{mathbf r_i^Tmathbf r_i}$
  5. $mathbf d_{i+1}=mathbf r_{i+1} + beta_imathbf d_i$
  6. $i=i+1$

Теперь с помощью метода сопряженных градиентов мы можем решить уравнение (2) и найти $delta$, которая будет минимизировать (1). В нашем случае: $mathbf A equiv H, mathbf b equiv -nabla h(theta), mathbf x equiv delta$.
Остановка CG алгоритма. Останавливать метод сопряженных градиентов можно исходя из разных критериев. Мы будем делать это на основе относительного прогресса в оптимизации квадратичной функции $M$:

$s_i=frac{M(delta_i)-M(delta_{i-w})}{M(delta_i)-M(0)}qquadqquad(3)$

где $w$ — размер «окна» по которому будем считать значение прогресса, $k=max(10, i/10)$. Условием остановки возьмем: $s_i < 0.0001$.
А теперь можно заметить, что главная особенность HF оптимизации заключается в том, что нам не требуется находить гессиан напрямую, а надо лишь найти результат его произведения на вектор.

Умножение гессиана на вектор

Как уже говорилось ранее прелесть данного метода заключается в том, что у нас нет необходимости считай гессиан напрямую. Надо лишь посчитать результат произведения матрицы вторых производных на вектор. Для этого можно представить $H(theta)v$ как производную от $H(theta)$ по направлению $v$:

$H(theta)v=lim_{epsilonto0}frac{nabla h(theta+epsilon v)-nabla h(theta)}{epsilon}$

Но использование данной формулы на практике может вызвать ряд проблем связанные с вычислением при достаточно маленьком $epsilon$. Поэтому существует метод точного вычисления произведения матрицы на вектор. Введем дифференциальный оператор $R_vx$. Он обозначает производную некоторой величины $x$, зависящей от $theta$, по направлению $v$:

$R_vx=lim_{epsilonto0}frac{x(theta+epsilon v)-x(theta)}{epsilon}=frac{delta x}{deltatheta}v$

Отсюда видно, что бы посчитать произведение гессиана на вектор надо вычислить:

$Hv=R(nabla h(theta))qquadqquad(4)$

Некоторые улучшения HF оптимизации

1. Обобщенная матрица Ньютона-Гаусса (generalized Gauss-Newton matrix).
Неопределенность матрицы гессиана является проблемой для оптимизации не выпуклых (non-convex) функций, она может привести к отсутствию нижней границы для квадратичной функции $M$ и как следствие невозможность нахождения ее минимума. Эту проблему можно решить множеством способов. Например, введение доверительного интервала будет ограничивать оптимизацию или же дампинг (damping) на основе штрафа, который добавляет положительный полу-определенный (positive semi-definite) компонент к матрице кривизны $H$ и делает её положительно определенной.

Основываясь же на практических результатах наилучшим способом решить данную проблему является использование матрицы Ньютона-Гаусса $G$ вместо матрицы Гессе:

$G=frac{1}{|S|}sum_{(x,y)in S}J^TH_LJqquadqquad(5)$

где $J$ — якобиан, $H_L$ — матрица вторых производных от функции потерь $L(t, f(x,theta))$.
Для нахождения произведения матрицы $G$ на вектор $v$: $Gv=J^TH_LJv$, сначала найдем произведение якобиана на вектор:

$Jv=R_v(f(x,theta))$

потом вычислим произведение вектора $Jv$ на матрицу $H_L$ и в конце умножим матрицу $J$ на $H_LJv$.

2. Дампинг (damping).
Стандартный метод Ньютона может плохо оптимизировать сильно нелинейные целевые функции. Причиной этому может быть то, что на начальных стадиях оптимизации она может быть очень большой и агрессивной, так как начальная точка находится далеко от точки минимума. Для решения этой проблемы используется дампинг — метод изменения квадратичной функции $M$ или ограничения минимизации таким образом, что новая $delta$ будет лежать в таких пределах, в которых $M$ останется хорошей аппроксимацией $h$.
Регуляризация Тихонова (Tikhonov regularization) или дампинг Тихонова (Tikhonov damping). (Не путать с термином «регуляризация», который обычно используется в контексте машинного обучения) Это самый известный метод дампинга, который добавляет квадратичный штраф к функции $M$:

$hat{M}(delta)equiv M(delta)+frac{lambda}{2}delta^Tdelta=f(theta)+nabla h(theta)^Tdelta+frac{1}{2}delta^That{H}delta$

где $hat{H}=H+lambda I$, $lambda geq0$ — параметр дампинга. Вычисление $Hv$ производится так:

$hat{H}v=(H+lambda I)v=Hv+lambda vqquadqquad(6)$

3. Эвристика Левенберга-Марквардта (Levenberg-Marquardt heuristic).
Для дампинга Тихонова характерно динамическое подстраивание параметра $lambda$. Изменять $lambda$ будем по правилу Левенберга-Марквардта, который часто используется в контексте LM — метода (метод оптимизации, является альтернативой методу Ньютона). Для использования LM — эвристики необходимо посчитать так называемый коэффициент уменьшения (reduction ratio):

$rho=frac{h(theta_k+delta_k)-h(delta_k)}{M_k(delta_k)-M_k(0)}=frac{h(theta_k+delta_k)-h(delta_k)}{nabla h(theta_k)^Tdelta_k+frac{1}{2}delta_k^THdelta_k}qquadqquad(7)$

где $k$ — номер шага HF алгоритма, $delta_k$ — результат работы CG минимизации.
Согласно эвристики Левенберга-Марквардта получаем правило обновления $lambda$:

$begin{cases}if rho>3/4 then lambdagets(2/3)lambda\if rho<1/4 then lambdagets(3/2)lambdaend{cases}qquadqquad(8)$

4. Начальное условие для алгоритма сопряженных градиентов (preconditioning).
В контексте HF оптимизации у нас есть некоторая обратимая матрица трансформации $C$, с помощью которой мы изменяем $M$ таким образом, что $delta=C^{-1}gamma$ и вместо $delta$ минимизируем $gamma$. Применение данной особенности в алгоритме CG требует вычисление трансформированного вектора ошибки $y_i=P^{-1}r_i$, где $P=C^TC$.

Краткий алгоритм PCG (Preconditioned conjugate gradient):
Входные данные: $mathbf b$, $mathbf A$, $mathbf x_0$, $P$, $i=0$ — шаг алгоритма CG
Инициализация:

  1. $mathbf r_0=mathbf b - mathbf Amathbf x_0$ — вектор ошибки (residual)
  2. $mathbf y_0$ — решение уравнения $Py_0=r_0$
  3. $mathbf d_0=mathbf y_0$ — вектор направления поиска (search direction)

Повторяем пока не выполнится условие остановки:

  1. $alpha_i=frac{mathbf r_i^Tmathbf y_i}{mathbf d_i^Tmathbf Amathbf d}$
  2. $mathbf x_{i+1}=mathbf x_i + alpha_imathbf d_i$
  3. $mathbf r_{i+1}=mathbf r_i - alpha_imathbf Amathbf d_i$
  4. $mathbf y_{i+1}$ — решение уравнения $Py_{i+1}=r_i$
  5. $beta_i=frac{mathbf r_{i+1}^Tmathbf y_{i+1}}{mathbf r_i^Tmathbf y_i}$
  6. $mathbf d_{i+1}=mathbf y_{i+1} + beta_imathbf d_i$
  7. $i=i+1$

Выбор матрицы $P$ довольно не тривиальная задача. Так же, на практике использование диагональной матрицы (вместо матрицы с полным рангом) показывает довольно хорошие результаты. Один из вариантов выбора матрицы $P$ — это использование диагональной матрицы Фишера (Empirical Fisher Diagonal):

$P=diag(bar{F})=frac{1}{|S|}sum_{(x,y)in S}(nabla (L(t,f(x,theta)))^2qquadqquad(9)$

5. Инициализация CG — алгоритма.
Хорошей практикой является инициализация начальной $delta_0$, для алгоритма сопряженных градиентов, значением $delta_k$, найденным на предыдущем шаге алгоритма HF. При этом можно использовать некоторую константу распада: $delta_0=epsilondelta_k,  epsilon in (0, 1)$. Стоит отметить, что индекс $k$ относится к номеру шага алгоритма HF, в свою очередь индекс 0 в $delta_0$ относится к начальному шагу алгоритма CG.

Полный алгоритм Hessian-Free оптимизации:
Входные данные: $theta$, $lambda$ — параметр дампинга, $k$ — шаг итерации алгоритма
Инициализация:

  1. $delta_0=0$

Основной цикл HF оптимизации:

  1. Вычисляем матрицу $P$
  2. Находим $delta_k$ решая задачу оптимизации с помощью CG или PCG. $delta_k=CG(-nabla h(theta_k),H,epsilondelta_{k-1},P)$
  3. Обновляем параметр дампинга $lambda$ с помощью эвристики Левенберга-Марквардта
  4. $theta_{k+1}=theta_k+alphatheta_k$, $alpha$ — параметр скорости обучения (learning rate)
  5. $k=k+1$

Таким образом метод Hessian-Free оптимизации позволяет решать задачи нахождения минимума функции большой размерности. При этом не требуется нахождения матрицы Гессе напрямую.


Реализация HF оптимизации на TensorFlow

Теория это конечно хорошо, но давайте попробуем реализовать данный метод оптимизации на практике и посмотрим, что же из этого выйдет. Для написания HF алгоритма я использовал Python и библиотеку глубокого обучения TensorFlow. После этого, в качестве проверки работоспособности я обучил сеть прямого распространения с несколькими слоями на XOR и MNIST датасетах, используя HF метод для оптимизации.

Реализация (создание графа вычислений TensorFlow) для метода сопряженных градиентов.

  def __conjugate_gradient(self, gradients):
    """ Performs conjugate gradient method to minimze quadratic equation 
    and find best delta of network parameters.
    gradients: list of Tensorflow tensor objects
        Network gradients.
    return: Tensorflow tensor object
        Update operation for delta.
    return: Tensorflow tensor object
        Residual norm, used to prevent numerical errors. 
    return: Tensorflow tensor object
        Delta loss. """

    with tf.name_scope('conjugate_gradient'):
      cg_update_ops = []

      prec = None
      #вычисление матрицы P по формуле (9)
      if self.use_prec:
        if self.prec_loss is None:
          graph = tf.get_default_graph()
          lop = self.loss.op.node_def
          self.prec_loss = graph.get_tensor_by_name(lop.input[0] + ':0')
        batch_size = None
        if self.batch_size is None:
          self.prec_loss = tf.unstack(self.prec_loss)
          batch_size = self.prec_loss.get_shape()[0]
        else:
          self.prec_loss = [tf.gather(self.prec_loss, i)
            for i in range(self.batch_size)]
          batch_size = len(self.prec_loss)
        prec = [[g**2 for g in tf.gradients(tf.gather(self.prec_loss, i),
          self.W)] for i in range(batch_size)]
        prec = [(sum(tensor) + self.damping)**(-0.75)
          for tensor in np.transpose(np.array(prec))]
      #основной алгоритм сопряженных градиентов
      Ax = None
      if self.use_gnm:
        Ax = self.__Gv(self.cg_delta)
      else:
        Ax = self.__Hv(gradients, self.cg_delta)

      b = [-grad for grad in gradients]
      bAx = [b - Ax for b, Ax  in zip(b, Ax)]

      condition = tf.equal(self.cg_step, 0)
      r = [tf.cond(condition, lambda: tf.assign(r,  bax),
        lambda: r) for r, bax  in zip(self.residuals, bAx)]

      d = None
      if self.use_prec:
        d = [tf.cond(condition, lambda: tf.assign(d, p * r), 
          lambda: d) for  p, d, r in zip(prec, self.directions, r)]
      else:
        d = [tf.cond(condition, lambda: tf.assign(d, r), 
          lambda: d) for  d, r in zip(self.directions, r)]

      Ad = None
      if self.use_gnm:
        Ad = self.__Gv(d)
      else:
        Ad = self.__Hv(gradients, d)

      residual_norm = tf.reduce_sum([tf.reduce_sum(r**2) for r in r])

      alpha = tf.reduce_sum([tf.reduce_sum(d * ad) for d, ad in zip(d, Ad)])
      alpha = residual_norm / alpha

      if self.use_prec:
        beta = tf.reduce_sum([tf.reduce_sum(p * (r - alpha * ad)**2) 
          for r, ad, p in zip(r, Ad, prec)])
      else:
        beta = tf.reduce_sum([tf.reduce_sum((r - alpha * ad)**2) for r, ad 
          in zip(r, Ad)])

      self.beta = beta
      beta = beta / residual_norm

      for i, delta  in reversed(list(enumerate(self.cg_delta))):
        update_delta = tf.assign(delta, delta + alpha * d[i],
          name='update_delta')
        update_residual = tf.assign(self.residuals[i], r[i] - alpha * Ad[i],
          name='update_residual')
        p = 1.0
        if self.use_prec:
          p = prec[i]
        update_direction = tf.assign(self.directions[i],
          p * (r[i] - alpha * Ad[i]) + beta * d[i], name='update_direction')
        cg_update_ops.append(update_delta)
        cg_update_ops.append(update_residual)
        cg_update_ops.append(update_direction)

      with tf.control_dependencies(cg_update_ops):
        cg_update_ops.append(tf.assign_add(self.cg_step, 1))
      cg_op = tf.group(*cg_update_ops)

    dl = tf.reduce_sum([tf.reduce_sum(0.5*(delta*ax) + grad*delta)
      for delta, grad, ax in zip(self.cg_delta, gradients, Ax)])

    return cg_op, residual_norm, dl

Код для вычисления матрицы $P$ для нахождения начального условия (preconditioning) приведен ниже. При этом, так как Tensorflow суммирует результат вычисления градиентов по всему множеству подаваемых примеров обучения, пришлось немного поизвращаться, что бы получить градиент отдельно для каждого примера, что сказалось на численной стабильности решения. Поэтому использование preconditioning возможно, как говорится, на свой страх и риск.

 prec = [[g**2 for g in tf.gradients(tf.gather(self.prec_loss, i),
   self.W)] for i in range(batch_size)]

Вычисление произведения гессиана на вектор (4). При этом используется дампинг Тихонова (6).

  def __Hv(self, grads, vec):
    """ Computes Hessian vector product.
    grads: list of Tensorflow tensor objects
        Network gradients.
    vec: list of Tensorflow tensor objects
        Vector that is multiplied by the Hessian.
    return: list of Tensorflow tensor objects
        Result of multiplying Hessian by vec. """

    grad_v = [tf.reduce_sum(g * v) for g, v in zip(grads, vec)]
    Hv = tf.gradients(grad_v, self.W, stop_gradients=vec)
    Hv = [hv + self.damp_pl * v for hv, v in zip(Hv, vec)]

    return Hv

Когда же я захотел использовать обобщенную матрицу Ньютона-Гаусса (5) я натолкнулся на небольшую проблему. А именно, TensorFlow не умеет считать произведение Якобиана на вектор как это делает другой фреймворк для глубокого обучения Theano (в Theano есть функция Rop специально предназначенная для этого). Пришлось делать аналог операции в TensorFlow.

  def __Rop(self, f, x, vec):
    """ Computes Jacobian vector product.
    f: Tensorflow tensor object
        Objective function.
    x: list of Tensorflow tensor objects
        Parameters with respect to which computes Jacobian matrix.
    vec: list of Tensorflow tensor objects
        Vector that is multiplied by the Jacobian.
    return: list of Tensorflow tensor objects
        Result of multiplying Jacobian (df/dx) by vec. """

    r = None
    if self.batch_size is None:
      try:
        r = [tf.reduce_sum([tf.reduce_sum(v * tf.gradients(f, x)[i])
             for i, v in enumerate(vec)])
             for f in tf.unstack(f)]
      except ValueError:
        assert False, clr.FAIL + clr.BOLD + 'Batch size is None, but used '
          'dynamic shape for network input, set proper batch_size in '
          'HFOptimizer initialization' + clr.ENDC
    else:
      r = [tf.reduce_sum([tf.reduce_sum(v * tf.gradients(tf.gather(f, i), x)[j]) 
           for j, v in enumerate(vec)])
           for i in range(self.batch_size)]

    assert r is not None, clr.FAIL + clr.BOLD +
      'Something went wrong in Rop computation' + clr.ENDC

    return r

И потом уже реализовывать произведение обобщенной матрицы Ньютона-Гаусса на вектор.

  def __Gv(self, vec):
    """ Computes the product G by vec = JHJv (G is the Gauss-Newton matrix).
    vec: list of Tensorflow tensor objects
        Vector that is multiplied by the Gauss-Newton matrix.
    return: list of Tensorflow tensor objects
        Result of multiplying Gauss-Newton matrix by vec. """

    Jv = self.__Rop(self.output, self.W, vec)
    Jv = tf.reshape(tf.stack(Jv), [-1, 1])
    HJv = tf.gradients(tf.matmul(tf.transpose(tf.gradients(self.loss,
      self.output)[0]), Jv), self.output, stop_gradients=Jv)[0]
    JHJv = tf.gradients(tf.matmul(tf.transpose(HJv), self.output), self.W,
      stop_gradients=HJv)
    JHJv = [gv + self.damp_pl * v for gv, v in zip(JHJv, vec)]

    return JHJv

Функция основного процесса обучения представлена ниже. Сначала производится минимизация квадратичной функции с помощью CG/PCG, потом происходит уже основное обновление весов сети. Так же подстраивается параметр дампинга на основе эвристики Левенберга-Марквардта.

  def minimize(self, feed_dict, debug_print=False):
    """ Performs main training operations.
    feed_dict: dictionary
        Input training batch.
    debug_print: bool
        If True prints CG iteration number. """

    self.sess.run(tf.assign(self.cg_step, 0))
    feed_dict.update({self.damp_pl:self.damping})

    if self.adjust_damping:
      loss_before_cg = self.sess.run(self.loss, feed_dict)

    dl_track = [self.sess.run(self.ops['dl'], feed_dict)]
    self.sess.run(self.ops['set_delta_0'])

    for i in range(self.cg_max_iters):
      if debug_print:
        d_info = clr.OKGREEN + 'r[CG iteration: {}]'.format(i) + clr.ENDC
        sys.stdout.write(d_info)
        sys.stdout.flush()

      k = max(self.gap, i // self.gap)

      rn = self.sess.run(self.ops['res_norm'], feed_dict)
      #ранняя остановка для предотвращения численной ошибки
      if rn < self.cg_num_err:
        break

      self.sess.run(self.ops['cg_update'], feed_dict)
      dl_track.append(self.sess.run(self.ops['dl'], feed_dict))

      #ранняя остановка алгоритма, основываясь на формуле (3)
      if i > k:
        stop = (dl_track[i+1] - dl_track[i+1-k]) / dl_track[i+1]
        if not np.isnan(stop) and stop < 1e-4:
          break

    if debug_print:
      sys.stdout.write('n')
      sys.stdout.flush()

    if self.adjust_damping:
      feed_dict.update({self.damp_pl:0})
      dl = self.sess.run(self.ops['dl'], feed_dict)
      feed_dict.update({self.damp_pl:self.damping})

    self.sess.run(self.ops['train'], feed_dict)

    if self.adjust_damping:
      loss_after_cg = self.sess.run(self.loss, feed_dict)
      #коэффициент уменьшения (7)
      reduction_ratio = (loss_after_cg - loss_before_cg) / dl

      #эвристика Левенберга-Марквардта (8)
      if reduction_ratio < 0.25 and self.damping > self.damp_num_err:
        self.damping *= 1.5
      elif reduction_ratio > 0.75 and self.damping > self.damp_num_err:
        self.damping /= 1.5

Тестирование работы HF оптимизации

Протестируем написанный HF оптимизатор, для этого будем использовать простой пример c XOR датасетом и более сложный с MNIST датасетом. Для того, что бы посмотреть результаты обучения и визуализировать некоторую информацию воспользуемся TesnorBoard. Хочется так же заметить, что получился довольно сложный граф вычислений TensorFlow.

Hessian-Free оптимизация с помощью TensorFlow - 159
Граф вычислений TensorFlow.

Архитектура и обучение сети на XOR датасете.
Создадим простую сеть размером: 2 входных нейрона, 2 скрытых и 1 выходной. В качестве функции активации будем использовать сигмоиду. В качестве функции потерь используем log-loss.

  #определение функции потерь
  """ Log-loss cost function """
  loss = tf.reduce_mean(( (y * tf.log(out)) + 
    ((1 - y) * tf.log(1.0 - out)) ) * -1, name='log_loss')

  #XOR датасет
  XOR_X = [[0,0],[0,1],[1,0],[1,1]]
  XOR_Y = [[0],[1],[1],[0]]
  
  #создание оптимизатора
  sess = tf.Session()
  hf_optimizer = HFOptimizer(sess, loss, y_out, 
    dtype=tf.float64, use_gauss_newton_matrix=True)

  init = tf.initialize_all_variables()
  sess.run(init)

  #цикл тренировки
  max_epoches = 100
  print('Begin Training')
  for i in range(max_epoches):
    feed_dict = {x: XOR_X, y: XOR_Y}
    hf_optimizer.minimize(feed_dict=feed_dict)
    if i % 10 == 0:
      print('Epoch:', i, 'cost:', sess.run(loss, feed_dict=feed_dict))
      print('Hypothesis ', sess.run(out, feed_dict=feed_dict))

Теперь сравним результаты обучения с использованием HF оптимизации (с матрицей Гессе), HF оптимизации (с матрицей Ньютона-Гаусса) и обычным градиентным спуском с параметром скорости обучения равным 0.01. Количество итераций равно 100.

Hessian-Free оптимизация с помощью TensorFlow - 160
Loss для градиентного спуска (красная линия). Loss для HF оптимизации с матрицей Гессе (оранжевая линия). Loss для HF оптимизации с матрицей Ньютона-Гаусса (синяя линия).

При этом видно что быстрее всего сходится HF оптимизация с использованием матрицы Ньютона-Гаусса, для градиентного же спуска 100 итерации оказалось очень мало. Для того, что бы функция потерь при градиентном спуске была сопоставима с HF оптимизацией потребовалось около 100000 итераций.

Hessian-Free оптимизация с помощью TensorFlow - 161
Loss для градиентного спуска, 100000 итераций.

Архитектура и обучение сети на MNIST датасете.
Для решения задачи распознавания рукописных чисел создадим сеть размером: 784 входных нейрона, 300 скрытых и 10 выходных. В качестве функции потерь будем использовать кросс-энтропию. Размер мини батча подаваемого в ходе обучения равен 50-ти.

  with tf.name_scope('loss'):
    one_hot = tf.one_hot(t, n_outputs, dtype=tf.float64)
    xentropy = tf.nn.softmax_cross_entropy_with_logits(labels=one_hot, logits=y_out)
    loss = tf.reduce_mean(xentropy, name="loss")

  with tf.name_scope("eval"):
    correct = tf.nn.in_top_k(tf.cast(y_out, tf.float32), t, 1)
    accuracy = tf.reduce_mean(tf.cast(correct, tf.float64))

  n_epochs = 10
  batch_size = 50

  with tf.Session() as sess:
    """ Initializing hessian free optimizer """
    hf_optimizer = HFOptimizer(sess, loss, y_out, dtype=tf.float64, batch_size=batch_size, 
      use_gauss_newton_matrix=True)
    init = tf.global_variables_initializer()
    init.run()
    #основной процесс обучения
    for epoch in range(n_epochs):
      n_batches = mnist.train.num_examples // batch_size
      for iteration in range(n_batches):
        x_batch, t_batch = mnist.train.next_batch(batch_size)
        hf_optimizer.minimize({x: x_batch, t: t_batch})
        if iteration%10==0:
          print('Batch:', iteration, '/', n_batches)
          acc_train = accuracy.eval(feed_dict={x: x_batch, t: t_batch})
          acc_test = accuracy.eval(feed_dict={x: mnist.test.images,
                        t: mnist.test.labels})
          print('Loss:', sess.run(loss, {x: x_batch, t: t_batch}))
          print('Target', t_batch[0])
          print('Out:', sess.run(y_out_sm, {x: x_batch, t: t_batch})[0])
          print(epoch, "Train accuracy:", acc_train, "Test accuracy:", acc_test)

      acc_train = accuracy.eval(feed_dict={x: x_batch, t: t_batch})
      acc_test = accuracy.eval(feed_dict={x: mnist.test.images,
                        t: mnist.test.labels})
      print(epoch, "Train accuracy:", acc_train, "Test accuracy:", acc_test)

Так же как и в случае с XOR сравним результаты обучения с использованием HF оптимизации (с матрицей Гессе), HF оптимизации (с матрицей Ньютона-Гаусса) и обычным градиентным спуском с параметром скорости обучения равным 0.01. Количество итераций равно 200, т.е. если размер мини батча равен 50, то 200 — это не полная эпоха (не все примеры из обучающей выборки использованы). Я это сделал для того, что бы быстрее все протестировать, но даже из этого видна общая тенденция.

Hessian-Free оптимизация с помощью TensorFlow - 162
Рисунок слева точность для тестовой выборки. Рисунок справа точность для тренировочной выборки. Точность для градиентного спуска (красная линия). Точность для HF оптимизации с матрицей Гессе (оранжевая линия). Точность для HF оптимизации с матрицей Ньютона-Гаусса (синяя линия).

Hessian-Free оптимизация с помощью TensorFlow - 163
Loss для градиентного спуска (красная линия). Loss для HF оптимизации с матрицей Гессе (оранжевая линия). Loss для HF оптимизации с матрицей Ньютона-Гаусса (синяя линия).

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

Hessian-Free оптимизация с помощью TensorFlow - 164
Одна полная эпоха обучения. Рисунок слева точность для тестовой выборки. Рисунок справа точность для тренировочной выборки. Точность для градиентного спуска (бирюзовая линия). Loss для HF оптимизации с матрицей Ньютона-Гаусса (розовая линия).

Hessian-Free оптимизация с помощью TensorFlow - 165
Одна полная эпоха обучения. Loss для градиентного спуска (бирюзовая линия). Loss для HF оптимизации с матрицей Ньютона-Гаусса (розовая линия).

При использовании метода сопряженных градиентов с начальным условием для алгоритма сопряженных градиентов (preconditioning) сами вычисления значительно замедлились и сходились не быстрее чем при обычном CG.

Hessian-Free оптимизация с помощью TensorFlow - 166
Loss для HF оптимизации с использованием PCG алгоритма.

Из всех этих графиков можно заметить, что наилучший результат показала HF оптимизация с использованием матрицы Ньютона-Гаусса и стандартного метод сопряженных градиентов.
Полный код можно посмотреть на GitHub.


Итоги

В итоге была создана реализация HF алгоритма на Python с помощью библиотеки TensorFlow. В ходе создания я столкнулся с некоторыми проблемами при реализации основных особенностей алгоритма, а именно: поддержки матрицы Ньютона-Гаусса и preconditioning-а. Связано это было с тем, что все таки TensorFlow не такая гибкая библиотека как хотелось бы и не сильно предназначена для исследований. Для экспериментальных целей все таки лучше использовать Theano, так как она дает большую свободу. Но я изначально решил для себя сделать все это с помощью TensorFlow. Работа программы была протестирована и можно было увидеть, что наилучшие результаты дает HF алгоритм с использованием матрицы Ньютона-Гаусса. Использование же начального условия для алгоритма сопряженных градиентов (preconditioning) давало не стабильные численные результаты и сильно замедляло вычисления, но как мне кажется, это связано с особенностью именно TensorFlow (для реализации preconditioning пришлось сильно извращаться).


Источники

В данной статье теоретические аспекты Hessian — Free оптимизации описаны довольно кратко, что бы можно было понять основную суть алгоритмов. Если же необходимо более подробное описание материала, то ниже привожу источники откуда я брал основную теоретическую информацию, на основе которой сделал Python реализацию HF метода.

1) Training Deep and Recurrent Networks with Hessian-Free Optimization (James Martens and Ilya Sutskever, University of Toronto) — полное описание HF — оптимизации.
2) Deep learning via Hessian-free optimization (James Martens, University of Toronto) — статья с результатами использования HF — оптимизации.
3) Fast Exact Multiplication by the Hessian (Barak A. Pearlmutter, Siemens Corporate Research) — подробное описание умножения матрицы Гессе на вектор.
4) An introduction to the Conjugate Gradient Method without the Agonizing Pain (Jonathan Richard Shewchuk, Carnegie Mellon University) — подробное описание метода сопряженных градиентов.

Автор: M00nL1ght

Источник

Поделиться

* - обязательные к заполнению поля