Neural ODEs, Latent ODEs, and ODE-RNNs

"The discrete networks count their layers like steps on a staircase; I have no stairs. I am a slope. Ask me for my hidden state at time 4.137 and I will hand it to you, smooth and exact, because between every observation I have been quietly integrating, never once stopping to ask whether the clock ticked on the hour."

A Neural ODE Flowing Continuously Through Time

In Section 13.1 we built structured state-space models, the S4 and S5 family, which discretize a continuous linear state-space system $\dot{\mathbf{x}} = \mathbf{A}\mathbf{x} + \mathbf{B}u$ onto a fixed grid and then run the resulting linear recurrence with a fast parallel scan, and the linear-attention models of Section 13.3 likewise kept their recurrence discrete and linear. Those constructions kept the dynamics linear so the discretization had a closed form. This section removes the linearity restriction and keeps the time continuous: we let the right-hand side of the differential equation be an arbitrary neural network and solve the resulting nonlinear ODE numerically. The reward for giving up the closed-form scan is full nonlinearity and a state that is genuinely defined at every real-valued time, which is what we need for the irregularly sampled series introduced all the way back in Section 1.4 and engineered in Section 2.3. We use the unified notation of Appendix A: $\mathbf{h}(t)$ the continuous hidden state, $f_\theta$ the neural vector field with parameters $\theta$, and $t_0, t_1, \dots$ the (possibly irregular) observation times.

Why does a continuous reformulation of depth deserve its own section rather than a remark that "ResNets look like Euler"? Because the reformulation changes what a model is and what training costs. A discrete RNN or ResNet commits in advance to a number of steps and pays memory linear in that number. A Neural ODE commits only to a vector field and lets an adaptive solver decide how many steps the integration needs, spending more steps where the dynamics are stiff and fewer where they are smooth, and it can be trained with memory that does not grow with the number of those steps at all. For temporal data the payoff is sharper still: the discrete recurrence of Chapter 10 assumes evenly spaced steps, and real clinical, financial, and event streams are not evenly spaced. The continuous-time models here treat the gap between observations as a real number to integrate over, not a grid cell to fake, and that is a categorical, not cosmetic, improvement.

The competencies this section installs are four. To see a residual step as one Euler step and write the continuous-depth limit as an ODE solved by a black-box solver. To state the adjoint sensitivity method, explain why it gives constant memory, and weigh it against backpropagating through the solver. To explain how a Latent ODE and an ODE-RNN evolve a latent state continuously between irregular observations and why that beats a discrete RNN on such data. And to implement a Neural ODE both from scratch with a hand-written RK4 integrator and with torchdiffeq, fitting a continuous trajectory to scattered points. These carry directly into the Neural CDEs of Section 13.5.

1. The Continuous-Depth Idea: From a Residual Step to an ODE Beginner

Begin with the residual block that made very deep networks trainable. A residual layer does not output a transformed state directly; it outputs a correction added to its input, $\mathbf{h}_{\ell+1} = \mathbf{h}_\ell + f(\mathbf{h}_\ell, \theta_\ell)$, where $\ell$ indexes depth. Read that update with the eye of a numerical analyst and it is unmistakable: it is the forward Euler discretization of an ordinary differential equation. Euler's method for $\mathrm{d}\mathbf{h}/\mathrm{d}t = f(\mathbf{h}, t)$ with step size $\Delta t$ is $\mathbf{h}(t + \Delta t) = \mathbf{h}(t) + \Delta t \, f(\mathbf{h}(t), t)$, and the residual update is exactly this with $\Delta t = 1$ and depth playing the role of time. A ResNet is a fixed-step Euler integration of a vector field, with one quirk: each layer carries its own parameters $\theta_\ell$, so the vector field changes as you descend.

Now take the conceptual limit. Let the residual blocks become many and shrink the implicit step so the discrete depth index $\ell$ blurs into a continuous depth (or time) variable $t$, and tie the per-layer parameters into a single shared $f_\theta$ that may itself depend on $t$. The sequence of hidden vectors $\mathbf{h}_1, \mathbf{h}_2, \dots$ becomes a continuous trajectory $\mathbf{h}(t)$ governed by an initial value problem:

The network no longer specifies how to transform the state at each of finitely many layers; it specifies the instantaneous rate of change of the state, and the actual transformation from input $\mathbf{h}(t_0)$ to output $\mathbf{h}(t_1)$ is recovered by integrating that rate:

The integral has no closed form for a general neural $f_\theta$, so we evaluate it with a numerical ODE solver, a black box that, given the initial state, the vector field, and the integration interval, returns the state at the endpoint to a requested tolerance. This is the Neural ODE. Depth has become integration time; the layer count has become however many internal steps the solver chooses to take; and the model's entire content is the vector field $f_\theta$. Figure 13.4.1 contrasts the discrete residual staircase with the continuous flow.

Two immediate properties distinguish this from a fixed-depth network. First, the number of function evaluations is not fixed by architecture: an adaptive solver such as Dormand-Prince (the workhorse dopri5) chooses its own step sizes to meet a tolerance, taking many small steps through stiff regions and few large ones through smooth regions, so the model performs adaptive computation that a fixed-layer network cannot. Second, the map from $\mathbf{h}(t_0)$ to $\mathbf{h}(t_1)$ is, under mild conditions on $f_\theta$, a smooth invertible flow: trajectories of an ODE cannot cross, which gives Neural ODEs a built-in smoothness prior and makes them naturally suited to continuous normalizing flows, a use we note but do not pursue here.

2. Training: The Adjoint Sensitivity Method versus Backprop Through the Solver Intermediate

To train a Neural ODE we need the gradient of a scalar loss $L = L(\mathbf{h}(t_1))$, measured at the integration endpoint, with respect to the parameters $\theta$ and the initial state $\mathbf{h}(t_0)$. There are two ways to get it, and the choice is the central engineering decision of the method.

The direct way is backpropagation through the solver. A numerical solver is, after all, just a composition of arithmetic operations (each step evaluates $f_\theta$ a few times and combines the results), so we can record every one of those operations on the autograd tape and differentiate the whole unrolled integration exactly as we differentiated the unrolled recurrence in Section 9.3. This gives an exact gradient of the discretized computation, but it stores every intermediate state and every internal solver evaluation, so its memory grows linearly with the number of solver steps, which for a stiff problem and a tight tolerance can be large and unpredictable. It is the continuous analogue of full backpropagation through time, and it inherits the same $O(\text{steps})$ memory wall.

The elegant way, the one that made Neural ODEs practical, is the adjoint sensitivity method. Instead of recording the forward integration, define an adjoint state $\mathbf{a}(t) = \partial L / \partial \mathbf{h}(t)$, the sensitivity of the loss to the hidden state at time $t$. This adjoint obeys its own ordinary differential equation, run backward in time. The quantity $\mathbf{a}(t)^{\top}\,\partial f_\theta / \partial \mathbf{h}$ is a vector-Jacobian product, the same primitive autograd already computes, so the adjoint ODE is just that vector-Jacobian product integrated backward in time:

Integrating this adjoint ODE from $t_1$ back to $t_0$ transports the loss sensitivity to the start, and the parameter gradient is obtained by integrating one more quantity alongside it:

The decisive feature is what the adjoint method does not store. It does not keep the forward trajectory at all. To run the backward integration it needs $\mathbf{h}(t)$ along the way, and it recovers each value by re-integrating the original ODE backward from the endpoint $\mathbf{h}(t_1)$, in lockstep with the adjoint. The three quantities, the reconstructed state $\mathbf{h}(t)$, the adjoint $\mathbf{a}(t)$, and the accumulating gradient, are packed into one augmented ODE and solved together by a single backward call to the solver. Memory is therefore $O(1)$ in the number of solver steps: no matter how many steps the forward integration took, the backward pass holds only the current augmented state, not a tape of all of them. Figure 13.4.2 lays the two methods side by side.

The trade-off is now explicit. The adjoint method buys constant memory but pays for it twice: it runs a second integration (the backward solve roughly doubles compute), and reconstructing $\mathbf{h}(t)$ by reverse integration can drift from the true forward trajectory when the dynamics are stiff or chaotic, injecting a reconstruction error into the gradient that backpropagation through the solver does not have. The practical rule of thumb: use the adjoint when memory is the binding constraint or the integration is long; use backpropagation through the solver (or a checkpointed hybrid) when the problem is small enough to fit in memory and you want the most accurate gradient. Both are one line away in torchdiffeq (odeint versus odeint_adjoint), which is why the library makes the choice a flag rather than a rewrite.

3. Continuous-Time Sequence Models: Latent ODE and ODE-RNN Intermediate

Everything so far integrated from one time to another with no data arriving in between. The payoff for temporal AI comes when the integration limits are the observation times of an irregularly sampled series and observations update the state as they arrive. This is the home ground of the Latent ODE and the ODE-RNN of Rubanova, Chen, and Duvenaud (2019), and it is where continuous time stops being elegant and starts being necessary.

Recall the problem first raised in Section 1.4 and made concrete in Section 2.3: real temporal data is frequently sampled at irregular, uneven times. A patient's vitals are measured when a clinician happens to check, not on a fixed grid; a user's clicks arrive whenever they click; a sensor reports on change, not on a clock. A discrete RNN of Chapter 10 assumes one update per evenly spaced step, so feeding it irregular data forces an ugly choice: bin the timeline and impute the empty bins, or append the elapsed time as a feature and hope the network learns what a gap means. Both are workarounds for a model that has no native notion of "time elapsed".

The continuous-time models fix this at the root. The ODE-RNN keeps a hidden state but evolves it in two alternating modes. Between two consecutive observations at times $t_{i-1}$ and $t_i$, it lets the state flow according to a Neural ODE, $\mathbf{h}_i^{-} = \texttt{ODESolve}(\mathbf{h}_{i-1},\, f_\theta,\, t_{i-1},\, t_i)$, so the elapsed real time $t_i - t_{i-1}$ literally sets the integration length. Then, at the observation, it applies an ordinary RNN update to fold in the new measurement, $\mathbf{h}_i = \texttt{RNNCell}(\mathbf{h}_i^{-},\, \mathbf{x}_i)$. Flow, then jump; flow, then jump. A long gap produces a long flow and a correspondingly larger drift of the state; a short gap barely moves it. The model never sees a grid; it sees real elapsed times and integrates over them. This is precisely the continuous generalization of the discrete recurrence, and it realizes the temporal thread from the Kalman predict-update cycle of Chapter 7: the ODE flow is the continuous-time "predict" step, and the RNN cell is the "update" step that assimilates a measurement.

The Latent ODE goes one step further into a generative, fully continuous latent. An encoder (often an ODE-RNN run backward over the observed points) produces a distribution over an initial latent $\mathbf{z}(t_0)$; a single Neural ODE then deterministically integrates that latent forward through the entire timeline, $\mathbf{z}(t) = \texttt{ODESolve}(\mathbf{z}(t_0), f_\theta, t_0, t)$; and a decoder maps the latent at each observation time to the data space. The whole latent trajectory is one smooth ODE solution, so it can be queried at any time, observed or not, which makes the Latent ODE a true continuous-time generative model: it interpolates between observations and extrapolates past them by simply integrating further, with calibrated uncertainty inherited from the variational encoder. Figure 13.4.3 contrasts the discrete RNN's grid assumption with the ODE-RNN's flow-and-jump.

The contrast with the discrete RNN is therefore not a tuning detail but a difference in what the model can represent. A discrete RNN's state is undefined between steps and its dynamics are blind to how much time a step spanned; an ODE-RNN's state is defined for every real $t$ and its between-observation evolution is governed by genuine elapsed time. On irregularly sampled benchmarks (the clinical PhysioNet and MIMIC tasks, the human-activity datasets) the continuous-time models consistently outperform grid-and-impute RNN baselines, and they do so while producing a state you can query at any timestamp, which is exactly what a forecaster asked "what is the value at 3:47 am" needs.

4. Strengths and Costs Advanced

The continuous-depth view delivers real advantages, and they cluster around three themes. Adaptive computation: an adaptive solver spends function evaluations where the dynamics are hard and saves them where the dynamics are easy, so a single model can be cheap on smooth inputs and accurate on stiff ones without changing its architecture, something no fixed-depth network does. Constant-memory training: the adjoint method makes training memory independent of the integration length, which is what lets these models train on long horizons that would overflow a backpropagation-through-time tape. A smoothness prior and continuous queries: because the state is the solution of an ODE it is smooth in time and defined everywhere, so the model interpolates and extrapolates by construction and is a natural fit for irregular data and for invertible flows.

The costs are equally real and worth naming so you reach for the method only when it earns its keep. Solver cost and latency: each forward pass is a numerical integration that may take dozens to thousands of function evaluations, and each evaluation is a full neural-network call, so a Neural ODE can be far slower per example than a fixed-depth network of comparable width, and the cost is data-dependent and hard to predict. Stiffness: if the learned vector field has fast and slow components together (a stiff system), an explicit adaptive solver must take very small steps to stay stable, the function-evaluation count explodes, and training slows to a crawl or the solver fails; controlling stiffness (through regularizing the dynamics, weight decay on $f_\theta$, or stiff solvers) is an active practical concern. Adjoint reconstruction error: as noted in subsection two, reconstructing the forward trajectory backward can drift for unstable dynamics, corrupting the gradient. And there is a representational limit: a plain Neural ODE defines a homeomorphism, so it cannot represent maps that require trajectories to cross (the classic failure on data that an ordinary network handles trivially), which motivated augmented Neural ODEs that add extra dimensions to give trajectories room to move.

Read the strengths and costs together and a clean decision emerges. Continuous-time models are the right tool when the data is irregularly sampled or you must query the state off-grid, when memory forces constant-memory training, or when a smoothness prior matches the problem. They are the wrong tool when a fixed-depth network suffices, when latency budgets are tight, or when the dynamics are stiff and you cannot tame them. Section 13.5 shows how Neural Controlled Differential Equations push the continuous-time idea further to handle the incoming data stream itself as a continuous control path, which sharpens the irregular-sampling story still more.

5. Worked Example: A Neural ODE From Scratch and With torchdiffeq Advanced

We now make the section executable. The plan mirrors the from-scratch-then-library discipline used throughout the book: first integrate a Neural ODE with a hand-written RK4 solver so every arithmetic step is visible, then solve the identical problem with torchdiffeq's odeint and odeint_adjoint, and finally fit a continuous latent trajectory to a set of irregularly spaced points. Code 13.4.1 is the from-scratch RK4 integrator and a tiny vector-field network.

import torch
import torch.nn as nn

torch.manual_seed(0)

class VectorField(nn.Module):
    """The neural right-hand side f_theta(h, t): dh/dt = f_theta(h, t)."""
    def __init__(self, dim, width=32):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, width), nn.Tanh(),
            nn.Linear(width, dim),
        )
    def forward(self, t, h):                 # signature (t, h) matches torchdiffeq
        return self.net(h)                   # autonomous field: no explicit t dependence here

def rk4_odesolve(f, h0, t_grid):
    """Integrate dh/dt = f(t, h) from t_grid[0] to t_grid[-1] with classical RK4.
    Returns the state at every grid time, so trajectories can be supervised."""
    hs = [h0]
    h = h0
    for i in range(len(t_grid) - 1):
        t, dt = t_grid[i], t_grid[i + 1] - t_grid[i]
        k1 = f(t, h)                         # four slope evaluations per step
        k2 = f(t + dt / 2, h + dt / 2 * k1)
        k3 = f(t + dt / 2, h + dt / 2 * k2)
        k4 = f(t + dt,     h + dt * k3)
        h = h + dt / 6 * (k1 + 2 * k2 + 2 * k3 + k4)   # weighted average of slopes
        hs.append(h)
    return torch.stack(hs)

field = VectorField(dim=2)
h0 = torch.tensor([[2.0, 0.0]])              # one initial state (batch of 1, dim 2)
t_grid = torch.linspace(0.0, 5.0, steps=50) # 49 RK4 steps from t=0 to t=5
traj = rk4_odesolve(field, h0, t_grid)       # full state trajectory, shape (50, 1, 2)
print("from-scratch RK4: steps =", len(t_grid) - 1,
      "| h(5) =", traj[-1].detach().numpy().round(4))

from-scratch RK4: steps = 49 | h(5) = [[ 0.1473 -0.6218]]

Now the library pair. Code 13.4.2 hands the identical VectorField to torchdiffeq.odeint, which runs an adaptive Dormand-Prince solver and chooses its own step count, and to odeint_adjoint, which computes the gradient by the constant-memory adjoint method of subsection two. The from-scratch solver and its supervision loop ran about 18 lines; the library does the same integration in one call and switches to constant-memory adjoint gradients by changing a single function name.

from torchdiffeq import odeint, odeint_adjoint

t_eval = torch.linspace(0.0, 5.0, steps=50)

# Adaptive solver, gradient by backprop THROUGH the solver (stores the trajectory).
traj_lib = odeint(field, h0, t_eval, method="dopri5", rtol=1e-6, atol=1e-8)

# Same solve, gradient by the ADJOINT method (constant memory in solver steps).
traj_adj = odeint_adjoint(field, h0, t_eval, method="dopri5", rtol=1e-6, atol=1e-8)

print("torchdiffeq dopri5 | h(5) =", traj_lib[-1].detach().numpy().round(4))
print("adjoint vs through-solver endpoint match:",
      torch.allclose(traj_lib[-1], traj_adj[-1], atol=1e-4))

# Both give a usable gradient; the adjoint path never stored the trajectory.
loss = (traj_adj[-1] ** 2).sum()
loss.backward()                              # constant-memory backward solve
gnorm = sum(p.grad.norm().item() for p in field.parameters())
print("adjoint gradient norm =", round(gnorm, 4))

torchdiffeq dopri5 | h(5) = [[ 0.1471 -0.6219]]
adjoint vs through-solver endpoint match: True
adjoint gradient norm = 3.2814

Finally we use the machine for its intended purpose: fitting a continuous latent trajectory to irregularly sampled observations, the use case of subsection three. Code 13.4.3 generates noisy points at uneven times, then trains a Neural ODE (with adjoint gradients) so its integrated trajectory passes through them, demonstrating that a single continuous solution can be supervised on a non-uniform time grid.

import torch.optim as optim

# Irregularly spaced observation times (note the uneven gaps) and noisy targets
# from a true decaying spiral, the kind of off-grid data that breaks fixed-step RNNs.
t_obs = torch.tensor([0.0, 0.3, 0.7, 1.5, 2.1, 3.4, 4.0, 5.0])
true = torch.stack([2.0 * torch.cos(t_obs) * torch.exp(-0.3 * t_obs),
                    2.0 * torch.sin(t_obs) * torch.exp(-0.3 * t_obs)], dim=1)
obs = true + 0.05 * torch.randn_like(true)         # add observation noise

model = VectorField(dim=2)
z0 = torch.zeros(1, 2, requires_grad=True)         # learnable initial latent state
opt = optim.Adam(list(model.parameters()) + [z0], lr=0.05)

for step in range(300):
    opt.zero_grad()
    pred = odeint_adjoint(model, z0, t_obs, method="dopri5")  # solve AT the irregular times
    loss = ((pred[:, 0, :] - obs) ** 2).mean()     # fit the trajectory to the scattered points
    loss.backward()                                # constant-memory adjoint backward
    opt.step()
    if step % 100 == 0:
        print(f"step {step:3d}  fit loss = {loss.item():.5f}")

print("final fit loss =", round(loss.item(), 5))

step   0  fit loss = 1.84213
step 100  fit loss = 0.01627
step 200  fit loss = 0.00731
final fit loss = 0.00498

Read the three code blocks as one argument. Code 13.4.1 made the solver concrete by hand, so RK4 is no longer a black box. Code 13.4.2 replaced the hand-written loop with one torchdiffeq call and flipped to constant-memory adjoint gradients with a one-word edit, confirming both routes agree. Code 13.4.3 then put the continuous solution to work on genuinely irregular data, training a latent trajectory to fit off-grid points, which is the capability that distinguishes these models from every discrete recurrence in Part III. The progression from hand-written integrator to library to applied irregular-data fit is the whole pedagogical arc of continuous-time modeling in miniature.

from torchdiffeq import odeint_adjoint
# field: nn.Module with forward(t, h);  t_obs: 1-D tensor of (irregular) times
traj = odeint_adjoint(field, h0, t_obs, method="dopri5", rtol=1e-6, atol=1e-8)
((traj[-1] - target) ** 2).mean().backward()   # constant-memory adjoint gradient