Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

Structured Simulators and Neural Operators › Neural Differential Equations

Michael BrenndoerferJune 22, 202610 min read

World Models HandbookAll chapters

Part 4 of 12Space, Physics, and Causality

Chapter 22Structured Simulators and Neural Operators › Neural Differential Equations

Reading section 5 of 11 · Neural Differential Equations · 10 minSection 5 of 11 · 10 min

Section 5 of 11 · 10 min read

Structured Simulators and Neural Operators › Neural Differential EquationsOne section · Sections in this chapter5 / 11 ↓

Sections in this chapter11

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Both previous approaches keep an explicit integrator and learn either parameters or a force law. Neural differential equations go one step further back: the learned object is the vector field. There is no hand-written equation and no assumed interaction structure; there is only the commitment that state evolves as a continuous-time flow, and a neural network is installed as the right-hand side of the flow equation. A neural ODE (Chen and colleagues, 2018) defines state evolution by

\frac{dz(t)}{dt} = f_\psi(z(t), t),

where:

$z(t)$ : the state of the system at continuous time $t$
$f_\psi$ : a neural network, with parameters $\psi$ , that defines the time derivative of the state

Trajectories are produced by a compatible ODE solver, and the model's job is to learn the right-hand side from data. The change of perspective is subtle but consequential: in a differentiable simulator the integrator is an implementation detail of your physics, while here the integrator is a generic engine and the entire modeling content lives in the vector field. The commitment is that state evolution can be represented by a sufficiently regular continuous-time vector field; the regularity needed for existence, uniqueness, and differentiability depends on the model and solver.

Neural ODE and the adjoint trick

A continuous-time dynamical system whose vector field is a neural network. Gradients are computed either by backpropagating through solver steps or with an adjoint method. Ideal continuous-adjoint storage can be independent of the number of accepted steps when the trajectory is reconstructible; practical implementations may retain workspace or checkpoints.

The connection to the adjoint method is especially clean. The continuous-adjoint formulation can avoid storing every accepted solver step by reconstructing or recomputing the trajectory during the backward solve. Its idealized memory cost is independent of the number of stored steps, excluding solver workspace and checkpoints, but numerical error and backward instability can make its gradients differ from differentiating the actual discrete solver. The appropriate choice is therefore implementation- and problem-dependent.

What does continuous time buy you, concretely?

Decoupling the model from the integrator. Because the learned object is a vector field rather than a discrete map, it can be evaluated with another compatible solver or time grid. Accuracy and stability can still change with the solver, tolerance, and sampling regime, so deployment at a new resolution requires validation rather than an assumption of invariance.
Irregular sampling. Observation times need not be uniform. The solver simply integrates between whatever times data arrives, which is awkward for fixed-window discrete models and natural here. This connects directly to the temporal state and memory questions of Part III, Chapter 4. Real sensors do not tick at metronomic intervals, and a model class that treats the timestamps as part of the problem rather than a nuisance is materially better suited to them.
Dynamical-systems tools. You can compute fixed points, Jacobians, Lyapunov exponents, or conservation properties of $f_\psi$ , and you can constrain it: divergence-free fields for incompressible flow, Hamiltonian structure for conservative systems (the Hamiltonian and Lagrangian neural network line of work). This is a two-way street: the tools analyze what was learned, and the constraints shape what can be learned. An exact autonomous Hamiltonian vector field conserves its Hamiltonian along the continuous flow, but a numerical rollout can still drift unless the integrator and step-size regime preserve the relevant structure accurately.

The main training design choice is one-step supervision versus unrolled supervision, and it is the same open-loop-versus-compounding issue we have flagged before. Supervising $f_\psi$ to match $\frac{z_{t+1} - z_t}{\Delta t}$ at every pair of adjacent states is easy and stable, but it is exactly the one-step trap from our discussion of rollout error: tiny one-step biases integrate into large trajectory biases. A vector field that is locally correct to within a percent can still point systematically off-trajectory, and over a hundred integration steps those small angular errors push the solution onto a path the true system never visits. Training through short unrolled segments pushes the model toward trajectories that stay on track, and is usually worth the extra cost.

Let us learn a Van der Pol oscillator from a single trajectory. Its dynamics,

\dot{z}_1 = z_2, \qquad \dot{z}_2 = \mu(1 - z_1^2)\, z_2 - z_1,

where:

$z_1, z_2$ : the two state variables, position-like and velocity-like respectively
$\mu$ : the nonlinearity parameter; with $\mu > 0$ the system exhibits a self-sustained limit cycle, and here $\mu = 1.5$
$\mu(1 - z_1^2) z_2$ : the state-dependent damping term, which removes energy for large $|z_1|$ and injects energy near the origin, producing the limit cycle

produce a self-sustained limit cycle: trajectories spiral out from near the origin and settle onto a closed loop. The Van der Pol system is a demanding target for exactly the reason it is a good test: its long-term behavior is a geometric object, a closed attracting loop, and reproducing it requires getting the energy balance right everywhere around the cycle, not just on average. A learned vector field that injects slightly too much energy will spiral outward forever; one that removes slightly too much will collapse to the origin. The learned model sees only states, never the equations.

In[17]:

Code

import numpy as np

mu_vdp, dt_ode = 1.5, 0.05


def vdp_rhs(z):
    return np.array([z[1], mu_vdp * (1.0 - z[0] ** 2) * z[1] - z[0]])


def rk4_np(f, z, h):
    k1 = f(z)
    k2 = f(z + 0.5 * h * k1)
    k3 = f(z + 0.5 * h * k2)
    k4 = f(z + h * k3)
    return z + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)


z = np.array([1.0, 1.0])  # start off the limit cycle
traj = [z.copy()]
for _ in range(2000):
    z = rk4_np(vdp_rhs, z, dt_ode)
    traj.append(z.copy())
vdp_traj = np.array(traj)

The model is a small MLP defining the vector field. Training unrolls 8 RK4 steps from a batch of starting states and penalizes the mismatch at every step, so gradients flow through the integrator. The unrolling length of 8 is itself a design decision: long enough that compounding errors are visible to the loss, short enough that gradients remain well conditioned. We supervise at every step of the segment rather than only the terminal step, which anchors the trajectory while still letting later steps correct compounding drift.

In[18]:

Code

import torch
import torch.nn as nn

torch.manual_seed(0)

f_node = nn.Sequential(
    nn.Linear(2, 64),
    nn.Tanh(),
    nn.Linear(64, 64),
    nn.Tanh(),
    nn.Linear(64, 2),
)
for p in f_node.parameters():
    nn.init.normal_(p, std=0.1)


def rk4_torch(f, z, h):
    k1 = f(z)
    k2 = f(z + 0.5 * h * k1)
    k3 = f(z + 0.5 * h * k2)
    k4 = f(z + h * k3)
    return z + (h / 6.0) * (k1 + 2 * k2 + 2 * k3 + k4)


seg_len = 8
starts = np.random.default_rng(3).integers(
    0, len(vdp_traj) - seg_len - 1, size=512
)
seg_data = torch.tensor(
    np.stack([vdp_traj[s : s + seg_len + 1] for s in starts]),
    dtype=torch.float32,
)

opt = torch.optim.Adam(f_node.parameters(), lr=1e-2)
node_losses = []
for step in range(600):
    batch = seg_data[torch.randint(0, seg_data.shape[0], (64,))]
    zb = batch[:, 0]
    loss = 0.0
    for t in range(seg_len):
        zb = rk4_torch(f_node, zb, dt_ode)
        loss = loss + torch.mean((zb - batch[:, t + 1]) ** 2)
    opt.zero_grad()
    loss.backward()
    opt.step()
    node_losses.append(loss.item() / seg_len)

Then we roll the learned vector field out for the full 100 seconds from the same initial condition and see whether it reproduces the limit cycle. This is a much harder test than the training loss suggests: 100 seconds is 2000 integrator steps, roughly 250 times the unrolled training length. It tests long-horizon consistency along and near the visited trajectory, not global recovery of the vector field.

In[19]:

Code

with torch.no_grad():
    zt = torch.tensor([[1.0, 1.0]])
    learned_traj = [zt]
    for _ in range(2000):
        zt = rk4_torch(f_node, zt, dt_ode)
        learned_traj.append(zt)
learned_rollout = torch.cat(learned_traj).numpy()
print(f"Final segment training loss: {node_losses[-1]:.2e}")

Out[20]:

Visualization

Closed-loop phase portrait with the learned trajectory tracing a similar loop. — Phase portrait of the ground-truth Van der Pol trajectory and of the neural ODE rollout from the same initial condition. The learned rollout traces a closed loop close to the true limit cycle.

Semilog plot of training loss decreasing over iterations. — Unrolled multi-step training loss for the neural vector field. The loss decreases overall across 600 optimization steps while supervising every step of each eight-step segment.

The learned rollout traces a loop close to the visited limit cycle. From segments of one clean simulated trajectory, this run learns a vector field whose long rollout visually approximates the oscillation. The figure does not by itself establish global vector-field recovery or precise amplitude and period errors. Two failure modes are worth knowing about before you rely on this pattern in earnest.

Stiffness. If the true dynamics have fast transients relative to the slow behavior of interest, explicit solvers need tiny steps, and training through them is expensive. The problem is not merely slow training: an adaptive solver may take very different numbers of internal steps for nearby inputs, which makes the training landscape itself jagged. Recognizing stiffness and switching to implicit solvers is a real skill, and the adjoint method helps because it does not need to store the fine-grained forward trajectory.

Drift and non-conservation. Nothing in the architecture prevents the learned vector field from slowly spiraling inward or outward where the true system conserves energy. The Van der Pol example survives this because the limit cycle is attracting, so small energy errors are corrected by the dynamics themselves; a conservative system with no attractor gives the learned model nowhere to hide, and unbounded drift accumulates. Structural variants address this: Hamiltonian and Lagrangian neural networks build conservation into the parameterization, and augmented neural ODEs add extra dimensions so the model can express trajectories that a fixed-dimensional field cannot fit without distortion. Uncertainty about where the model is trustworthy, rather than a single confident trajectory, is the deeper fix, and that is the subject of the next chapter.

Neural ODEs are also a conceptual bridge to Part V. Residual recurrent state updates can sometimes be interpreted as discretizations of continuous-time latent dynamics, with an encoder mapping observations into that state. Recurrent state-space models in general are broader: they may be gated, stochastic, discontinuous, noninvertible, or otherwise not embeddable in a same-dimensional ODE flow.