NeuralODE - Neural Ordinary Differential Equations

Neural Ordinary Differential Equations (NeuralODE) is a continuous normalizing flow model proposed by Chen et al., 2018. It works by evolving ordinary differential equations according to neural network defined functions. Since, ODEs under mild conditions are diffeomorphic as described by the Picard-Lindelof theorem, it means that NeuralODEs are invertible transformations. Therefore, they provide a way of representing only non-self-intersecting geometry. In this tutorial, we walk through the theory behind the NeuralODE representation and how to actually use their geodiff implementation in practice.

Theoretical Background¶

The Neural Ordinary Differential Equation (NeuralODE) architecture as proposed in Chen et al., 2018 is based on the idea of parameterizing a continuous-time dynamical system using a neural network. Rather than composing a finite sequence of discrete invertible transformations (as in NICE or RealNVP), a Neural ODE models a continuously invertible transformation obtained by integrating a learned vector field over time.

Formally, given an input $\mathbf{x}_0 \in \mathbb{R}^D$, the Neural ODE defines a time-dependent transformation $\mathbf{f}_{\theta}: \mathbb{R}^D \times [0, 1] \rightarrow \mathbb{R}^D$ through an ordinary differential equation (ODE):

where $f_{\theta}$ is a neural network with parameters $\theta$ that outputs a velocity (or rate of change) at each point $\mathbf{x}(t)$ and time $t$. Integrating this ODE from $t = 0$ to $t = 1$ yields the overall mapping:

Hence, the Neural ODE defines a transformation $\Phi_\theta : \mathbf{x}_0 \mapsto \mathbf{x}(1)$ that is implicitly determined by a continuous flow.

For representing non-self-intersecting geometry we want invertible transformations. For Neural ODEs, this property arises naturally from the theory of ordinary differential equations, and can be formally guaranteed by the Picard–Lindelöf theorem (also known as the Cauchy–Lipschitz theorem or the existence and uniqueness theorem).

The theorem states that if the vector field $f_\theta(\mathbf{x}, t)$ is Lipschitz continuous in $\mathbf{x}$ and continuous in $t$, then for every initial condition $\mathbf{x}_0$, there exists a unique differentiable solution $\mathbf{x}(t)$ to the ODE defined by Eq. (1). This guarantee implies that the flow generated by integrating $f_\theta$ forward in time is bijective and continuously invertible: each point $\mathbf{x}_0$ follows a unique trajectory through state space, and no two trajectories can intersect. Consequently, the mapping defined by the flow at time $t = 1$,

is a diffeomorphism - a smooth, invertible transformation with a smooth inverse given by integrating the same ODE backward in time.

In the context of geometry representation, we interpret $\mathbf{x}(t)$ as the evolving coordinates of points on a baseline closed manifold (for example, a unit circle or sphere). The ODE defines how these baseline points flow through a neural velocity field $f_\theta(\mathbf{x}, t)$ to form the final geometry at $t = 1$:

where the parameter samples $T \in [0, 1]^{d-1}$ define coordinates on the input parameter domain. The function $\mathrm{ClosedTransform}(T)$ maps these parameters to a closed baseline manifold - for example, the unit circle in 2D or the unit sphere in 3D. Integrating the neural velocity field $f_{\theta}$ from $t = 0$ to $t = 1$ then evolves this baseline into the final smooth and non-self-intersecting geometry $\mathbf{X}(1)$.

Implementation using geodiff¶

We now look at how geodiff allows us to easily use the NeuralODE parameterization to represent shapes.

The NeuralODE class initializer and its expected arguments are shown below:

    def __init__(
        self,
        geometry_dim: int,
        ode_net: nn.Module
    ) -> None:

To construct a NeuralODE object, we need to supply:

• Geometry Dimension 2 for 2D and 3 for 3D.

• Coupling Net a torch network module to be used as the function driving the ODE.

We use the torchdiffeq package created by Chen, n.d. to implement NeuralODEs for geometry representation.

Fitting Shapes¶

We will use a square as our target geometry. We then use the ChamferLoss to compute the geometric difference loss between the target shape and the shape represented by our NeuralODE parameterization. Since the implementation is written in PyTorch we can use the autograd capabilities to compute the gradients of the loss w.r.t. the NeuralODE network parameters and use an optimizer to modify our geometry.

We start by obtaining points on our shapes and normalizing them appropriately such that $x \in [0, 1]$. This is not needed for NeuralODE but is used to offer a comparison with the classical shape representation methods like Hicks-Henne and CST.

We now create a NeuralODE object by specifying the geometry dimension and a ODE function network.

We use the ChamferLoss provided by geodiff to compute a geometric loss between the target shape and the shape represented by the NeuralODE object. PyTorch’s autograd capabilities then allow us to compute gradients of the loss w.r.t. the NeuralODE network parameters and modify them to fit the target shape.

We can now visualize the shape represented by our NeuralODE object using its visualize method.