ODE_for_Neural_Optimal_Transport

Репозиторий проекта по применению Neural ODE к хадаче Оптимального Транспорта.

Команда

• Никита Гущин (руководитель)
• Давид Ли
• Петр Чижов
• Эльфат Сабитов

Введение

Задачу ОТ можно формализовать :

$$\max_f\min_T \int_{\mathcal{X}} c(x, T(x))d\mathbb{P}(x) + \int_{\mathcal{Y}} f(y) d\mathbb{Q}(y) - \int_{\mathcal{X}} f(T(x))d\mathbb{P}(x)$$

В рамках проекта команда в качестве оператора отображения использовалась модель нейронных ОДУ, траектории которых можно сделать прямыми, а значит сделать процесс отобржения обратимым.

$$T(x) = x_0 + \int\limits_0^1 V_{\theta}(x(t), t)dt$$

Ограничение на кривизну: $ \dfrac{d^2T(x)}{dt^2} = 0$

В ходе экспериментов, для которых использовалась библиотека Torchdyn и реализация Optimal Transport от команды Евгения Бурнаева, выяснилось, что регуляризация хоть и выпрямляет траектории, но мешает обучению, поэтому в дальнейшем была использована функция потерь:

$$\max_f\min_V { \mathbb{E} [\int_0^1||V(x(t), t)||^2 dt] + \int_{\mathcal{Y}} f(y) d\mathbb{Q}(y) - \int_\mathcal{X} f(T(x))d\mathbb{P}(x) }, $$

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

Структура

В папке src_not хранится реализация Neural Optimal Transport и дискриминатора для него - resnet. В models хранится реализация NeuralODE (ode.py), простого MLP для игрушечных датасетов (simple_mlp.py) и реализация Unet (models.py).

Основные эксперименты хранятся в папке notebooks, в них использовались датасеты colored mnist, bags and shoes.