Computational Visualization Center

Problem and Goals

• Models for learning unknown dynamics in dynamical systems: $\frac{d}{dt}y(t)=f(y(t),t)\backslash frac\{d\}\{dt\}y(t)\; =\; f(y(t),t)$
• Adapt the Neural ODE model with the adjoint method.
• To enforce the stability of the dynamics via Lypanunov constraints on the forward and adjoint equations. The forward and backwards (adjoint equations) for a Neural ODE are $\frac{d}{dt}z(t)=f_{\theta }(z(t),t)\backslash frac\{d\}\{dt\}\; z(t)\; =\; f\_\backslash theta(z(t),t)$ $\frac{d}{dt}a(t)=-a(t{)}^T\frac{\mathrm{\partial }f(z(t),t)}{\mathrm{\partial }z(t)}\backslash frac\{d\}\{dt\}\; a(t)\; =\; -a(t)^T\; \backslash frac\{\backslash partial\; f(z(t),t)\}\{\backslash partial\; z(t)\}$

The Lyapunov stability constraints for the two equation are: $\frac{d}{dt}W(a(t))=\mathrm{\nabla }W(a(t){)}^T{a}^{\mathrm{\prime }}(t)=-\mathrm{\nabla }W(a(t){)}^{T}a(t{)}^T\frac{\mathrm{\partial }f(z(t),t)}{\mathrm{\partial }z(t)}\backslash frac\{d\}\{dt\}W(a(t))\; =\; \backslash nabla\; W(a(t))^T\; a\text{'}(t)\; =\; -\; \backslash nabla\; W(a(t))^T\; a(t)^T\; \backslash frac\{\backslash partial\; f(z(t),t)\}\{\backslash partial\; z(t)\}$ $=-\mathrm{\nabla }W(a(t){)}^T{\frac{\mathrm{\partial }L}{\mathrm{\partial }z(t)}}^T\frac{\mathrm{\partial }f(z(t),t)}{\mathrm{\partial }z(t)}=-\mathrm{\nabla }W(a(t){)}^T{\frac{\mathrm{\partial }L}{\mathrm{\partial }z(t)}}^T\frac{\mathrm{\partial }}{\mathrm{\partial }z(t)}\frac{dz(t)}{dt}=-\backslash nabla\; W(a(t))^T\; \backslash frac\{\backslash partial\; L\}\{\backslash partial\; z(t)\}^T\; \backslash frac\{\backslash partial\; f(z(t),t)\}\{\backslash partial\; z(t)\}\; =\; -\backslash nabla\; W(a(t))^T\backslash frac\{\backslash partial\; L\}\{\backslash partial\; z(t)\}^T\; \backslash frac\{\backslash partial\}\{\backslash partial\; z(t)\}\; \backslash frac\{dz(t)\}\{dt\}$

Neural ODEs

A neural ODE is a machine learning model where the output is a solution to an ordinary differential equation (ODE) governed by a learned dynamics function, which is a neural network. $z(0)=g_1(x)(\mathrm{w}\mathrm{h}\mathrm{e}\mathrm{r}\mathrm{e}x\mathrm{i}\mathrm{s}\mathrm{t}\mathrm{h}\mathrm{e}\mathrm{i}\mathrm{n}\mathrm{p}\mathrm{u}\mathrm{t})z(0)\; =\; g\_1(x)\; \backslash qquad\; (\backslash mathrm\{where\}\backslash ;\; x\; \backslash ;\; \backslash mathrm\{is\backslash ;the\backslash ;input\})$ $\frac{d}{dt}z(t)=f_{\theta }(z(t),t),0\le t\le T\backslash frac\{d\}\{dt\}\; z(t)\; =\; f\_\backslash theta\; (z(t),t),\; \backslash quad\; 0\backslash leq\; t\; \backslash leq\; T$ $\mathrm{O}\mathrm{u}\mathrm{t}\mathrm{p}\mathrm{u}\mathrm{t}:y=g_2(z(T))\backslash mathrm\{Output\}:\; \backslash ;\; y\; =\; g\_2(z(T))$ Neural ODEs are in some way inspired by the similarity of residual neural networks and the Euler discretization of an ODE:$z_{n+1}=z_n+f_{{\theta }_n}(z_n)(\mathrm{r}\mathrm{e}\mathrm{s}\mathrm{i}\mathrm{d}\mathrm{u}\mathrm{a}\mathrm{l}\mathrm{n}\mathrm{e}\mathrm{t}\mathrm{w}\mathrm{o}\mathrm{r}\mathrm{k})z\_\{n+1\}\; =\; z\_n\; +\; f\_\{\backslash theta\_n\}(z\_n)\backslash qquad\backslash quad\backslash qquad\backslash ;\; (\backslash mathrm\{residual\backslash ;\; network\})$ $z_{n+1}=z_n+\mathrm{\Delta }t{f}_{\theta }(z_n,t_n)(\mathrm{E}\mathrm{u}\mathrm{l}\mathrm{e}\mathrm{r}\mathrm{d}\mathrm{i}\mathrm{s}\mathrm{c}\mathrm{r}\mathrm{e}\mathrm{t}\mathrm{i}\mathrm{z}\mathrm{a}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n})z\_\{n+1\}\; =\; z\_n\; +\; \backslash Delta\; t\; f\_\backslash theta(z\_n,\; t\_n)\backslash qquad\; (\backslash mathrm\{Euler\backslash ;\; discretization\})$

by using neural ODEs, more advanced ODE solvers such as Runge-Kutta solvers and adaptive solvers like DOPRI.

Neural ODE adjoints

In order to train the Neural ODE, the gradient of the loss with respect to the parameter, $\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }\theta }\backslash frac\{\backslash partial\; \backslash mathcal\{L\}\}\{\backslash partial\; \backslash theta\}$, must be computed. The typical method of backpropagation used in neural networks can have high computation/memory cost when using ODE solvers. An alternative method for neural ODE is using the adjoint ODE. The adjoint ODE is solved backwards in time for the gradients. The adjoint state is defined by $a(t)=\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }z(t)},a(t)\; =\; \backslash frac\{\backslash partial\; \backslash mathcal\{L\}\}\{\backslash partial\; z(t)\},$ and it solves the adjoint ODE, $\frac{d}{dt}a(t)=-a(t{)}^T\frac{\mathrm{\partial }{f}_{\theta }(z(t),t)}{\mathrm{\partial }z(t)},T\ge t\ge 0.\backslash frac\{d\}\{dt\}a(t)\; =\; -a(t)^T\backslash frac\{\backslash partial\; f\_\backslash theta(z(t),t)\}\{\backslash partial\; z(t)\},\; \backslash qquad\; T\backslash geq\; t\; \backslash geq\; 0.$

From the adjoint, the loss gradient can be computed by $\frac{\mathrm{\partial }\mathcal{L}}{\mathrm{\partial }\theta }=-{\int }_0^{T}a(t{)}^T\frac{\mathrm{\partial }{f}_{\theta }(z(t),t)}{\mathrm{\partial }\theta }\mathrm{.}\backslash frac\{\backslash partial\; \backslash mathcal\{L\}\}\{\backslash partial\; \backslash theta\}\; =\; -\; \backslash int\_0^T\; a(t)^T\; \backslash frac\{\backslash partial\; f\_\backslash theta(z(t),t)\}\{\backslash partial\; \backslash theta\}.$

Stability Criteria

Lyapunov Stability

An ODE is Lyapunov stable if .

Asymptotic Stability

An ODE is asymptoticly stable if ${\lim }_{z\to \mathrm{\infty }}z(t)=0\mathrm{f}\mathrm{o}\mathrm{r}\mathrm{a}\mathrm{n}\mathrm{y}\mathrm{i}\mathrm{n}\mathrm{i}\mathrm{t}\mathrm{i}\mathrm{a}\mathrm{l}\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{d}\mathrm{i}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}z(0)\mathrm{.}\backslash lim\_\{z\backslash rightarrow\; \backslash infty\}\; z(t)=\; 0\; \backslash ;\; \backslash mathrm\{for\backslash ;\; any\backslash ;\; initial\backslash ;\; condition\backslash ;\}\; z(0).$

Exponential Stability

An ODE is exponentially stable with equilibrium $z_{e}z\_e$ if $\mathrm{\parallel }z(t)\mathrm{\parallel }\le C{e}^{-\alpha t}\mathrm{\parallel }z(0)\mathrm{\parallel }\mathrm{f}\mathrm{o}\mathrm{r}\mathrm{s}\mathrm{o}\mathrm{m}\mathrm{e}C,\alpha >0\backslash |z(t)\backslash |\backslash leq\; C\; e^\{-\backslash alpha\; t\}\backslash |z(0)\backslash |\; \backslash ;\; \backslash mathrm\{for\backslash ;\; some\backslash ;\}\; C,\backslash alpha\; >\; 0$

Lyapunov Functions

A Lyapunov function is a function $V:{\mathbb{R}}^d\to \mathbb{R}V:\backslash mathbb\{R\}^d\; \backslash rightarrow\; \backslash mathbb\{R\}$ such that $V(z)>0V(z)>0$ for $z\mathrm{\ne }0z\backslash neq0$ and $V(0)=0V(0)=0$. An ODE is asymptoticaly stable if there exists a Lyapunov function $VV$ such that

An ODE is exponentially stable if there exists a Lyapunov function $VV$ such that $\frac{d}{dt}V(z(t))=\mathrm{\nabla }V(z(t){)}^T{z}^{\mathrm{\prime }}(t)\le -\alpha V(z(t)),\backslash frac\{d\}\{dt\}V(z(t))\; =\; \backslash nabla\; V(z(t))^T\; z\text{'}(t)\backslash leq-\backslash alpha\; V(z(t)),$ $\frac{\frac{d}{dt}V(z(t))}{V(z(t))}\le -\alpha \backslash frac\{\backslash frac\{d\}\{dt\}V(z(t))\}\{V(z(t))\}\backslash leq\; -\backslash alpha$ $\frac{d}{dt}log(V(z(t)))\le -\alpha \backslash frac\{d\}\{dt\}\; log(V(z(t)))\backslash leq\; -\backslash alpha$ $log(V(z(t)))-log(V(z(0)))\le -\alpha tlog(V(z(t)))\; -\; log(V(z(0)))\; \backslash leq\; -\backslash alpha\; t$ $log(V(z(t))\le -\alpha t+log(V(z(0)))log(V(z(t))\; \backslash leq\; -\backslash alpha\; t\; +\; log(V(z(0)))$ $V(z(t))\le e^{-\alpha t}V(z(0))V(z(t))\backslash leq\; e^\{-\backslash alpha\; t\}\; V(z(0))$

To show that exponential stability implies the above Lyapunov condition: Assume that $\mathrm{\parallel }z(t)\mathrm{\parallel }\le C{e}^{-\alpha t}\mathrm{\parallel }z(0)\mathrm{\parallel }\backslash |z(t)\backslash |\backslash leq\; Ce^\{-\backslash alpha\; t\}\; \backslash |z(0)\backslash |$.

Related Work on Stablizing Neural ODEs

Learning Stable Deep Dynamics Models

This paper describes a method to learn the dynamics function and guarentee that it is stable. In order to guarentee it is stable, in addition to learning the dynamics function and Lyapunov function is also learned. The dynamics function is projected onto the half-space that satisfying the Lyapunov stability condition, so that the projected dynamics function is guaranteed to be stable. The projection is given by: $f(z)=\widehat{f}(z)-\mathrm{\nabla }V(z)\frac{RELU(\mathrm{\nabla }V(z{)}^T\widehat{f}(z)+\alpha V(z))}{\mathrm{\parallel }\mathrm{\nabla }V(z){\mathrm{\parallel }}^2}f(z)\; =\; \backslash hat\; f(z)\; -\; \backslash nabla\; V(z)\; \backslash frac\{RELU(\backslash nabla\; V(z)^T\; \backslash hat\; f(z)\; +\; \backslash alpha\; V(z))\}\{\backslash |\backslash nabla\; V(z)\backslash |^2\}$ It is stable since $\frac{d}{dt}V(z(t))=\mathrm{\nabla }V(z{)}^{T}f(z)=\mathrm{\nabla }V(z{)}^T\mathrm{\nabla }\widehat{f}(z)\backslash frac\{d\}\{dt\}V(z(t))\; =\; \backslash nabla\; V(z)^T\; f(z)\; =\; \backslash nabla\; V(z)^T\; \backslash nabla\; \backslash hat\; f(z)$ $-\mathrm{\nabla }V(x{)}^T\mathrm{\nabla }V(z)\frac{RELU(\mathrm{\nabla }V(z{)}^T\widehat{f}(z)+\alpha V(z))}{\mathrm{\parallel }\mathrm{\nabla }V(z){\mathrm{\parallel }}^2}-\; \backslash nabla\; V(x)^T\; \backslash nabla\; V(z)\; \backslash frac\{RELU(\backslash nabla\; V(z)^T\; \backslash hat\; f(z)\; +\; \backslash alpha\; V(z))\}\{\backslash |\backslash nabla\; V(z)\backslash |^2\}$ $=\mathrm{\nabla }V(z{)}^T\mathrm{\nabla }\widehat{f}(z)-RELU(\mathrm{\nabla }V(z{)}^T\widehat{f}(z)+\alpha V(z))=\; \backslash nabla\; V(z)^T\; \backslash nabla\; \backslash hat\; f(z)\; -\; RELU(\backslash nabla\; V(z)^T\; \backslash hat\; f(z)\; +\; \backslash alpha\; V(z))$

If $\mathrm{\nabla }V(z{)}^T\widehat{f}(z)+\alpha V(z)\le 0\backslash nabla\; V(z)^T\; \backslash hat\; f(z)\; +\; \backslash alpha\; V(z)\; \backslash leq\; 0$, then $\frac{d}{dt}V(z(t))=\mathrm{\nabla }V(z{)}^T\mathrm{\nabla }\widehat{f}(z)\le -\alpha V(z)\mathrm{.}\backslash frac\{d\}\{dt\}V(z(t))\; =\; \backslash nabla\; V(z)^T\; \backslash nabla\; \backslash hat\; f(z)\; \backslash leq\; -\backslash alpha\; V(z).$

If $\mathrm{\nabla }V(z{)}^T\widehat{f}(z)+\alpha V(z)>0\backslash nabla\; V(z)^T\; \backslash hat\; f(z)\; +\; \backslash alpha\; V(z)\; >\; 0$, then $\frac{d}{dt}V(z(t))=\mathrm{\nabla }V(z{)}^T\mathrm{\nabla }\widehat{f}(z)-(\mathrm{\nabla }V(z{)}^T\widehat{f}(z)+\alpha V(z))=-\alpha V(z)\mathrm{.}\backslash frac\{d\}\{dt\}V(z(t))\; =\; \backslash nabla\; V(z)^T\; \backslash nabla\; \backslash hat\; f(z)\; -\; (\backslash nabla\; V(z)^T\; \backslash hat\; f(z)\; +\; \backslash alpha\; V(z))\; =\; -\backslash alpha\; V(z).$

So for either case, it is exponentially stable.

Adjoint equation with Lyapunov projection

Let us consider the adjoint equation,

The projection for the adjoint would be:

LyaNet

LyaNet is a Neural ODE with a modified loss function so that zero loss guarantees the exponential stability of the forward ODE. Consider the Neural ODE taking input $xx$ to output $yy$: $z(0)=g_{\theta }(x)z(0)=g\_\backslash theta(x)$ $y=h_{\theta }(z(1))y\; =\; h\_\backslash theta(z(1))$

Let the standard loss function be $\mathrm{\ell }(y,\widehat{y})\backslash ell(y,\backslash hat\; y)$. The Lyapunov function is taken to be $V(z(t))=\mathrm{\ell }(h_{\theta }(z(t)),\widehat{y})V(z(t))\; =\; \backslash ell(h\_\backslash theta(z(t)),\backslash hat\; y)$. The new loss function is $\mathcal{V}(\theta )={\int }_0^{1}RELU(\frac{d}{dt}V(z(t))+\alpha V(z(t)))dt\mathrm{.}\backslash mathcal\{V\}(\backslash theta)\; =\; \backslash int\_0^1\; RELU\; \backslash left(\backslash frac\{d\}\{dt\}V(z(t))\; +\backslash alpha\; V(z(t))\; \backslash right)\; dt.$

Suppose there exists optimal parameters ${\theta }^{\ast }\backslash theta^*$ such that $\mathcal{V}({\theta }^{\ast })=0\backslash mathcal\{V\}(\backslash theta^*)=0$. Then $\frac{d}{dt}V(z(t))+\alpha V(z(t))\le 0,\backslash frac\{d\}\{dt\}V(z(t))\; +\backslash alpha\; V(z(t))\backslash leq\; 0,$

so it is exponentially stable, and also: $V(z(t))\le e^{-\alpha t}V(z(0))V(z(t))\backslash leq\; e^\{-\backslash alpha\; t\}\; V(z(0))$ $\mathrm{\ell }(h_{\theta }(z(t)),\widehat{y})\le e^{-\alpha t}\mathrm{\ell }(h_{\theta }(z(0)),\widehat{y})\mathrm{.}\backslash ell(h\_\backslash theta(z(t)),\; \backslash hat\; y)\; \backslash leq\; e^\{-\backslash alpha\; t\}\; \backslash ell(h\_\backslash theta(z(0)),\; \backslash hat\; y).$

The Lyapunov loss function is an integral that must be approximated. There are two proposed methods for approximating it. One is to use a Monte Carlo method of sampling time points $t_{i}t\_i$ and states $z_{i}z\_i$ and averaging the value at different samples. The other is path integration. The ODE is numerically integrated for an approximate solution at a number of time steps. Then discrete approximations of the derivative and integral in the Lyapunov loss function are used.

Possible approach for adjoint

Consider using two regularizations, one for the forward and one for the adjoint: ${\mathcal{L}}_{\mathrm{L}\mathrm{y}\mathrm{a}\mathrm{p}}={\int }_0^1\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(\frac{d}{dt}V(z(t))+\alpha V(z(t)))dt\backslash mathcal\{L\}\_\backslash mathrm\{Lyap\}\; =\; \backslash int\_0^1\; \backslash mathrm\{ReLU\}\; \backslash left(\; \backslash frac\{d\}\{dt\}\; V(z(t))\; +\; \backslash alpha\; V(z(t))\; \backslash right)\; dt$ For adjoint: ${\mathcal{L}}_{\mathrm{a}\mathrm{d}\mathrm{j}}={\int }_0^1\mathrm{R}\mathrm{e}\mathrm{L}\mathrm{U}(\frac{d}{dt}W(a(t))+\alpha W(a(t)))dt\backslash mathcal\{L\}\_\backslash mathrm\{adj\}\; =\; \backslash int\_0^1\; \backslash mathrm\{ReLU\}\; \backslash left(\; \backslash frac\{d\}\{dt\}\; W(a(t))\; +\; \backslash alpha\; W(a(t))\; \backslash right)\; dt$ The problem will be determining how

Dueling Neural ODE Inner and Outer Loop

The inner loop contains the forward and adjoint ODE solves with stability constraints. The outer loop focuses on the filtering of gradients, i.e. selection of gradients with highest impact on training.

• To enforce stability in the gradient calculations used for updating in the training process.

Some Results

Comparison of standard NODE, LyaNet, and PINN for solving spring equation (NYU data) over a domain of initial conditions

Note: PINN is trained for 200000 epochs and the final test loss is shown for reference.

Results of training LyaNet to solve spring equation (NYU data). Two types of loss are computed, MSE and Lyapunov, for each the training and testing datasets (top figure). The differences between the training and testing loss (a measure of generalization) is plotted in the bottom figure.

Results of previous model but with controller that removes some gradients

References

• Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, David Duvenaud. Neural Ordinary Differential Equations. 2019. https://arxiv.org/abs/1806.07366
• Gaurav Manek and J. Zico Kolter. Learning stable deep dynamics models. Neural information processing systems. 2019. https://arxiv.org/pdf/2001.06116 https://github.com/locuslab/stable_dynamics
• Ivan Dario Jimenez Rodriguez, Aaron D. Ames, Yisong Yue. LyaNet: A Lyapunov dramework for training neural ODEs. 2022. https://arxiv.org/abs/2202.02526
• Sanchez-Lengeling and Aspuru-Guzik. Inverse molecular design using machine learning: generative models for matter engineering. Science. 2018.
• Advanced Generative Adversarial Autoencoder Model for de Novo Generation of New Molecules with Desired Molecular Properties in Silico. Mol. Pharmaceutics. 2017.
• Bajaj, et al. An Efficient Higher-Order Fast Multipole Boundary Element Solution for Poisson–Boltzmann-Based Molecular Electrostatics. SIAM J Sci Comput. 2011.
• Lelièvre and Stoltz. Partial differential equations and stochastic methods in molecular dynamics. Actua Numerica. 2016.
• J.T. Oden, Ludovic Chamoin, Serge Prudhomme. A stochastic coupling method for atomic-to-continuum Monte-Carlo simulations. Comput. Methods Appl. Mech. Engrg. 2008.