Дифференциальные уравнения и нейронные сети вместе? Не может быть или может... Neural ODE – подход в глубоком обучении, объединяющий идеи нейронных сетей и обыкновенных дифференциальных уравнений. Выглядит пугающе, давайте проверим!

Как уже стало понятно, в данной статье речь пойдет о ДУ (вам предстоит много раз встретить их упоминание) и нейронных сетях. Но перед тем как приступать непосредственно к статье, начнем с азов, чтобы не было паники в голове.

Для людей, которых не обескураживают словосочетания: дифференциальные уравнения ~~(опять они),~~ метод Эйлера, эволюционные алгоритмы - можете смело переходить к идеи Neural ODE ~~(а может вам проще прочитать оригинальную статью)~~

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

Что такое дифференциальные уравнения?

Начнем с базы!

Предполагается, что читатели знакомы с понятием производной. Дифференциальное уравнение от одной переменной вида F(x, y, y’, y’’, … y^(n)) = 0 называется обыкновенным дифференциальным уравнением (ODE), где у-неизвестная функция, а y^n производная функции n порядка.

Переменная х интерпретируется как время, а у - величина, которая меняется со временем. То есть ДУ (дифференциальные уравнения) описывают динамические системы, которые меняются с течением времени. Применение ДУ множество, начиная от описания движением планет и заканчивая описанием распространения эпидемии.

Можно накидать много терминов и сложных уравнений, но постараемся обойтись без этого. Не оставлю без внимания лишь одно, чтобы в дальнейшем не было недопонимания. Задача Коши - ДУ y’=f(x, y), в котором требуется найти частичное решение исходя из начальных условий (например, у(х)=0).

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

Необходимо найти функцию y (не ее производную!). Вспомним такие слова как первообразная и интеграл. Кто не знает, советую ознакомится. Берем интерал с двух сторон и получаем:

где С - это некоторая постоянная. Мы нашли общее решение ДУ. Если бы у нас стояла задача Коши (было задано начальное условие), то можно было бы найти знание константы С, таким образом получая частное решение.

Метод Эйлера

Основная идея метода Эйлера заключается в приближенном представлении непрерывной траектории системы путем разбиения времени на небольшие интервалы. На каждом интервале вычисляется производная состояния системы, и эта производная используется для обновления состояния системы.

Для обыкновенного дифференциального уравнения первого порядка формула Эйлеровой дискретизации имеет вид:

Алгоритм представлен на рисунке ниже.

Простые блок-схемы конечно хорошо, но хотелось бы увидеть пример. Сказано, сделано!

Пусть: У нас есть уравнение с условием y(0)=1:

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

Эволюционный метод решения ОДУ

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

Теперь давайте свяжем их: эволюционный метод для решения ОДУ может использовать метод Эйлера, где каждый кандидат в популяции представляет собой начальные условия и шагает вперед по времени. Процесс эволюции улучшает популяцию, приближаясь к лучшему решению ОДУ.

Пример:

У нас есть исходное уравнение и начальное условие, что y(0) = 1:

Мы уже знаем, как решать уравнения такого вида (~~надеюсь~~), берем интеграл с обеих сторон, получаем.

На сцену выходит эволюционный алгоритм. Предполагаем, что у нас есть начальная популяция (начинаем со случайных значений):

Кандидат 1:

a=0.5; b=2; C=2

Кандидат 2:

a=1; b=-1; C=-1

Кандидат 3:

a=-0,2; b=0; C=0

Рассчитываем, насколько близко решение каждого кандидата к начальным условиям (Конечно же, используя метод Эйлера!).

	Кандидат 1	Кандидат 2	Кандидат 3
y =	2	-1	0
\|y(0) - y\|	1	2	1

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

Мутация:

Скрещивание:

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

Neural ODE

Теперь мы готовы перейти к самому ~~сложному~~ вкусному!

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

В Neural Ordinary Differential Equations (Neural ODEs) обратный проход (backpropagation) использует принципы дифференцирования, но в этом случае, это основывается на дифференциальных уравнениях и методе сопряженных состояний (adjoint method). Обычные нейронные сети используют стандартные алгоритмы (градиентный спуск и его вариации) для обратного распространения градиентов. Важной особенностью Neural ODEs является то, что обратный проход протекает через процесс решения дифференциальных уравнений, что может быть более эффективным с точки зрения вычислений.

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

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

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

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

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

Достоинства такого подхода:

Память. Нам не нужно хранить промежуточные величины, которые были получены при прямом проходе (forward pass), что существенно экономит память.
Адаптивность. Метод Эйлера самый простой для решения ODE, существуют множество других способов решений ODE. Современные методы дают гарантии относительно того, как быстро или медленно может возрастать ошибка аппроксимации (приближения) данных. (Ошибка аппроксимации - это разница между точными данными и их приближенным значением, полученным при использовании методов обработки или решения). Они могут адаптироваться на лету, изменяя свою стратегию при изменении условий.
Биективность отображения. Благодаря непрерывности мы можем отображать данные в пространство более высокой размерности и обратно, сохраняя информацию.
Непрерывность. Мы не зациклены на дискретности и можем получить данные в любой момент времени.

Backpropagation в ODE

Авторы рассматривают модель (ODE solver) как черный ящик и для вычисления градиентов используют метод чувствительности сопряжённых уравнений (adjoint sensitivity method)

Данный подход линейно масштабируется в зависимости от размера входных данных.

Функция потерь можно представить следующим уравнением (допустим у нас два состояние t0 и t1):

Откуда взялся интеграл? Интеграл представляет собой решение дифференциального уравнения, описывающего эволюцию скрытого состояния z(t) в течение времени t. То есть то что написано в формуле выше выглядит страшно, но по факту это y_pred.

Для оптимизации L нужно знать градиенты. Нужно определить, как значение L зависит от скрытого состояния z(t) в каждый момент времени. Это значение называется сопряжённым состоянием (adjoint).

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

Его динамика задается другим дифференциальными уравнением, которое можно считать непрерывным аналогом дифференцирования сложной функции (chain rule). Это уравнение описывает как меняется сопряженное состояние со временем.

Получается: для обратного прохода нам надо знать состояние z(t) в каждый момент времени при вычислении производной при обратном проходе. Но вместо того, чтобы хранить ее, мы можем просто ее заново посчитать с помощью ODE Solver вместе с adjoint (то есть вызываем ODE Solver два раза, он быстрый, поэтому не затратно по времени). Таким образом, мы идем с финального состояния системы z(T) до начального состояния системы z(t0) на каждом этапе вызывая ODE Solver.

Neural ODEs. Supervised learning

Авторы статьи используют неявный метод Рунге-Кутты с адаптивным размером шага. В экспериментах была использована небольшая сеть c residual блоками, которая уменьшает размер входных данных дважды, затем применяет 6 стандартных residual блоков.

У авторов было несколько вариаций архитектур:

Вариант сети ODE-Net заменяет эти блоки на модуль ODESolve.
Сеть с аналогичной архитектурой, но градиенты обратно распространяются через интегратор Рунге-Кутты, сеть назвали RK-Net.

В таблице представлены результаты тестирования: ошибка на тестовой выборке, количество параметров, объем занимаемой памяти время. L обозначает количество слоев в ResNet, а L с волной - количество вычислений, которое Solver ОДУ нужно за один проход. В итоге авторы сделали вывод, что ODE-Nets и RK-Nets могут достигнуть примерно такой же производительности, как у ResNet.

Также авторы выделили несколько свойств:

Изменение уровня численной ошибки влечет за собой изменение количества шагов в процессе forward pass.
Время, затраченное на forward pass, пропорционально количеству вызова функций.
Количество вычислений в процессе backpropagation составляет примерно половину от количества в forward pass. Метод сопряженного градиента может быть более вычислительно эффективным, чем прямое распространение градиента через ODESolver.
По мере увеличения эпох обучения ODE-Net, требуется все больше вычислений (с уменьшением шага), что может указывать на адаптацию к увеличению сложности модели.

Пример

Советую посмотреть на репозиторий авторов статьи, в папке examples можно найти несколько чудесных примеров. Я же ограничусь примитивным, чтобы было понятно как строить архитектуры и применять метод.

Допустим у нас есть задача классификации (возьмем датасет MNIST). Построим примитивную модель (заметьте, тут нет и речи о хорошей модели).

Главное не забыть импортировать библиотеку:

from torchdiffeq import odeint

Построим сетку на основе линейных слоев ~~(да-да, я знаю, что есть свертки)~~. Предположим, мы хотим сжать картинку до некоторого скрытого пространства, получив скрытые фичи.

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchdiffeq import odeint

class ODEBlock(nn.Module):

    def __init__(self, odefunc):
        super(ODEBlock, self).__init__()
        self.odefunc = odefunc
        self.integration_time = torch.tensor([0, 1]).float()

    def forward(self, x):
        self.integration_time = self.integration_time.type_as(x)
        out = odeint(self.odefunc, x, self.integration_time)
        return out[1]

class ODEFunc(nn.Module):

    def __init__(self, input_dim=28):
        super(ODEFunc, self).__init__()

        self.net = nn.Sequential(
            nn.Flatten(start_dim=1),
            nn.Linear(input_dim*input_dim, 64),
            nn.ReLU(),
            nn.Linear(64, input_dim*input_dim),
        )

    def forward(self, t, y):
        b, c, h, w = y.shape
        output = self.net(y**3)
        output = output.view(output.size(0), c, h, w) 
        return output

Отлично! Наша супер крутая сетка почти готова, осталось построить модель:

feature_layers = [ODEBlock(ODEFunc()) for _ in range(2)]
fc_layers = [nn.ReLU(), 
             nn.AdaptiveAvgPool2d((1, 1)), 
             nn.Flatten(), 
             nn.Linear(1, 10)
            ]
model = nn.Sequential(*feature_layers, *fc_layers)

Теперь мы можем запускать обучение. Это оставлю за кадром.

Hidden text

Сетка училась долго и результат accuracy=0.3. Реализация со свертками мне понравилась гораздо больше.

Надеюсь, вы узнали что-то новое. Если нет, извините :(

Литература

Классная статья с примером реализации Neural ODE

Neural ODE: встреча с дифференциальными Уравнениями