Нейронные дифференциальные уравнения
Предисловие
Дифференциальные уравнения — естественный язык физики и природы. Вспомните, как в школе нам объясняли второй закон Ньютона о движении:
Однако методы машинного обучения и другие методы, основанные на данных, доказали свою способность успешно предсказывать поведение сложных систем. Возникает вопрос: а как мы можем объединить машинное обучение с дифференциальными уравнениями для создания более надежных моделей?
Вспоминаем обыкновенные дифференциальные уравнения
В общем случае дифференциальное уравнение может быть записано в таком виде:
\( \frac{d}{dt}z(t) := \dot z(t) = f\left(t,z(t), \theta\right) \)
где \( z(t) \) представляет состояние, \( f(\cdot) \) обозначает так называемое векторное поле, \( t \) соответствует времени, а \( \theta \)— параметрам. Говоря простым языком, ОДУ моделируют скорость изменения переменных.
Например, дифференциальное уравнение, описывающее простой маятник, выглядит следующим образом:
\( \dot z(t) = \begin{bmatrix} z_2(t)\\ -\frac{g}{L} \sin(z_1(t)) \end{bmatrix} \)
где \( z_1(t):=\theta \) соответствует углу отклонения маятника от вертикали, \( z_2(t):=\dot \theta \) представляет угловую скорость, \( g \) обозначает гравитационную постоянную, а \( L \) — длину маятника. Интуитивно, ОДУ маятника показывают изменения угла и угловой скорости маятника.
При работе с ОДУ крайне важно указывать начальное состояние. Без знания начального условия \( z(t_0) \) невозможно решить дифференциальное уравнение, то есть проинтегрировать его. Говоря об интегрировании, следует отметить, что большинство значимых дифференциальных уравнений не имеет аналитических решений. Однако есть численные методы для приближенного нахождения решений в дискретных точках. Эти методы включают интегрирование уравнения на заданном интервале:
\( z(t_1) = z(t_0) + \int_{t_0}^{t_1} f(t, z(t), \theta)\ dt = \text{ODESolve}(z(t_0), f, \theta, t_0, t_1) \)
И сейчас мы рассмотрим два простых численных метода для решения ОДУ, а затем коротко обсудим более продвинутые методы.
Метод Эйлера
Один из самых простых методов численного интегрирования — это метод Эйлера. Он вытекает из базового определения касательной аппроксимации градиента в точке:
\( \frac{dz(t)}{dt} \approx \frac{z(t+\Delta t) — z(t)}{\Delta t} \)
где \( \Delta t \) — фиксированная длина шага. Перегруппировав это выражение, получаем выражение для интегратора Эйлера:
\( z(t+\Delta t) = z(t) + \Delta t \frac{dz(t)}{dt} = z(t) + \Delta t \ f\big(t, z(t), \theta \big). \)
Интегрирование ОДУ при помощи метода Эйлера мы начинаем с начального времени \( t_0 \), начального состояния \( z(t_0) \) и шага \( \Delta \). Затем мы итеративно вычисляем траекторию ОДУ, пока не достигнем конечного времени \( t_N := t_0 + N \Delta t \):
\( z(t_1) = z(t_0) + \Delta t\ f(t_0, z(t_0), \theta),\quad t_1 = t_0 + \Delta t \\ z(t_2) = z(t_1) + \Delta t\ f(t_1, z(t_1), \theta),\quad t_2 = t_1 + \Delta t \\ \dots \\ z(t_N) = z(t_{N-1}) + \Delta t\ f(t_{N-1}, z(t_{N-1}), \theta),\quad t_N = t_{N-1} + \Delta t \)
Одна из основных идей этого метода — мы рассчитываем траекторию ОДУ, последовательно делая небольшие шаги в направлении касательной.
Методы Рунге-Кутты
Несмотря на доступность для понимания и легкость реализации метода Эйлера, он не часто используется на практике из-за своей низкой точности. Метод Рунге-Кутты четвертого порядка (RK4) более точен, но при этом вычислительно эффективен и намного чаще применяется на практике. RK4 вычисляет состояния на следующем временном шаге путем выполнения четырех промежуточных шагов:
\( z(t + \Delta t) = z(t) + \frac{1}{6} \Delta t(k_1 + k_2 + k_3 + k_4) \)
где
\( k_1 = f(t, z(t), \theta) \\ k_2 = f(t + \frac{\Delta t}{2}, z(t) + \Delta t\ \frac{k_1}{2}, \theta) \\ k_3 = f(t + \frac{\Delta t}{2}, z(t) + \Delta t\ \frac{k_2}{2}, \theta)\\ k_4 = f(t + \Delta t, z(t) + \Delta t\ k_3, \theta) \)
Мы не будем углубляться в подробности получения этого уравнения. Но мы коротко отметим: вывод уравнений основан на разложении \( z(t + \Delta t) \) и \( f(t + \Delta t, z + \Delta z) \) в ряд Тейлора.
Интегрирование ОДУ с использованием RK4 такое же, как и у метода Эйлера и у всех других методов с постоянным шагом: задаем начальное время \( t_0 \), начальное состояние \( z(t_0) \), размер шага \( \Delta t \) и \( N \); затем в цикле последовательно вычисляем состояние ОДУ в дискретные моменты времени \( t_0, t+\Delta t, \dots , t+ N \Delta t \).
Методы с переменным размером шага
Для достижения более высокой точности по сравнению с RK4 нам необходимо использовать методы с переменным размером шага. Они динамически корректируют размер шага, чтобы удовлетворить заранее определенным абсолютным и относительным допускам ошибки интегрирования. Одним из примеров методов интегрирования с адаптивным размером шага является алгоритм Рунге-Кутта-Фельберга (RK45), который оценивает ошибку интегрирования при помощи RK5. Если ошибка больше допустимых значений — уменьшаем размер шага пропорционально соотношению допустимой ошибки интегрирования и реальной.
Однако одним из недостатков этих точных методов является увеличение вычислительного времени: нам заранее неизвестно количество шагов, необходимых для интегрирования ОДУ.
Нейронные дифференциальные уравнения
Теперь, когда мы понимаем, что такое ОДУ и как их решать, давайте рассмотрим процесс получения самого ОДУ. Для простых систем (как маятник) можно относительно просто вывести выражение векторного поля \( f \). Однако для более сложных систем (как роботы) вывод \( f \) требует обширных знаний в предметной области. Более того, даже если нам удастся получить векторное поле, расчет параметров \( \theta \) векторного поля, который соответствовал бы действительности, может оказаться нетривиальной задачей. Это связано с многочисленными предположениями и упрощениями, сделанными при формулировке физических законов (например, предположение об идеальном газе).
С другой стороны, при наличии у нас входных-выходных данных мы можем попытаться аппроксимировать векторное поле \( f \) каким-нибудь универсальным аппроксиматором. Один из популярных выборов для этого аппроксиматора — нейронная сеть. Это приводит к концепции нейронных обыкновенных дифференциальных уравнений (НОДУ):
\( \dot z(t) = NN\big(t, z(t), \theta_{NN}\big) \)
Да, получается, нейронные ОДУ — это не что иное, как ОДУ, где в качестве векторного поля выступает нейронная сеть.
Сравнение с ResNet
Применяем метод Эйлера для интегрирования НОДУ и получаем уравнение, близкое ResNet:
\( z(t+\Delta t) = z(t) + \Delta t \ NN(t, z(t), \theta_{NN}). \)
Единственное отличие — коэффициент масштабирования \( \Delta t \). НОДУ, интегрированное методом Эйлера, сводится к ResNet, если мы спрячем шаг интегрирования внутри нейронной сети \( NN(\cdot) \), то есть \( \tilde{NN}(\cdot) = \Delta t\ NN(\cdot) \) и
\( z(t+\Delta t) = z(t) +\tilde{NN}(t, z(t), \theta_{NN}) \)
Чуть раньше мы выяснили, что метод Эйлера уступает другим методам интегрирования, таким как RK4 или методам с переменным размером шага. Поэтому использование НОДУ с этими улучшенными интеграторами может привести к более точным и надежным моделям по сравнению с ResNet. Кроме того, НОДУ можно интерпретировать как бесконечномерные ResNet: ResNet дает решение только в дискретных точках и имеет низкую точность, а НОДУ с адаптивным шагом интегрирования имеет высокую точность (поля направленности здесь более гладкие как на Рис. 3) и может дать решение в любой точке вдоль траектории. Для достижения подобного результата с ResNet нам нужно сделать его очень глубоким.
Обучение и градиенты
Если вы задумались о тренировке НОДУ, но пока не понимаете, как можно тренировать бесконечномерные ResNet, не волнуйтесь! Вычисление градиентов для нейронных дифференциальных уравнений довольно простое и требует минимального объема памяти.
Начнем по порядку. Предположим, у нас есть набор данных \( {D} = \{x_i, y_i\} \), где \( x \) представляет собой входные данные, а \( y \) — целевые значения. Например, пусть \( x \) — угол и угловая скорость маятника, а \( y \) — тоже угол и угловая скорость маятника, но через 1 секунду. Наша цель в обучении нейронного дифференциального уравнения предсказывать целевые значения при помощи некоторой функции потерь \( {L}(\hat y, y) \). Мы можем выразить предсказанное значение \( \hat y_i \) как решение НОДУ:
\( \hat y_i := z(t_1) = \text{ODESolve}(z(t_0), NN, \theta_{NN}, t_0, t_1). \)
Здесь \( \text{ODESolve} \) может быть любым интегратором, \( y(t_0) := x_i \) представляет начальное состояние нашего нейронного дифференциального уравнения, а \( t_0 \) и \( t_1 \) — начальное и конечное время интегрирования соответственно. При использовании интеграторов с переменным размером шага мы выполняем множество промежуточных вычислений. Это приводит к большому вычислительному графу. Следовательно, обратное распространение градиентов через этот граф может быть вычислительно затратным и медленным.
На самом деле получение градиентов \( \hat y_i \) по отношению к состояниям \( z(t) \), параметрам \( \theta_{NN} \) и времени \( t \) довольно простое: нам нужно интегрировать дополнительное ОДУ (обратное по времени), которое вернет все эти градиенты. Эта мощная техника известна как adjoint method. Для понимания связи градиента НОДУ и решения сопряженного ОДУ можно почитать вот эти два хорошие объяснения: метод множителей Лагранжа и метод, описанный в оригинальной статье.
Мы в качестве примера рассмотрели задачу регрессии, но НОДУ также могут использоваться и для задачи классификации (здесь в конце интегрирования разные классы должны быть линейно разделимы).
Преимущества и недостатки
Согласно оригинальной статье, НОДУ имеют несколько преимуществ по сравнению с глубокими ResNet. Они позволяют более эффективно использовать память во время тренировки и требуют меньшего количества параметров. Кроме того, в случае НОДУ процесс обратного распространения ошибки вычислительно более эффективен.
НОДУ очень хорошо подходят для работы с временными рядами, особенно при неравномерной выборке набора данных (такое часто встречается на практике). Для работы с неравномерными данными и использования рекуррентных нейронных сетей требуется особая обработка. В отличие от рекуррентных сетей, НОДУ естественным образом обрабатывают неравномерные данные путем адаптации временной сетки интегратора.
Из личного опыта могу сказать, что тренировка нейронных дифференциальных уравнений занимает намного больше времени по сравнению с тренировкой рекуррентных нейронных сетей. Если имеете дело с равномерными данными — лучше начните с рекуррентных нейронных сетей 🙂
Библиотеки для работы с НОДУ
Сейчас есть много библиотек для работы с НОДУ. Для тренировки НОДУ с Pytorch можно использовать torchdiffeq или torchdyn. Если вы предпочитаете JAX — рекомендуем diffrax. Существуют и другие библиотеки, но они, как правило, более узконаправленные.
Домашнее задание
Чтобы поближе познакомиться с НОДУ и пощупать их на практике, предлагаем вам сделать следующее задание:
- Реализуйте интегратор RK4, он более точен для численного интегрирования дифференциальных уравнений по сравнению с методом Эйлера.
- Реализуйте правую часть ОДУ маятника, выбрав длину \( L=1 \), а \( g=9.81 \).
- Случайным образом выберите начальные условия \( y(t_0) \) для маятника и смоделируйте его поведение, начиная с выбранного начального условия. Используйте шаг времени \( \Delta t = 0.01 \) и проводите симуляцию в течение 1 секунды.
- Реализуйте adjoint method для НОДУ. Это позволит эффективно вычислять градиенты во время процесса обучения (этот пункт — пункт по желанию).
- Создайте модель НОДУ и обучите ее предсказывать траекторию маятника. На вход модели подается начальное состояние, а на выходе получаются состояния на каждом шаге интегратора.
В комментариях можно задать вопрос или поделиться успехами 🙂
