Neural ODEs

Summary

Neural ODEs (Neural Ordinary Differential Equations) represent one of the most elegant advances in recent years in the field of deep learning. Introduced by Chen et al. in 2018 in the paper “Neural Ordinary Differential Equations,” they propose a fundamental paradigm shift: instead of stacking discrete neural layers as traditional neural networks do, Neural ODEs model the transformation of data as a continuous process governed by an ordinary differential equation. This approach offers constant memory during training, adaptive computational complexity evaluation, and a natural ability to handle irregularly spaced time series. In this complete guide, we will explore the mathematical foundations, the intuition behind this concept, practical implementation with PyTorch and the torchdiffeq library, as well as the concrete use cases that make Neural ODEs an essential tool for any practitioner of modern machine learning.

Mathematical Principle

From Discrete Layers to Continuous Dynamics

To understand Neural ODEs, we must first examine how a classic deep neural network works, particularly a residual network (ResNet). In a ResNet, each layer computes:

h_{t+1} = h_t + f(h_t, θ_t)

where $h_t$ is the hidden state at layer $t$, $f$ is a differentiable function (typically a small neural network), and $θ_t$ are the parameters specific to that layer. If we look at this equation carefully, we recognize the form of a Euler discretization of an ordinary differential equation, with a time step $\Delta t = 1$.

Neural ODEs push this observation to its logical extreme. Instead of considering discrete layers with parameters $θ_t$ that change at each step, they define a continuous dynamics where the derivative of the hidden state with respect to time (or network depth) is parameterized by a single shared neural network:

$$\frac{dh(t)}{dt} = f(h(t), t, θ)$$

Here, $f$ is a single neural network whose parameters $θ$ are shared across the entire “time” $t$. The output of the network at time $t_1$, given an input at $t_0$, is obtained by numerical integration of this differential equation:

$$h(t_1) = h(t_0) + \int_{t_0}^{t_1} f(h(t), t, θ) \, dt$$

In practice, this integral is solved by a numerical ODE solver (such as Runge-Kutta of order 4-5, called “dopri5”), which automatically adapts the integration step size to achieve a given precision.

The Adjoint Method for Gradient Computation

The real tour de force of Neural ODEs lies in how they compute gradients during backpropagation. In a classic deep network, backpropagation requires storing all intermediate activations, which costs memory proportional to the number of layers: $O(L)$.

For Neural ODEs, Chen et al. use the adjoint sensitivity method. The idea is to define an adjoint state:

$$a(t) = \frac{\partial L}{\partial h(t)}$$

where $L$ is the loss function. This adjoint state satisfies its own differential equation:

$$\frac{da(t)}{dt} = -a(t)^\top \frac{\partial f(h(t), t, θ)}{\partial h}$$

To obtain the gradients with respect to parameters $θ$, we solve an additional ODE backward:

$$\frac{dL}{dθ} = \int_{t_1}^{t_0} a(t)^\top \frac{\partial f(h(t), t, θ)}{\partial θ} \, dt$$

This approach has a remarkable consequence: the memory used is constant, $O(1)$, independent of the network complexity or the number of solver evaluations. There is no need to store intermediate activations; the system simply needs to be reintegrated backward. This is a decisive advantage for very deep models or long time series.

Solver Tolerance: The Precision/Speed Trade-off

Unlike traditional neural networks where the number of layers is fixed in advance, Neural ODEs let the ODE solver decide how many evaluations are needed to achieve a given precision. This precision is controlled by two hyperparameters:

• atol (absolute tolerance): the maximum error accepted in absolute value at each integration step.
• rtol (relative tolerance): the maximum error accepted proportionally to the magnitude of the current state.

Lower tolerances (e.g., atol=1e-7, rtol=1e-7) force the solver to take more integration steps, increasing precision but also computation time. Conversely, higher tolerances (atol=1e-3, rtol=1e-3) give faster computation at the cost of a coarser approximation. This mechanism offers an adaptive trade-off unique in deep learning.

Intuition

To truly grasp what Neural ODEs are, a physical analogy is particularly illuminating.

Imagine a classic neural network as a staircase. Each step corresponds to a layer. To go from the input to the output at the bottom, you must place your foot on each step, one at a time. You cannot skip a step, and each step takes you to a well-defined discrete level. If you want more precision in your descent, you must build more steps — which requires more memory to remember each step.

Now imagine a continuous ramp instead of a staircase. On a ramp, your descent is smooth, without interruptions. You slide along a slope whose incline can change gradually. There are no discrete levels: you are in a continuum. This is exactly what a Neural ODE does: the transformation of data is like a slide along a ramp, where the shape of the ramp is defined by the network $f(h(t), t, θ)$.

And to know how to modify the shape of this ramp (i.e., how to adjust the parameters $θ$ to minimize the error), we use the adjoint method. The intuition here is that, rather than going back up the staircase step by step remembering each step (which would cost a lot of memory), we use the same ODE solver to go back up the slope smoothly in reverse. The solver works just as well forward as backward, eliminating the need to store intermediate states.

This analogy explains why Neural ODEs are particularly well-suited to continuous data: time series with observations at irregular intervals, physical trajectories, dynamic processes where time flows naturally and non-discretely. The choice of solver and its tolerances is like choosing how finely we explore the slope — we can use large steps for a fast descent, or small steps for maximum precision.

Python Implementation with torchdiffeq

The torchdiffeq library is the reference implementation for Neural ODEs in the PyTorch ecosystem. It provides the key function odeint, which integrates a differential equation whose right-hand side function is a neural network.

Installation

pip install torchdiffeq

Basic Neural ODE Architecture

Let’s start by defining a network that will serve as the dynamics function $f(h(t), t, θ)$:

import torch
import torch.nn as nn
from torchdiffeq import odeint

class ODEFunction(nn.Module):
    """
    Dynamics function f(h(t), t, theta) parameterized by a neural network.
    This is the heart of the Neural ODE.
    """
    def __init__(self, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(hidden_dim, 64),
            nn.Tanh(),
            nn.Linear(64, 64),
            nn.Tanh(),
            nn.Linear(64, hidden_dim)
        )

    def forward(self, t, h):
        """
        Compute dh/dt = f(h, t, theta).
        Time t is provided by the solver but can be ignored or used.
        """
        return self.net(h)


class NeuralODE(nn.Module):
    """
    Complete Neural ODE: integrates the ODEFunction over an interval [t0, t1].
    """
    def __init__(self, hidden_dim, method="dopri5", atol=1e-7, rtol=1e-7):
        super().__init__()
        self.ode_func = ODEFunction(hidden_dim)
        self.method = method
        self.atol = atol
        self.rtol = rtol

    def forward(self, h0, t_span):
        """
        Integrates h(t) from t_span[0] to t_span[-1].
        Returns the complete trajectory if len(t_span) > 2,
        or simply the final state if t_span = [t0, t1].
        """
        h = odeint(
            self.ode_func,
            h0,
            t_span,
            method=self.method,
            atol=self.atol,
            rtol=self.rtol
        )
        return h

ResNet vs Neural ODE Comparison on MNIST

To illustrate the concrete difference, let’s compare a classic ResNet and a Neural ODE on MNIST classification:

import torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# MNIST data preparation
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST("./data", train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)


class ResNet_MNIST(nn.Module):
    """Classic ResNet with 6 residual layers."""
    def __init__(self):
        super().__init__()
        self.fc_in = nn.Linear(28 * 28, 64)
        self.res_layers = nn.ModuleList([
            nn.Sequential(
                nn.Linear(64, 64),
                nn.ReLU(),
                nn.Linear(64, 64)
            ) for _ in range(6)
        ])
        self.fc_out = nn.Linear(64, 10)

    def forward(self, x):
        h = x.view(-1, 28 * 28)
        h = self.fc_in(h)
        for layer in self.res_layers:
            h = h + layer(h)  # residual connection
        return self.fc_out(h)


class NeuralODE_MNIST(nn.Module):
    """Neural ODE applied to MNIST."""
    def __init__(self):
        super().__init__()
        self.fc_in = nn.Linear(28 * 28, 64)
        self.ode_func = ODEFunction(hidden_dim=64)
        self.fc_out = nn.Linear(64, 10)

    def forward(self, x):
        h = x.view(-1, 28 * 28)
        h = self.fc_in(h)
        # Integration from t=0 to t=1
        t_span = torch.tensor([0.0, 1.0])
        h = odeint(self.ode_func, h, t_span, method="dopri5",
                   atol=1e-7, rtol=1e-7)[-1]  # we take the final state
        return self.fc_out(h)


# Memory comparison
def compare_memory():
    """Compare ResNet vs Neural ODE memory consumption."""
    resnet = ResNet_MNIST()
    neural_ode = NeuralODE_MNIST()

    resnet_params = sum(p.numel() for p in resnet.parameters())
    node_params = sum(p.numel() for p in neural_ode.parameters())

    print("ResNet parameters: " + str(resnet_params))
    print("Neural ODE parameters: " + str(node_params))
    # Neural ODE uses far fewer parameters because weights are shared
    # through time instead of being replicated per layer.

The crucial point is the following: in the ResNet, each residual layer has its own parameters. In the Neural ODE, the ode_func network is reused at every solver evaluation. The solver dynamically decides how many times to evaluate it. This gives a more compact model in terms of parameter count, with constant backpropagation memory.

Application to Irregular Time Series

One of the most powerful use cases for Neural ODEs is the processing of irregularly spaced time series. In many fields (medicine, finance, environmental sciences), observations are not collected at regular intervals. Classical RNNs require interpolating or filling missing values. Neural ODEs, on the other hand, naturally model a continuous trajectory and can be evaluated at any time $t$.

class IrregularTimeSeriesODE(nn.Module):
    """
    Neural ODE for irregular time series modeling.
    The network can be evaluated at any set of time points.
    """
    def __init__(self, input_dim, hidden_dim):
        super().__init__()
        self.encoder = nn.Linear(input_dim, hidden_dim)
        self.ode_func = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim)
        )
        self.decoder = nn.Linear(hidden_dim, input_dim)

    def forward(self, x, t_obs):
        """
        x: tensor of shape (batch, input_dim) - observed initial state
        t_obs: 1D tensor of observation times [t_0, t_1, ..., t_n]
        Returns predictions at each observation time.
        """
        h0 = self.encoder(x)
        t_obs = t_obs.to(x.device)
        # The solver integrates exactly at the requested times
        trajectory = odeint(self.ode_func, h0, t_obs, method="dopri5",
                            atol=1e-7, rtol=1e-7)
        predictions = self.decoder(trajectory)
        return predictions


# Example usage with irregular observations
def demo_irregular_timeseries():
    batch_size = 32
    input_dim = 5
    # Irregular observation times: not uniformly spaced
    t_obs = torch.tensor([0.0, 0.3, 0.7, 1.5, 2.0, 3.1, 4.8, 6.0])

    model = IrregularTimeSeriesODE(input_dim=input_dim, hidden_dim=32)
    initial_state = torch.randn(batch_size, input_dim)

    # The model returns the trajectory at each of the 8 irregular times
    output = model(initial_state, t_obs)
    print("Output shape: " + str(output.shape))
    # Each time gets its prediction, even though intervals vary
    # from 0.3 to 1.7 - the solver automatically handles this irregularity.

This ability to evaluate the model at arbitrary time points makes Neural ODEs superior to classical sequential architectures (LSTM, GRU) when dealing with heterogeneous temporal data.

Key Hyperparameters

Configuring Neural ODEs’ hyperparameters is crucial for good performance. Here are the main parameters to consider:

Solver Choice (method)

In practice, dopri5 is the default choice in the vast majority of cases. Its adaptive nature means it automatically adjusts the integration step size: it takes large steps when dynamics are slow and small steps when they are fast.

Tolerances (atol, rtol)

• atol (absolute tolerance): default 1e-7. Determines the maximum absolute error tolerated by the solver at each step.
• rtol (relative tolerance): default 1e-7. Determines the maximum relative error tolerated.

In general, the same value is used for both: atol=rtol. For quick experiments or debugging, you can increase these values to 1e-3 or 1e-4. For publication-quality results, lower them to 1e-7 or 1e-8. Warning: tolerances that are too low can lead to an exponential number of evaluations of the function $f$, significantly slowing training.

Integration Time (t_span)

The interval $[t_0, t_1]$ over which the ODE is integrated influences the effective depth of the network. A longer interval $[0, 10]$ allows a richer transformation than a short interval $[0, 0.1]$. In practice, $[0, 1]$ is a good starting point, and you can adjust according to the complexity of the problem.

Advantages and Limitations

Advantages

1. Constant $O(1)$ memory: thanks to the adjoint method, backpropagation memory does not depend on the number of solver evaluations. This is a major advantage over classical deep networks whose memory grows linearly with the number of layers.
2. Adaptive computational complexity: the ODE solver automatically decides the number of evaluations needed. “Easy” samples are processed quickly, while “difficult” samples receive more computation. This is a form of dynamic resource allocation.
3. Irregular time series modeling: unlike RNNs that require regularly spaced observations, Neural ODEs can be evaluated at any time $t$, making them natural for heterogeneous temporal data.
4. Reduced parameter count: the weights of the dynamics network $f$ are shared across the entire integration interval, unlike deep networks where each layer has its own parameters.
5. Rigorous connection to continuous mathematics: formalization via differential equations opens the door to theoretical analysis (stability, convergence) and mature numerical tools for ODE solving.

Limitations

1. Training speed: numerical integration at each pass (forward and backward) is more expensive than a simple forward pass in a classical network. Training can be significantly slower.
2. Optimization instability: solver tolerances can create discontinuous or noisy gradients, complicating optimizer convergence. Strategies such as regularization or progressive tolerance adjustment are sometimes needed.
3. Difficulty with very stiff dynamics: certain functions $f$ produce “stiff” differential equations where dynamics vary at very different time scales. These cases require specialized solvers (BDF) and are harder to optimize.
4. Lack of interpretability of hidden states: although the continuous trajectory is mathematically elegant, it is often harder to interpret what the state $h(t)$ represents at a given point compared to activations of a discrete layer.
5. Sensitive to initialization: the choice of initial parameters for the network $f$ strongly influences the behavior of the integrated trajectory. Poor initialization can lead to exploding or stagnant trajectories.

4 Concrete Use Cases

1. Irregularly Sampled Medical Data Modeling

In intensive care medicine, a patient’s vital signs are measured at irregular intervals — some measurements are taken every minute, others every hour. Neural ODEs can model the continuous evolution of the patient’s state between observations, enabling prediction of critical events (cardiac arrest, sepsis) even when data is widely spaced. Studies have shown that Neural ODEs outperform LSTMs and GRUs on real intensive care data.

2. Dynamical Systems and Physics-Informed Learning

Neural ODEs are ideal for learning the laws of physics from observation data. For example, from observed trajectories of a pendulum or a planetary system, a Neural ODE can learn the underlying dynamics function. This approach, called “Neural Laplace” or “Hamiltonian Neural Networks” in certain variants, enables the discovery of implicit physical laws without explicit supervision.

3. Generative Modeling with Continuous Normalizing Flows

Neural ODEs form the basis of Continuous Normalizing Flows (CNFs), a data generation method that transforms a simple distribution (such as a Gaussian) into a complex distribution via a continuous, reversible flow. Unlike GANs or VAEs, CNFs offer exact likelihood through the instantaneous change of variables theorem. This approach is used for image generation, density modeling, and variational inference.

4. Optimal Control and Robotics

In the field of physical system control (robots, drones, autonomous vehicles), Neural ODEs offer a natural way to model the continuous dynamics of the system and optimize control commands. By combining the adjoint method with optimal control techniques (such as Pontryagin’s maximum principle), we can learn control policies that respect the differential equations of the physical system while minimizing a given cost.

See Also

• Exploring Dynamic Polynomials with Python: Complete Guide and Practical Applications
• Mastering the Floor Operator in Python: The Arithmetic Calculation Revolution