What it actually does
A residual block computes h + f(h), which is Euler’s method for an ODE with a step size of one. Take that seriously: shrink the step, let the solver choose how many it needs, and the network stops being a fixed stack of layers and becomes a vector field that something else integrates. Two things follow that are not merely aesthetic. The trajectory is defined at every time, not just at layer boundaries, which is exactly what an irregularly sampled series needs. And because the dynamics are reversible, the backward pass can be recomputed instead of stored, so memory stops growing with depth.
- The underlying process is genuinely continuous in time, so a differential equation is the right description rather than a convenient one
- The learned vector field is smooth enough for the solver to integrate at a tractable step count
- Integration error is small enough that it does not contaminate the gradient — an assumption that needs checking rather than believing
- The input-to-output map can be a homeomorphism, or the state has been augmented so that it need not be
- Solver settings at serving time match those at training time, since they are part of the model
Loss function
The loss is applied only at the endpoint, which is worth pausing on: there is no per-layer supervision and no intermediate output, because there are no layers to attach one to. Everything unusual is in the constraint that the forward pass is an integral rather than a composition. Two consequences follow. The solver becomes part of the model — change its tolerance and you change the function, which is unlike anything else in this reference. And because a solution to an ODE cannot cross itself, the map from input to output is a homeomorphism, so a plain Neural ODE provably cannot represent a function that separates nested or interleaved regions. That is a genuine expressivity limit rather than an optimization difficulty, and the standard fix is to augment the state with extra dimensions so the trajectory has room to go around.
- the learned vector field: given a state and a time, it returns a derivative, never a next state
- the hidden state as a continuous function of time rather than a value per layer
- integration time, which plays the role depth plays in a conventional network
- whatever the task loss is — the objective is ordinary; the forward pass is not
How it is minimized
Adam on the endpoint loss, with gradients from the adjoint sensitivity method or from backpropagating through the solver
- the adjoint: the gradient of the loss with respect to the state at time t
- the adjoint is integrated BACKWARDS in time, alongside a recomputed forward state
- a vector-Jacobian product, never a materialized Jacobian
- accumulated as a third integral in the same augmented backward solve
The adjoint method is the idea worth taking away, and it generalizes well past this model. Backpropagating through a solver stores every intermediate state, so memory grows with the number of steps — and an adaptive solver chooses that number, which means memory becomes data-dependent and unpredictable. The adjoint instead solves an augmented system backwards in time: the state, the adjoint and the parameter gradient integrated together, with the forward state recomputed on the way rather than stored. Memory becomes constant in depth, at the cost of roughly doubling the compute and, more seriously, of numerical error, because recomputing the forward trajectory backwards is not exact for a stiff or chaotic system and the reconstructed states drift from the ones the forward pass actually visited. That trade — constant memory for a recomputation that may not be faithful — is the whole engineering story, and the honest advice is to use the adjoint when memory binds and direct backpropagation when it does not.
solver
Fixed-step (Euler, RK4) or adaptive (Dopri5). Adaptive is more accurate and makes cost data-dependent (RK4 / Dopri5)
tolerance (rtol, atol)
For adaptive solvers. Part of the model definition, not a runtime setting — changing it changes the function (1e-7 to 1e-3)
integration time
Plays the role depth plays elsewhere; longer means more solver steps and more expressive dynamics (1.0, usually normalized)
augmented dimensions
Extra state dimensions that let trajectories avoid crossing. Without them some functions are simply unrepresentable (0 to 10)
vector-field width
Capacity of the network computing the derivative; evaluated many times per forward pass (32 to 256)
max solver steps
A hard cap. Without it a stiff system can take unboundedly many steps and hang the training run (1000 to 10000)
Trains like any other model until it does not, and the failures are solver failures rather than optimization ones — which is what makes them unfamiliar. The characteristic one is creeping stiffness: as training progresses the learned dynamics get sharper, an adaptive solver needs more steps to hold tolerance, and the cost per iteration climbs steadily until a run that started at seconds per step is taking minutes. Nothing reports this except wall-clock time, and the number of function evaluations is the diagnostic people forget to log. The second is adjoint drift, where reconstructing the forward trajectory backwards accumulates error and the gradients become quietly wrong; it shows up as training that plateaus for no visible reason, and tightening tolerance or switching to direct backpropagation fixes it. The third is the topological limit — a plain Neural ODE cannot separate nested regions at all, so it will fail on such a task no matter how long it trains, and augmenting the state is the fix rather than more capacity.
O(NFE · cost of f) per forward pass, where the number of function evaluations is chosen by the solver and grows as the dynamics stiffen. The adjoint backward roughly doubles that while using memory constant in depth; direct backpropagation uses memory linear in the number of steps. That memory-versus-faithfulness trade is the defining cost characteristic of the model.
How, where, and why it is used
Time-Series Forecasting
ViableThe latent-ODE formulation: encode the observed history into an initial state, integrate the learned dynamics forward to whatever times you need, and decode. Because the state is defined continuously, observations at arbitrary times are read off the same trajectory — there is no resampling, no imputation, and no notion of a missing step, because there are no steps.
- Irregularly sampled clinical time series, where measurements arrive when someone took them and a fixed grid is a fiction
- Sensor networks with asynchronous or dropout-prone sampling, where resampling destroys the timing information
- Physical and biological systems where the true dynamics are continuous and a differential equation is the natural description
- Forecasting at arbitrary horizons from one trained model, since the horizon is an integration limit rather than an output dimension
Within the irregular-sampling niche this is genuinely the right tool, and the argument is structural rather than empirical: a discrete recurrent model has to pretend the data lies on a grid, and every way of doing that — resampling, imputation, a delta-time feature — discards or fabricates something. A continuous-time model does not have the problem at all. Outside that niche the case collapses. On regularly sampled data it is substantially more expensive than an LSTM or a patch transformer for no accuracy gain, and the extra machinery buys nothing. The honest trigger is irregular or asynchronous observation times, not sophistication.
- Feed actual observation times rather than an index, since the whole point is that spacing carries information
- Normalize the time axis so integration intervals are order one, or the solver spends its step budget on scale rather than on dynamics
- Augment the state with extra dimensions, which both relieves the topological limit and gives the dynamics room to be smooth
- Cap solver steps explicitly, or a stiff learned field can hang a training run indefinitely
- Log the number of function evaluations per batch — it is the only early warning that the dynamics are stiffening
Rolling-origin backtesting at the actual observation times, with an ODE-free baseline that consumes a delta-time feature — that comparison is what establishes whether continuous time is contributing anything, and it is routinely skipped. Report the number of function evaluations alongside accuracy, since a model that is marginally better and five times slower is usually not the better model.
- Using it on regularly sampled data, where a discrete model is cheaper and just as good
- Stiffening dynamics silently inflating cost per step across training
- Skipping the delta-time baseline, so the continuous formulation is never actually justified
- Treating solver tolerance as a runtime knob when it is part of the model definition
Anomaly Detection
AdaptedThrough the continuous normalizing flow variant, which is where the model earns a density. Because the change in log-density along a trajectory is the integral of the vector field’s divergence, an ODE that transports a simple base distribution to the data distribution gives an exact log-likelihood — and a low-likelihood point is the anomaly score, with no reconstruction or threshold heuristic in between.
- Density-based outlier detection where an exact likelihood is wanted rather than a reconstruction proxy
- Scientific and physical monitoring, where the density itself is the quantity of interest
- Irregularly sampled trajectory anomaly detection, combining the continuous-time and density arguments
- Cases where a calibrated probability rather than a score is required downstream
It gives an exact likelihood rather than a bound or a proxy, which is genuinely rarer than it sounds — a VAE gives a bound, an autoencoder gives a reconstruction residual that is not a density at all, and this gives the actual number. The reason it is adapted rather than primary is cost: evaluating the divergence exactly is quadratic in dimension, so practical implementations use a stochastic trace estimator that reintroduces variance into the very quantity you wanted exactly. Combined with solver cost, density evaluation is orders of magnitude more expensive than an isolation forest, which on ordinary tabular data is competitive. Reach for this when an exact density is genuinely required and the dimension is modest.
- Use the stochastic trace estimator in dimensions above roughly ten, and report that the likelihood is now an estimate rather than exact
- Tighten solver tolerance for likelihood evaluation relative to training, since density is far more sensitive to integration error than a point prediction
- Standardize inputs, because the base distribution is standard and a badly scaled input makes the flow work to fix the scale
- Calibrate the likelihood threshold on held-out normal data rather than assuming a level
Precision and recall at a fixed alert budget, with an isolation forest baseline on the same data — the cost difference is several orders of magnitude and the accuracy difference frequently is not. Verify likelihood stability under tolerance changes: if the score moves when the solver tolerance moves, the density is being reported more precisely than it is known.
- Quoting an exact likelihood while using a stochastic trace estimator, which makes it an estimate with real variance
- Solver tolerance too loose for the density, so scores shift between runs of the same model
- Enormous cost against an isolation forest that performs comparably on tabular data
- Training data containing the anomalies, which the density then assigns high probability
Optimization
Primary fitThis entry is substantially an optimization and numerical-methods topic wearing a modelling hat. The adjoint sensitivity method is the centrepiece: gradients of a loss defined through an ODE solve, obtained by integrating an augmented system backwards rather than by storing a tape, giving memory constant in depth. Around it sit adaptive step control with embedded error estimates, stiffness as the reason a solver takes more steps, and the general principle that the numerical method is part of the model rather than an implementation choice beneath it.
- Adjoint sensitivity as memory-constant differentiation through an implicit computation — the transferable idea here
- Adaptive step control with embedded error estimation, where the same evaluations give both the step and its error
- Stiffness as a diagnosable cost driver, measured by function evaluations rather than by wall clock
- Checkpointed or interleaved schemes trading recomputation against storage, the same frontier as gradient checkpointing
Worth studying because it is the clearest case in this reference of differentiating through something that is solved rather than computed, and because the trade it exposes is real and uncomfortable: the adjoint buys constant memory by recomputing the forward trajectory backwards, and that recomputation is not exact. On a stiff or chaotic system the reconstructed states drift from the ones the forward pass visited and the gradients are quietly wrong — a failure with no error message. Knowing that the correct answer is often "use direct backpropagation and pay the memory" is more useful than knowing the adjoint derivation.
- Log function evaluations per batch as the primary cost metric; wall clock conflates stiffness with hardware
- Use direct backpropagation when memory permits, and reserve the adjoint for when it genuinely does not
- Tighten tolerance until the loss stops changing, which is the practical test that integration error is no longer contaminating the gradient
- Cap solver steps hard, since a stiff field can otherwise consume an unbounded step budget and hang the run
Compare adjoint gradients against direct backpropagation on a small problem; a divergence between them is integration error, not a bug, and it is the signal that tolerance is too loose. Sweep tolerance and watch both the loss and the evaluation count — the point where the loss stabilizes is the tolerance the model actually needs.
- Trusting adjoint gradients on a stiff system, where the backward reconstruction drifts and the gradients are wrong with no error
- Treating the solver as infrastructure rather than as part of the model, so tolerance changes silently change the function
- Optimizing wall clock while the real cost driver is a rising evaluation count
Control & Operations
ViableSystem identification: learn the continuous dynamics of a physical system from trajectory data, then use the learned field inside a model-predictive controller or a simulator. The learned model is an ODE, which is the same object the control literature already works with, so it drops into existing machinery rather than requiring a new one.
- Learning dynamics for model-predictive control from observed trajectories
- Digital twins of physical processes where the underlying description is genuinely a differential equation
- Hybrid models where a known physical ODE is augmented by a learned residual term
- Continuous-time reinforcement learning, where the dynamics model is an ODE rather than a discrete transition
The output is in the form control theory expects, which is a larger practical advantage than it sounds — a learned discrete transition model has to be converted or re-derived before a continuous controller can use it, and this does not. The hybrid form is the most defensible use: keep the known physics and learn only the residual, which preserves the guarantees the known part carries. Against it: a controller calling an adaptive solver inside its inner loop has unpredictable latency, which is a hard problem in a real-time loop, and a learned field carries no stability guarantee at all, so it can produce trajectories that are physically impossible with nothing to prevent it.
- Prefer the hybrid form — known physics plus a learned residual — since it keeps whatever the known part guarantees
- Use a fixed-step solver in a control loop; an adaptive one has unpredictable latency and that is disqualifying in real time
- Constrain the learned field where physics permits, such as enforcing energy dissipation, rather than hoping it is learned
- Validate on trajectories much longer than those trained on, since compounding is where a learned field fails
Long-horizon trajectory error against held-out rollouts, plus closed-loop control performance — open-loop prediction error and closed-loop regret diverge sharply, and only the second matters. Check for physically impossible behaviour explicitly, because nothing in the training objective forbids it.
- Unpredictable solver latency breaking a real-time control loop
- A learned field with no stability guarantee producing divergent or impossible trajectories
- Validating only on short horizons, where compounding error has not yet appeared
What it costs to run
Substantially more than a comparable discrete model, and variably so: cost is the number of function evaluations times the cost of the vector field, and an adaptive solver chooses that number based on how stiff the learned dynamics have become. A run that starts fast can become slow without any change to the code.
One solve per input, with the same variability. A fixed-step solver gives predictable latency at the cost of accuracy, which is usually the right trade in production. Adaptive solvers should not be placed behind a latency SLA without a hard step cap.
Driven by the domain rather than the model — physical systems drift slowly, clinical practice changes on its own schedule. Nothing about the architecture makes retraining more or less frequent than its discrete equivalents.
- Track function evaluations per request as a first-class metric; a rising count means the dynamics have stiffened and it precedes every latency problem
- Alert on solver step-cap hits, which mean the solution was truncated rather than converged and the output is not what the model defines
- Monitor latency at the tail rather than the mean, since adaptive solvers produce a long and input-dependent tail
- Re-verify that outputs are stable under a tolerance change; drift there means integration error has become material
- The solver and its tolerance are part of the model. Serving with different settings than training used gives a different function, with nothing to indicate it
- Adaptive step counts make latency data-dependent, so a single stiff input can take an order of magnitude longer than the median — a hard step cap is mandatory behind any SLA
- Adjoint gradients can be silently wrong on stiff dynamics. If training plateaus inexplicably, compare against direct backpropagation before touching the architecture
- A plain Neural ODE provably cannot represent a map that separates nested regions; if the task needs that, augment the state rather than adding capacity
- Solver libraries differ in their step-selection heuristics, so the same model under a different solver library is not quite the same model
Pros and cons, in context
Handles irregular sampling natively
The state exists at every time, so observations at arbitrary times need no resampling, imputation or delta-time feature. Decisive on clinical and asynchronous sensor data; worth nothing on a regular grid.
Memory constant in depth via the adjoint
Training cost in memory does not grow with the number of solver steps, which is what makes very deep effective dynamics affordable. The transferable idea in this entry, and it applies wherever you differentiate through something solved rather than computed.
Exact likelihoods through the continuous flow variant
Not a bound like a VAE’s and not a proxy like a reconstruction error — the actual density. Genuinely rare, and the reason this appears in density estimation at all.
The output is a differential equation
Which is the object physics and control theory already work with, so the learned model composes with existing machinery instead of needing conversion. A larger practical advantage than it first appears.
Cost is variable and can rise during training
An adaptive solver takes more steps as dynamics stiffen, so cost per iteration climbs with nothing reporting it but wall-clock time. Genuinely disorienting the first time, and the reason function evaluations must be logged.
A provable expressivity limit
Trajectories cannot cross, so a plain Neural ODE cannot represent a map separating nested regions — no amount of capacity or training fixes it. Augmenting the state is the fix, and knowing that in advance saves a great deal of confusion.
Adjoint gradients can be quietly wrong
Reconstructing the forward trajectory backwards is inexact on stiff or chaotic systems, so gradients drift with no error message and training plateaus for no visible reason. The constant-memory benefit has a real and under-advertised cost.
The solver is part of the model
Change tolerance and you change the function. Unlike almost anything else here, an inference-time configuration difference is a model difference, which makes the training-to-serving contract wider than people expect.
Loses decisively outside its niche
On regularly sampled data a discrete model is cheaper, simpler and just as accurate. The honest trigger is irregular observation times or a genuine need for continuous dynamics, not elegance.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""A Neural ODE and its adjoint, transcribed the way the paper reads.
A residual block computes h + f(h), which is Euler's method with a step size
of one. Take that seriously:
dh/dt = f(h, t; theta) the network emits a DERIVATIVE, not a state
h(t1) = h(t0) + integral a solver turns that into a trajectory
"Depth" becomes integration time and the layer count becomes a solver
decision. The adjoint then gives gradients by integrating an augmented system
BACKWARDS, so memory does not grow with the number of steps.
Plain loops, no libraries, explicit finite-difference Jacobian products.
"""
import math
def vector_field(state, time, params):
"""One hidden layer emitting a derivative. The whole network.
Note it takes `time` as an input: the dynamics may be non-autonomous, and
a field that ignores t cannot represent anything whose behaviour changes
over the integration interval.
"""
hidden = []
for row, bias in zip(params['w1'], params['b1']):
total = bias + time * params['w_time'][len(hidden)]
for weight, value in zip(row, state):
total += weight * value
hidden.append(math.tanh(total))
derivative = []
for row, bias in zip(params['w2'], params['b2']):
total = bias
for weight, value in zip(row, hidden):
total += weight * value
derivative.append(total)
return derivative
def euler_step(state, time, step, params):
"""h + dt * f(h, t). Exactly a residual block with a scaled update."""
derivative = vector_field(state, time, params)
return [h + step * d for h, d in zip(state, derivative)]
def rk4_step(state, time, step, params):
"""Four evaluations, weighted. Fourth-order accurate against Euler's first.
The weights are not arbitrary: they are chosen so the Taylor expansion of
the numerical step matches the true solution to fourth order. Four times
the work per step, but the step can be far larger for the same error, so
it is usually cheaper overall.
"""
k1 = vector_field(state, time, params)
mid = [h + 0.5 * step * d for h, d in zip(state, k1)]
k2 = vector_field(mid, time + 0.5 * step, params)
mid = [h + 0.5 * step * d for h, d in zip(state, k2)]
k3 = vector_field(mid, time + 0.5 * step, params)
end = [h + step * d for h, d in zip(state, k3)]
k4 = vector_field(end, time + step, params)
return [
h + step * (a + 2.0 * b + 2.0 * c + d) / 6.0
for h, a, b, c, d in zip(state, k1, k2, k3, k4)
]
def integrate(initial, t0, t1, num_steps, params, method=rk4_step):
"""Fixed-step integration, keeping every state.
Keeping the trajectory is what direct backpropagation needs, and it is
exactly what the adjoint below avoids: memory here grows with num_steps.
"""
step = (t1 - t0) / num_steps
state = list(initial)
trajectory = [list(state)]
for index in range(num_steps):
state = method(state, t0 + index * step, step, params)
trajectory.append(list(state))
return state, trajectory
def jacobian_vector_product(state, time, adjoint, params, epsilon=1e-6):
"""a^T df/dh by finite differences.
A real implementation gets this from automatic differentiation; done
literally it is one extra field evaluation per dimension, which is why
this is the expensive part of the adjoint.
"""
base = vector_field(state, time, params)
result = [0.0] * len(state)
for dimension in range(len(state)):
nudged = list(state)
nudged[dimension] += epsilon
perturbed = vector_field(nudged, time, params)
for i in range(len(state)):
result[dimension] += adjoint[i] * (perturbed[i] - base[i]) / epsilon
return result
def adjoint_backward(final_state, loss_gradient, t0, t1, num_steps, params):
"""Integrate the augmented system BACKWARDS.
Three things travel together: the state (recomputed, not stored), the
adjoint dL/dh(t), and the accumulating parameter gradient. Memory is
constant in num_steps, which is the entire point.
The uncomfortable part: the state is RECONSTRUCTED by integrating
backwards, and that is not exact. On stiff or chaotic dynamics the
reconstructed trajectory drifts from the one the forward pass visited, and
the gradients are then quietly wrong with no error anywhere.
"""
step = (t1 - t0) / num_steps
state = list(final_state)
adjoint = list(loss_gradient)
parameter_gradient = 0.0
for index in range(num_steps):
time = t1 - index * step
# Reconstruct the state one step back. This is the approximation.
state = euler_step(state, time, -step, params)
# da/dt = -a^T df/dh, integrated backwards.
product = jacobian_vector_product(state, time, adjoint, params)
adjoint = [a + step * p for a, p in zip(adjoint, product)]
# dL/dtheta accumulates as a third integral in the same solve.
derivative = vector_field(state, time, params)
parameter_gradient += step * sum(a * d for a, d in zip(adjoint, derivative))
return adjoint, parameter_gradient
def count_function_evaluations(num_steps, method_name):
"""The cost metric that actually matters.
Wall clock conflates stiffness with hardware. Function evaluations do not,
and a rising count across training is the ONLY early warning that the
learned dynamics are stiffening - which is the characteristic way a
Neural ODE run becomes slow without any code changing.
"""
per_step = {'euler': 1, 'rk4': 4, 'dopri5': 6}
return num_steps * per_step[method_name]
O(NFE * cost of f) per forward pass, with NFE the number of function evaluations. Illustrative, not a measured benchmark: the finite-difference Jacobian product costs one extra field evaluation per state dimension, so the literal adjoint is O(d) times the forward pass rather than O(1) as a real vector-Jacobian product would be.
How this language actually wants it written.
"""The same model, with the solver as a type and its cost made visible.
Three things change. The solver becomes an explicit object rather than a
function argument, because tolerance and method are PART OF THE MODEL -
change them and the function changes, which is unlike anything else in this
reference. Adaptive stepping arrives with an embedded error estimate. And
function evaluations are counted and returned, because a rising count is the
only early warning that the learned dynamics are stiffening.
"""
from __future__ import annotations
import math
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, NamedTuple, Sequence
VectorField = Callable[[Sequence[float], float], list[float]]
class Method(Enum):
"""Fixed-step or adaptive. Not interchangeable at serving time."""
EULER = 'euler'
RK4 = 'rk4'
DOPRI5 = 'dopri5'
@property
def evaluations_per_step(self) -> int:
return {'euler': 1, 'rk4': 4, 'dopri5': 6}[self.value]
@property
def adaptive(self) -> bool:
return self is Method.DOPRI5
class StepBudgetExceeded(RuntimeError):
"""Raised when the solver hits its step cap.
Its own type because the consequence is specific: the solution was
TRUNCATED rather than converged, so the output is not what the model
defines. Without this cap a stiff learned field consumes an unbounded
budget and hangs the training run with no diagnostic at all.
"""
class ToleranceMismatch(ValueError):
"""Raised when serving settings differ from the trained ones.
The solver is part of the model. Serving at a different tolerance is a
different function, and nothing else in the system will say so.
"""
class ShapeMismatch(ValueError):
"""Raised on a dimension violation instead of computing nonsense."""
@dataclass(frozen=True)
class SolverConfig:
"""Frozen, because these settings define the function being computed.
That is the unusual claim worth stating plainly: for almost every other
model an inference-time setting is a performance choice. Here it changes
what the model computes.
"""
method: Method = Method.DOPRI5
rtol: float = 1e-5
atol: float = 1e-7
max_steps: int = 10_000
initial_step: float = 0.01
def __post_init__(self) -> None:
if self.rtol <= 0.0 or self.atol <= 0.0:
raise ShapeMismatch('tolerances must be positive')
if self.max_steps < 1:
raise ShapeMismatch('the step budget must admit at least one step')
def fingerprint(self) -> tuple[str, float, float]:
"""Comparable across training and serving, so a mismatch is catchable."""
return (self.method.value, self.rtol, self.atol)
def assert_matches(self, trained: SolverConfig) -> None:
"""Guard clause for a mismatch with no other symptom."""
if self.fingerprint() != trained.fingerprint():
raise ToleranceMismatch(
f'serving solver {self.fingerprint()} differs from the trained '
f'{trained.fingerprint()}; this is a different function'
)
class SolveResult(NamedTuple):
"""State AND cost.
Evaluations come back rather than being discarded because they are the
primary cost metric: wall clock conflates stiffness with hardware, and a
rising evaluation count is the only early warning that the dynamics have
stiffened.
"""
state: list[float]
evaluations: int
steps: int
rejected_steps: int
@property
def acceptance_rate(self) -> float:
"""A low rate means the step controller is fighting the dynamics."""
total = self.steps + self.rejected_steps
return self.steps / total if total else 1.0
# Dormand-Prince coefficients. The embedded pair is the point: the same six
# evaluations yield both a fifth-order solution and a fourth-order one, and
# their difference IS the error estimate - free, rather than requiring a
# second solve at half the step size.
_DOPRI_A: tuple[tuple[float, ...], ...] = (
(),
(1 / 5,),
(3 / 40, 9 / 40),
(44 / 45, -56 / 15, 32 / 9),
(19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729),
(9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656),
)
_DOPRI_C: tuple[float, ...] = (0.0, 1 / 5, 3 / 10, 4 / 5, 8 / 9, 1.0)
_DOPRI_B5: tuple[float, ...] = (35 / 384, 0.0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84)
_DOPRI_B4: tuple[float, ...] = (5179 / 57600, 0.0, 7571 / 16695, 393 / 640, -92097 / 339200, 187 / 2100)
class Solver:
"""Owns the configuration and the evaluation counter."""
def __init__(self, config: SolverConfig) -> None:
self._config = config
def integrate(
self, field: VectorField, initial: Sequence[float], t0: float, t1: float
) -> SolveResult:
if t1 == t0:
return SolveResult(list(initial), 0, 0, 0)
if self._config.method is Method.DOPRI5:
return self._adaptive(field, initial, t0, t1)
return self._fixed(field, initial, t0, t1)
def _fixed(
self, field: VectorField, initial: Sequence[float], t0: float, t1: float
) -> SolveResult:
steps = max(1, int(abs(t1 - t0) / self._config.initial_step))
if steps > self._config.max_steps:
raise StepBudgetExceeded(
f'{steps} fixed steps exceeds the budget of {self._config.max_steps}'
)
step = (t1 - t0) / steps
state = list(initial)
for index in range(steps):
state = self._step(field, state, t0 + index * step, step)
return SolveResult(state, steps * self._config.method.evaluations_per_step, steps, 0)
def _step(
self, field: VectorField, state: Sequence[float], time: float, step: float
) -> list[float]:
if self._config.method is Method.EULER:
return [h + step * d for h, d in zip(state, field(state, time))]
k1 = field(state, time)
k2 = field([h + 0.5 * step * d for h, d in zip(state, k1)], time + 0.5 * step)
k3 = field([h + 0.5 * step * d for h, d in zip(state, k2)], time + 0.5 * step)
k4 = field([h + step * d for h, d in zip(state, k3)], time + step)
return [
h + step * (a + 2.0 * b + 2.0 * c + d) / 6.0
for h, a, b, c, d in zip(state, k1, k2, k3, k4)
]
def _adaptive(
self, field: VectorField, initial: Sequence[float], t0: float, t1: float
) -> SolveResult:
"""Dormand-Prince with PI step control.
Guard clause on the step budget rather than an unbounded while loop:
a stiff field would otherwise reduce the step indefinitely and hang.
"""
state = list(initial)
time = t0
step = math.copysign(self._config.initial_step, t1 - t0)
evaluations = 0
accepted = 0
rejected = 0
while (t1 - time) * math.copysign(1.0, t1 - t0) > 1e-12:
if accepted + rejected >= self._config.max_steps:
raise StepBudgetExceeded(
f'the solver hit its {self._config.max_steps}-step budget at t={time}; '
'the solution is truncated, not converged'
)
# Do not overshoot the endpoint.
if abs(step) > abs(t1 - time):
step = t1 - time
stages: list[list[float]] = []
for stage in range(6):
offset = list(state)
for previous, coefficient in enumerate(_DOPRI_A[stage]):
for i in range(len(offset)):
offset[i] += step * coefficient * stages[previous][i]
stages.append(field(offset, time + _DOPRI_C[stage] * step))
evaluations += 6
fifth = [
h + step * math.fsum(b * k[i] for b, k in zip(_DOPRI_B5, stages))
for i, h in enumerate(state)
]
fourth = [
h + step * math.fsum(b * k[i] for b, k in zip(_DOPRI_B4, stages))
for i, h in enumerate(state)
]
# The embedded pair gives the error for free: their difference IS
# the local truncation estimate.
error = 0.0
for high, low, current in zip(fifth, fourth, state):
scale = self._config.atol + self._config.rtol * max(abs(current), abs(high))
error = max(error, abs(high - low) / scale)
if error <= 1.0:
state = fifth
time += step
accepted += 1
else:
rejected += 1
# Standard controller with clamping, so one bad step cannot make
# the next step absurdly large or small.
factor = 0.9 * (1.0 / max(error, 1e-10)) ** 0.2
step *= min(5.0, max(0.2, factor))
return SolveResult(state, evaluations, accepted, rejected)
@dataclass
class StiffnessMonitor:
"""Tracks evaluation counts across training.
A rising count is the characteristic way a Neural ODE run becomes slow
without any code changing, and nothing else reports it. This belongs on a
dashboard rather than in a comment.
"""
history: list[int] = field(default_factory=list)
def record(self, result: SolveResult) -> None:
self.history.append(result.evaluations)
def stiffening(self, window: int = 100, ratio: float = 2.0) -> bool:
if len(self.history) < 2 * window:
return False
early = math.fsum(self.history[:window]) / window
recent = math.fsum(self.history[-window:]) / window
return early > 0 and recent / early > ratio
What changed: Three changes. The solver becomes an explicit frozen configuration rather than a function argument, because tolerance and method are part of the model rather than a performance setting — change them and the function changes, which is unlike anything else in this reference, so there is a fingerprint and a mismatch check with its own exception type. Adaptive stepping arrives properly, and the embedded Dormand-Prince pair is the reason it is worth having: the same six evaluations yield both a fifth-order and a fourth-order solution, and their difference is the error estimate for free rather than requiring a second solve at half the step. The step budget is a hard guard rather than an unbounded loop, with its own exception, because a stiff learned field otherwise shrinks the step indefinitely and hangs a training run with no diagnostic — and hitting the cap means the solution was truncated rather than converged, which is a correctness statement rather than a timeout. Third, function evaluations and rejected steps come back in the result rather than being discarded, since a rising evaluation count is the only early warning that the dynamics are stiffening and a low acceptance rate means the controller is fighting them.
An adaptive solve costs six evaluations per attempted step, accepted or not, so the acceptance rate is a direct multiplier on cost. Illustrative, not a measured benchmark: the embedded error estimate is what makes adaptivity affordable, since computing it separately would roughly double the evaluations.
Faster, and the justification for each change.
"""Batched trajectories and a fused solver. The state gains a batch axis.
The structural observation: the solver loop over TIME is inherently
sequential, but nothing about it is sequential across the batch. So the state
becomes a matrix, every trajectory takes step t together, and the field
evaluation - which dominates everything - becomes a GEMM.
Three changes:
1. The vector field batches. One GEMM per layer for every trajectory at
once, instead of one matrix-vector product per trajectory.
2. The Runge-Kutta stages are fused: the stage offsets are computed in
place into a preallocated buffer, so a four-stage step allocates nothing.
3. The adjoint uses a real vector-Jacobian product rather than finite
differences, which turns the backward pass from O(d) field evaluations
per step into O(1).
"""
from __future__ import annotations
import numpy as np
from numpy.typing import NDArray
FLOAT = np.float32
# Dormand-Prince coefficients as arrays, so a stage combination is a
# contraction rather than a Python loop over previous stages.
_A = np.array([
[0, 0, 0, 0, 0, 0],
[1 / 5, 0, 0, 0, 0, 0],
[3 / 40, 9 / 40, 0, 0, 0, 0],
[44 / 45, -56 / 15, 32 / 9, 0, 0, 0],
[19372 / 6561, -25360 / 2187, 64448 / 6561, -212 / 729, 0, 0],
[9017 / 3168, -355 / 33, 46732 / 5247, 49 / 176, -5103 / 18656, 0],
], dtype=FLOAT)
_C = np.array([0, 1 / 5, 3 / 10, 4 / 5, 8 / 9, 1], dtype=FLOAT)
_B5 = np.array([35 / 384, 0, 500 / 1113, 125 / 192, -2187 / 6784, 11 / 84], dtype=FLOAT)
_B4 = np.array([5179 / 57600, 0, 7571 / 16695, 393 / 640, -92097 / 339200, 187 / 2100], dtype=FLOAT)
class BatchedField:
"""The vector field, evaluated for a whole batch of trajectories at once.
This is where all the time goes, so it is the only thing worth making
fast: a two-layer field is two GEMMs regardless of how many trajectories
are in flight.
"""
def __init__(self, w1: NDArray[np.float32], b1: NDArray[np.float32],
w_time: NDArray[np.float32], w2: NDArray[np.float32],
b2: NDArray[np.float32]) -> None:
self._w1 = np.ascontiguousarray(w1, dtype=FLOAT)
self._b1 = np.ascontiguousarray(b1, dtype=FLOAT)
self._w_time = np.ascontiguousarray(w_time, dtype=FLOAT)
self._w2 = np.ascontiguousarray(w2, dtype=FLOAT)
self._b2 = np.ascontiguousarray(b2, dtype=FLOAT)
def __call__(self, state: NDArray[np.float32], time: float) -> NDArray[np.float32]:
"""state is (batch, dim). Returns the derivative, same shape."""
hidden = state @ self._w1
hidden += self._b1
# Time enters as a rank-one update, not a tiled column: the field may
# be non-autonomous, and one broadcast is cheaper than a concatenate.
hidden += FLOAT(time) * self._w_time
np.tanh(hidden, out=hidden)
derivative = hidden @ self._w2
derivative += self._b2
return derivative
class BatchedSolver:
"""Dormand-Prince over a batch, with per-trajectory error control.
Note the design decision: the step is shared across the batch and the
error is the MAXIMUM over trajectories. Per-trajectory steps would be
better numerically and would destroy the batching entirely, so the whole
batch moves at the pace of its stiffest member.
"""
def __init__(self, batch: int, dim: int, rtol: float = 1e-5, atol: float = 1e-7,
max_steps: int = 10_000) -> None:
self._rtol = FLOAT(rtol)
self._atol = FLOAT(atol)
self._max_steps = max_steps
# Six stages preallocated: a step allocates nothing at all.
self._stages = np.empty((6, batch, dim), dtype=FLOAT)
self._offset = np.empty((batch, dim), dtype=FLOAT)
def integrate(
self,
field: BatchedField,
initial: NDArray[np.float32],
t0: float,
t1: float,
initial_step: float = 0.01,
) -> tuple[NDArray[np.float32], int, int]:
"""Returns (final state, evaluations, rejected steps)."""
state = np.array(initial, dtype=FLOAT, copy=True)
time = t0
step = float(np.copysign(initial_step, t1 - t0))
evaluations = 0
attempts = 0
rejected = 0
while (t1 - time) * np.sign(t1 - t0) > 1e-12:
if attempts >= self._max_steps:
raise RuntimeError(
f'solver hit its {self._max_steps}-step budget at t={time}; '
'the solution is truncated, not converged'
)
attempts += 1
if abs(step) > abs(t1 - time):
step = t1 - time
for stage in range(6):
# Stage offset as one contraction over previous stages,
# written into a preallocated buffer.
np.copyto(self._offset, state)
if stage:
self._offset += step * np.einsum(
'j,jbd->bd', _A[stage, :stage], self._stages[:stage], optimize=True
)
self._stages[stage] = field(self._offset, time + float(_C[stage]) * step)
evaluations += 6
fifth = state + step * np.einsum('j,jbd->bd', _B5, self._stages, optimize=True)
fourth = state + step * np.einsum('j,jbd->bd', _B4, self._stages, optimize=True)
# The embedded pair gives the error for free: their difference IS
# the local truncation estimate, so no second solve is needed.
scale = self._atol + self._rtol * np.maximum(np.abs(state), np.abs(fifth))
error = float(np.max(np.abs(fifth - fourth) / scale))
if error <= 1.0:
state = fifth
time += step
else:
rejected += 1
factor = 0.9 * (1.0 / max(error, 1e-10)) ** 0.2
step *= min(5.0, max(0.2, factor))
return state, evaluations, rejected
def vector_jacobian_product(
field: BatchedField,
state: NDArray[np.float32],
time: float,
adjoint: NDArray[np.float32],
epsilon: float = 1e-4,
) -> NDArray[np.float32]:
"""a^T df/dh in ONE extra field evaluation, not one per dimension.
A directional derivative along the adjoint: f(h + eps*a) - f(h), divided
by eps. The literal version cost O(d) evaluations per step, which made
the adjoint more expensive than the forward pass it was supposed to
replace; this is O(1) and is what makes the method viable at all.
A production implementation gets this exactly from reverse-mode AD; the
finite-difference form is here because it makes the O(1) claim visible.
"""
base = field(state, time)
nudged = field(state + FLOAT(epsilon) * adjoint, time)
return (nudged - base) / FLOAT(epsilon)
def adjoint_backward(
field: BatchedField,
solver: BatchedSolver,
final_state: NDArray[np.float32],
loss_gradient: NDArray[np.float32],
t0: float,
t1: float,
) -> tuple[NDArray[np.float32], int]:
"""Integrate the augmented system backwards, in one batched solve.
State, adjoint and parameter gradient travel together in one array, so
the whole backward pass is a single call to the same solver - memory
constant in the number of steps, which is the entire point.
The uncomfortable part is unchanged by batching: the state is
RECONSTRUCTED by integrating backwards, and on stiff dynamics it drifts
from the trajectory the forward pass actually visited. The gradients are
then quietly wrong, and only a comparison against direct backpropagation
will show it.
"""
batch, dim = final_state.shape
# Augmented state: [h, a] concatenated along the feature axis, so one
# solver call advances both rather than two interleaved ones.
augmented = np.concatenate((final_state, loss_gradient), axis=1)
def augmented_field(combined: NDArray[np.float32], time: float) -> NDArray[np.float32]:
state, adjoint = combined[:, :dim], combined[:, dim:]
derivative = field(state, time)
product = vector_jacobian_product(field, state, time, adjoint)
return np.concatenate((derivative, -product), axis=1)
result, evaluations, _ = solver.integrate(
BatchedFieldAdapter(augmented_field), augmented, t1, t0
)
return result[:, dim:], evaluations
class BatchedFieldAdapter:
"""Wraps a plain callable so the solver's contract stays a single type."""
def __init__(self, function) -> None:
self._function = function
def __call__(self, state: NDArray[np.float32], time: float) -> NDArray[np.float32]:
return self._function(state, time)
What changed: The structural observation is that the solver loop over time is inherently sequential but nothing about it is sequential across the batch, so the state gains a leading batch axis and every trajectory takes step t together — which turns the field evaluation, where all the time goes, into two GEMMs regardless of how many trajectories are in flight. Time enters as a rank-one broadcast rather than a tiled column, since the field may be non-autonomous and one broadcast is cheaper than a concatenate. The six Runge-Kutta stages are preallocated and each stage offset is one contraction over the previous stages written into an existing buffer, so a step allocates nothing. The largest single change is in the adjoint: the literal version computed the vector-Jacobian product by finite differences one dimension at a time, costing O(d) field evaluations per step and making the backward pass more expensive than the forward one it was meant to replace — a directional derivative along the adjoint gives the same quantity in one extra evaluation, which is what makes the method viable at all. The design decision worth naming is that the step size is shared across the batch and the error is the maximum over trajectories: per-trajectory steps would be numerically better and would destroy the batching, so the batch moves at the pace of its stiffest member.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
The field is two GEMMs for the whole batch, and each Runge-Kutta stage combination is one einsum over the stage axis instead of a Python loop over previous stages.
Trade-off: The six-stage buffer is batch times dimension times six in float32 and is live for the whole solve, so a large batch with a wide state makes the stage storage the dominant memory cost of the forward pass.
Pre-allocate output arrays and use in-place operations
Stages and the stage offset are written into existing buffers, and the field composes its layers in place, so an accepted or rejected step allocates nothing.
Trade-off: The offset buffer is reused across stages, so nothing can inspect an earlier stage offset after the fact — which is exactly what you want when diagnosing why a step keeps being rejected.
Batch work to amortize interpreter overhead
Trajectories become rows, so a solve costs a constant number of interpreter operations per step regardless of batch size.
Trade-off: One step size is shared across the batch and the error is the maximum over it, so the stiffest trajectory sets the pace for every other one — batching trades numerical efficiency per trajectory for throughput.
Fuse adjacent operations so intermediates are never materialized
The augmented backward system carries state, adjoint and gradient in one array so a single solver call advances all three, rather than interleaving separate solves.
Trade-off: The augmented state doubles the width the solver controls error on, so the shared tolerance now applies to the adjoint as well — which is usually right but means an accurate adjoint forces small steps on the state too.
The field becomes two GEMMs per evaluation for the whole batch, and the adjoint drops from O(d) evaluations per step to O(1). Illustrative, not a measured benchmark: since the batch shares one step size, cost is set by the stiffest trajectory in it — so a single pathological input can make an entire batch expensive.