Neural ODE08.07.2024 18:15

Аттрактор Лоренца

В статье описана идея нейронных обыкновенных дифференциальных уравнений (Neural ODEs) — подхода в области глубокого обучения, который объединяет методы численного решения дифференциальных уравнений с нейронными сетями. Neural ODEs позволяют моделировать непрерывное изменение скрытых состояний, что открывает новые возможности для анализа временных рядов, обработки сигналов и динамических систем.

Пользуясь случаем, оставлю ссылку на свой канал — notmagicneuralnetworks ✨

1. Предыстория

1.1. ResNet

Pанние архитектуры нейронных сетей состояли преимущественно из линейных слоев и функция активаций. Даже не в очень глубоких сетях остро стояла проблема затухающего градиента (vanishing gradients problem): когда градиенты передаются через множество слоев, они могут постепенно уменьшаться до незначительных значений. Это может привести к тому, что веса на начальных слоях практически не обновляются и не учатся эффективно, что замедляет или делает невозможным обучение.

Пример

Например, пусть после линейного слоя стоит сигмоидная функция активации $\sigma(x) = \frac{1}{1 + e^{-x}}$ . Во время back propagation ищется значение производной функции активации $\sigma'(x) = \sigma(x)(1 − \sigma(x))$ . На рисунке ниже представлены функции сигмоиды и ее производной.

Функция сигмоиды и ее производная

При поиске градиента, нужно проходить через производную функции активации. При этом, ее максимальное значение равно 0.25 . Если слоев несколько, то может оказаться так, что на начальных слоях градиенты почти равны нулю. Из-за этого веса на первых слоях могут обновляться очень медленно. Это и называется проблема затухания градиентов (vanishing gradients problem).

С этой проблемой боролись по-разному, например, использовали другие функции активации, в VGG применяли особые схемы обучения, а GoogLeNet добавили вспомогательные функции потерь. Чуть подробнее можно почитать здесь.

В 2015 году исследователями из компании Microsoft была предложена архитектура под названием Residual neural network (или ResNet). Она состояла из residual block (residual connection, skip connection), где входные данные передаются через слои (блоки) без изменений и соединяются с выходным данным блока:

Residual Block в рекуррентной записи:

$x_t = x_{t-1} + f(x_{t-1})$

где $x_{t-1}$ — данные с предыдущего слоя, — слой нейросети, $x_{t}$ — текущие данные.

def f(x, t, theta):
  return nnet(z, theta[t]))

def resnet(x, theta):
  for t in [1:T]:
    x = x + f(x, t, theta)
  return x

Благодаря такому перекидыванию данных через слои, ResNet хорошо решает проблему затухающего градиента.

Производная функции потерь, когда проходит через Residual Block будет равна

$\frac{dL}{dx} = \frac{dL}{d\varphi} \frac{d\varphi}{dx} = \frac{dL}{d\varphi} \left(1 + f'(x) \right)$ , где $\varphi = x + f(x)$

что в результате позволяет градиенту не затухать.

1.2. Euler method

Рекуррентная запись ResNet очень напоминает метод Эйлера — численный метод решения обыкновенных дифференциальный уравнений (ОДУ) с заданными начальными условиями.

Пусть задано некоторое дифференциальные уравнение с начальными условиями:

$\frac{dx}{dt} = f(x(t), t), \quad x(t_0) = x_0$

Нужно найти решение этого уравнения в конечный момент времени , то есть x(T) .

Чтобы найти решение методом Эйлера, нужно разбить отрезок времени [t_0, T] на равные промежутки. Обозначим $h = t_n - t_{n-1}$ — шаг разбиения временного отрезка. Тогда t_0 < t_1 < t_2 < ... < T — называются узлами. Очевидно, что чем меньше будет шаг разбиения, тем более точным окажется решение.

Решение ОДУ методом Эйлера находится по рекуррентной формуле

$x(t_{i+1}) = x(t_{i}) + h f(x(t_i), t_i)$

и она с точностью до повторяет формулу Residual Block.

Пример

2. Neural Ordinary Differential Equations

2.1. Постановка

Заметив такое сходство, авторы статьи Neural ordinary differential equations предложили представить нейронную сеть как обыкновенное дифференциальное уравнение. Можно сказать, что если у вас есть очень много рекуррентных слоев с малым разбиением по времени , то по-сути вы решаете ОДУ.

Переобозначим переменные в соответствии со статьей NeuralODE.

$\frac{dz}{dt} = f(z(t), t, \theta), \quad z(t_0) = z_0$

где z(t) — состояние (или точка) фазового пространства Neural ODE, — нейронная сеть, $\theta$ — обучаемые параметры нейронной сети, z_0 — начальные условия (входные данные нейронной сети).

На выходе из нейронной сети мы хотим получить значение z(t_1) — куда придет система в конечный момент времени t_1 .Решение может быть записано в интегральной форме:

$z (t_1 ) =z (t_0 ) + ∫^{t_1}_{t_0} f(z(t),t,\theta) dt$

Обозначим ODESolve — численный метода (например, Эйлера), с помощью которого можно решить ДУ. Запишем это решение как:

$z (t_1) =ODESolve(z(t_0), f, t_0 ,t_1, \theta)$

def f(z, t, theta):
  return nnet(z[t], theta))

def resnet(z, theta):
  return ODESolve(f, z, t0, t1, theta)

Хочется, чтобы истинный z(t_1) совпадал с предсказанным нейронной сетью $\hat z(t_1)$ .

Введем функцию потерь:

$L (z(t_1)) = L \left( z(t_0) + \int^{t_1}_{t_0} f(z(t),t,\theta)dt \right) = L \left( ODESolve( z(t_0), f,t_0, t_1 ,\theta) \right)$

Функция потерь может быть MSE, MAE и т.п. Обычно, она зависит только от значения в конечный момент времени, но, ее можно задать и на всю траекторию.

2.2. Численные методы решения ОДУ (ODE Solvers)

Вообще, ОДУ не обязательно решать методом Эйлера. Он являлся исторически первым и простейшим численным методом решения систем ОДУ первого порядка точности (то есть ошибка линейно зависит от шага интегрирования ).

Повысить точность и устойчивость вычисления решения можно с помощью модифицированного метода Эйлера (или midpoint method) второго порядка точности:

$x(t_{i+1}) = x(t_i) +hf \left(x(t_i)+\frac{1}{2}hf\left(x(t_i), t_i \right), \quad t_i + \frac{1}{2}h \right)$

Авторы статьи NeuralODE использовали метод Рунге-Кутты (Runge-Kutta methods). Вообще, методы Рунге-Кутты — это целый класс численных методов, включающий в себя в том числе и метод Эйлера. Однако, под классическим методом Рунге-Кутты подразумевается метод четвёртого порядка точности:

$x(t_{i+1}) = x(t_i) + \frac{1}{6} \left(k_1+2k_2+2k_3+k_4\right)$

$k_1 = f \left( x(t_i),t_i \right)$

$k_2 =f \left( x(t_i) + \frac{k_1}{2}, t_i+\frac{h}{2} \right)$

$k_3 = f \left(x(t_i)+\frac{k_2}{2}, t_i + \frac{h}{2} \right)$

$k_4 = f \left(x(t_i) + k_3, t_i+h \right)$

Все выше описанные методы относятся к методам с фиксированным шагом (fixed-step solvers). Однако, они обладают следующим недостатком: если выбрать шаг слишком большой, то решения окажутся сильно не точными, если же выбрать шаг слишком маленьким, получается очень большое количество вычислений.

Компромиссным решением являются адаптивные методы (adaptive solvers), которые автоматически регулируют шаг интегрирования в процессе вычисления решения. Они адаптируются к изменяющимся условиям задачи, чтобы обеспечить более точное решение и избежать излишних вычислительных затрат.

Например, один из широко использующихся методов — метод Дорманда–Принса (Dorman–Prince method). Он использует сразу два метода Рунге-Кутты 4 и 5 порядка. Метод 4 порядка (основной решатель) вычисляет итоговое решение ОДУ. Метод 5 порядка, обладающий лучшей точностью, нужен для оценки ошибки и регулирования шага основного решателя. Таким образом, когда решение более простое — выбирается шаг побольше, когда решение усложняется — шаг выбирается поменьше.

Любой из вышеупомянутых методов может быть использован в Neural ODE. Промежуток времени [t_0, t_1] играет роль слоев в нейросети. Каждая новая итерация в ODESolver — это новый рекуррентный слой.

2.3. Обучение: backpropagation & adjoint method

Конечно, можно считать градиент по всем слоям Neural ODE обычным backpropagation, но из-за большого количества слоев (маленького шага разбиения ) и сложного ODESolver (такого как метод Дорманда–Принса) это займет слишком много памяти и времени.

Ноутбук, в котором при обратном проходе используется Backpropagation through time (не мой): Пошаговое руководство по обучению модели Neural ODE с использованием BPTT.

В статье для поиска градиента по параметрам был предложен так называемый adjoint method (метод сопряженного состояния), который основан на принципе максимума Понтрягина.

Оптимальное управление движением. Принцип максимума Понтрягина.

1. Постановка задачи

Рассматривается управляемая система, описываемых с помощью ОДУ:

$\dot y = f(y, u), \quad y(t_0) = x^* \quad (1)$

где — вектор состояния размерности , — вектор управления, — непрерывная вектор-функция по совокупности аргументов и непрерывно дифференцируемая по .

Условием окончания процесса служит первое попадание в момент времени t_k на гладкое и без особых точек многообразие $M \subset \mathbb R^n$ .

Предполагается, что ресурсы управления ограничены, то есть управление $u(\cdot)$ принадлежит функциональному множеству .

Рассматривается экстремальная задача:

$\varphi_0(y(t_k)) \rightarrow \inf_{u(\cdot) \in U} \quad (2)$

Задача оптимального управления заключается в отыскании управления , которое переводит систему (1) из начального состояния на конечное многообразие, удовлетворяет ограничению на управление и при этом функционал (2) достигает своего минимума на траекториях системы (1).

Пара $\{ y^0(\cdot), u^0(\cdot), t \in [t_0, t_k^0] \}$ называется оптимальным процессом.

Простой пример чтобы понять текст выше:
Пусть имеется машинка, передвижение которой записано в виде дифференциального уравнения (1).
Управление машиной, например, таково, что мы можем двигаться только вперед или назад, причем с ограниченной скоростью. То есть функциональное множество $\{u_{-}, u_{+} \}$ .
Нам необходимо попасть из точки А (начального состояния), в точку B (конечное состояние) как можно быстрее. То есть функционал должен минимизировать время.
Задача оптимального управления — найти такое управление , при котором точка B достигается как можно быстрее.

2. Решение

Вводится функция Понтрягина:

$H (\psi, y, u) = \psi^T f(y, u)$

где $\psi(t)$ — решение системы ОДУ, сопряженных уравнениям в вариациях на оптимальном решении исходной системы:

$\dot{\psi} = - \frac{\partial f(y^0(t), u^0(t))}{\partial y}^\top \psi \quad (3)$

Теорема: Если $\{ y^0(\cdot), u^0(\cdot), t \in [t_0, t_k^0] \}$ — оптимальный процесс, то существует ненулевая пара $\{ \lambda_0 \geq 0, \phi(\cdot) \}$ такая, что выполняются следующие условия:

функция Понтрягина достигает максимума на множестве точек непрерывности оптимального управления $t \in T \subset [t_0; t_k^0]$
$\max_{u_{-} \leq u(t) \leq u_{+}} H(\psi(t), y^0(t), u(t)) = H( \psi(t), y^0(t), u^0(t)) \quad (4)$
вектор $\psi(t_k^0) + \lambda_0 \frac{\partial \varphi_0 (y^0(t_k^0))}{\partial y} \quad (5)$ ортогонален к многообразию в точке — условие трансверсальности
условие стационарности гамильтониана почти всюду на
$\mathcal{H}(t) = H(\psi(t), y^0(t), u^0(t)) \equiv 0 \quad (6)$

С помощью сформулированной выше теоремы поиск оптимального управления сводится к решению двухточечной краевой задачи для совокупности систем (1) и (3), где на левом конце задано условий y(t_0) = y^* , а на правом условий $y(t_k) \in M$ для y(t_k) , и n - m условий трансверсальности (5) для $\psi(t_k)$ . В процессе интегрирования двухточечной задачи при каждом $t \in T$ надо решать одномерную задачу оптимизации (4) по u(t) .

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

Если интересно подробнее прочитать про принцип максимума, то закинула в репозиторий материалы с примерами и ноутбук «О методе последовательных приближений для решения задач оптимального управления».
В этом месте очень хочется поблагодарить Бугрова Дмитрия Игоревича ❤️

Поиск градиента выглядит следующим образом:

При проходе вперед (forward pass), с помощью выбранного численного метода, решается дифференциальное уравнение
$\frac{dz}{dt} = f(z(t), t, \theta), \quad z(t_0) = z_0$
Запоминается решение только в конечный момент времени , а все промежуточные вычисление забываются.
Вводится сопряженная переменная adjoint: $a = \frac{dL}{dz(t)}$
Для нее записывается дифференциальное уравнение: $\frac{da(t)}{dt} = -a(t) \frac{df(z(t), t, \theta)}{dz}$
с конечным условием $a(t_1) = \frac{dL}{dz(t_1)}$
При обратном проходе оно решается в обратном времени тем же численным методом

$\frac{dL}{d \theta} = \int_{t_1}^{t_0} a(t) \frac{df(z(t), t, \theta)}{d \theta} dt$

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

Из-за того что не нужно запоминать промежуточные вычисления (в отличие от backpropagation), обратный проход занимает O (1) памяти. Хотя по времени, из-за численных методов, получается долго.

В реализации используют аугментированное состояние: в пространство состояний z(t) дополнительно вводятся переменные a(t) , , $\theta$ . Проход вперед и назад будет считаться по всем переменным.

Если сравнить с ПМП

Если сравнивать adjoint method с принципом Максимума, то сопряженная переменная adjoint a(t) — это $\psi(t)$ в принципе максимума (они же множители Лагранжа в функции Понтрягина). Параметр обучения $\theta$ , видимо, управление u(t) . Однако, находятся параметры путем интегрирования, а не максимизации функции Понтрягина.

Можно дополнительно почитать: Deep Implicit Layers, Знакомство с Neural ODE.

Ссылка на ноутбук.

Neural ODE могут использоваться, например, в задачах идентификации динамических систем; при работе с нерегулярными временными рядами. Подробнее про применение можно посмотреть в следующих видео:

Алексей Окунев | Нейронные дифференциальные уравнения для задач с выраженной динамикой
Research Seminar. Neural ODE: Part 1, Part 2, Part 3, Part 4.