Neural Differential Equations in Python: Pick a Solver

Search for neural differential equations in Python and the results are library documentation. You get the torchdiffeq README, the diffrax docs, the torchdyn README, a paper introducing torchode and a university tutorial notebook. Each one explains its own API well. None of them tells you the thing that decides whether your first training run on real process data converges or stalls: whether the system you are modelling is stiff.

For toy spirals and two-moons datasets it does not matter. For a reactor, a cultivation basin or anything with fast chemistry next to slow transport, it is usually the whole problem. This guide compares the four Python libraries through that lens, and ends with a selection procedure you can run before you write a training loop.

If you need the underlying idea first, start with what a neural ODE is and why it is continuous.

Four libraries, one question the docs skip

The Python ecosystem has settled on four options. Three are PyTorch: torchdiffeq, torchdyn and torchode. One is JAX: diffrax. All of them let you write the right-hand side of an ODE as a neural network, integrate it with a numerical solver and backpropagate through the result.

The question they skip is stiffness. The working definition, quoted from Ernst Hairer in the stiff neural ODE paper by Kim and colleagues, is blunt: stiff equations are problems for which explicit methods do not work. The same paper notes that the dopri method chosen in the original neural ODE work is exactly the kind of explicit solver that fails on them. Physically, stiffness shows up as widely separated time scales: one state relaxes in microseconds while another drifts over hours.

The libraries side by side

torchdiffeq: the reference, explicit by default

torchdiffeq comes from the authors of the original neural ODE paper and is the library most tutorials use. It exposes one main function, odeint, and an odeint_adjoint variant that backpropagates with O(1) memory by solving an adjoint ODE in the backward pass. The gotcha the README flags itself: for the adjoint, your vector field must be an nn.Module, because that is how the parameters are collected.

Its solver list is the part to read carefully. The adaptive solvers are dopri8, dopri5 (the default), bosh3, fehlberg2 and adaptive_heun. The fixed-step ones are euler, midpoint, rk4, explicit_adams and implicit_adams. Everything except implicit Adams is explicit. The recommendation for most problems is dopri5, or rk4 with a small step size. That is sound advice for most machine learning uses, and the wrong default for a stiff plant model.

The further documentation adds two useful levers: separate tolerances and solver for the adjoint pass, and a seminorm option that can make the adjoint solve cheaper with adaptive solvers.

torchdyn and torchode: the PyTorch alternatives

torchdyn is a higher-level layer. Its NeuralODE class wraps any nn.Module, it trains through PyTorch Lightning, and it extends the sensitivity methods available in torchdiffeq. Choose it for ergonomics, not for a different numerical core.

torchode is the performance option. It solves every problem in a batch independently, with its own step size and accept or reject decisions, it is compatible with PyTorch's JIT compiler, and it ships a PID step size controller that saved up to 5% of steps in the authors' evaluation when step sizes change quickly. The paper also notes that the 5th order Tsitouras method (tsit5) is often recommended over dopri5 today.

diffrax: the JAX option with implicit solvers

diffrax covers ordinary, stochastic and controlled differential equations, offers Tsit5, Dopri8, symplectic and implicit solvers, makes everything vmappable and provides several adjoint methods. The implicit solvers are what set it apart for physical systems. Its ODE solver reference lists implicit Runge-Kutta methods as suitable for stiff problems, with Kvaerno5 described as an A-L stable, stiffly accurate 5th order ESDIRK method with 7 stages. It also ships IMEX methods that integrate a non-stiff part explicitly and a stiff part implicitly, which matches many reaction and transport models closely.

For irregularly sampled sensor data the same library solves the controlled case; we cover that separately in neural controlled differential equations for irregular time series.

What stiffness does to training, not just solving

Most stiffness advice is about solving a known equation efficiently. With a neural ODE the equation changes every gradient step, and stiffness can appear during training even if the data looked benign.

Kim and colleagues trained on the ROBER chemical kinetics problem, where after a short induction period one species sits roughly five orders of magnitude below the other two. The baseline model failed to learn the dynamics: training losses did not decrease and the gradient norm fluctuated heavily. Their sensitivity analysis found equation scaling essential; the same network without scaling failed on ROBER.

Fronk and Petzold argue that scaling and regularisation only reframe the problem rather than solve it. Their demonstration is direct. Training a polynomial neural ODE on the van der Pol oscillator with mu equal to 1000, using classical RK4 with a fixed step matched to the data interval, became unstable as the model developed stiffness and the training halted. Single-step implicit methods (backward Euler, trapezoid, Radau3 and Radau5) trained successfully and recovered the stiff dynamics. Two caveats from the same work: the trapezoid method is not L-stable and can introduce oscillations, and the cost of the implicit approach is driven by how many Newton iterations each step needs, which grows with system dimension.

There is a quieter cost too, and it appears in batches. In the torchode benchmark, a batch of van der Pol oscillators solved jointly takes a common step size close to the smallest one any instance needs, so torchdiffeq and TorchDyn needed up to four times as many steps as solvers that treat instances separately. One stiff trajectory in a batch taxes all of them.

Adjoint or backpropagate through the solver

There are two ways to get gradients. The adjoint sensitivity method, which the UvA deep learning tutorial traces back to Pontryagin and colleagues in 1962, solves a second ODE backwards in time and keeps memory constant. The alternative, often called discretise-then-optimise, differentiates through the solver's own operations.

torchdiffeq warns that backpropagating through the solver internals is not numerically stable for all solvers, and encourages the adjoint. diffrax takes the opposite default: its solver selection guide says that when training via discretise-then-optimise, the default checkpointed adjoint, accurate solutions are often not needed and a low-order solver such as Heun is the most efficient choice.

Stiffness changes the arithmetic. Kim and colleagues point out that an implicit solver needs a linear solve whose LU cost grows with the cube of the number of states, so a standard continuous adjoint scales with the cube of states plus parameters. With the parameter count of a neural network that is a heavy bill. Their quadrature adjoint splits the work so that only the state-sized problems are stiff, bringing the cost down to the cube of the states plus a term linear in the parameters.

The torchode paper adds a practical warning in the other direction: solving a separate adjoint per batch instance made its backward pass the slowest of the libraries tested by more than an order of magnitude on a normalising flow benchmark, which is why it also ships a joint adjoint that solves the whole batch together.

A selection procedure for a physical system

None of the libraries publishes this as a checklist, but it follows from the sources above.

1. Check for stiffness before choosing a library. If your process has time scales several orders of magnitude apart, or your simulator already needs an implicit solver, assume the neural ODE will be stiff too.
2. Scale the states. Normalise each state to a comparable magnitude before training. On ROBER this alone separated failure from success.
3. Non-stiff: stay explicit. torchdiffeq with dopri5 is fine, torchode with tsit5 and a PID controller is faster on batches, and in diffrax a low-order solver like Heun is the efficient choice for discretise-then-optimise training.
4. Stiff: go implicit. In diffrax that means Kvaerno3, Kvaerno4 or Kvaerno5, almost always paired with a PIDController, because implicit solvers may need to reject steps whose implicit subproblem is intractable and retry with a smaller one. Consider IMEX if only part of the model is stiff.
5. Use 64-bit precision at tight tolerances. diffrax lists the symptoms of running tolerances of 1e-8 or tighter in 32-bit: NaN gradients, step counts jumping from 8 to 800, and the gradient call changing the forward result.
6. Pick the gradient method last, based on memory budget and system size, not habit.

From a trained model to a controller

A fitted model is only half the job if it has to act on a live process. At Equation Labs, continuous-time models are the front end of a control stack, and the measured figures on our own programme are what make the solver choice concrete: growth-model accuracy with R-squared above 0.95, a safety layer answering in under 2 ms, and end-to-end edge latency brought down to 285 ms from 1.2 s, validated on a 500 L pilot basin with zero safety violations across 400 simulated years.

Two lessons carry over. First, the solver you train with is not free to swap later, so a stiff model trained with a forced explicit workaround is a model you will be compensating for in deployment. Second, every implicit step is a Newton solve, and that cost lands on the latency budget; it is one reason neural ODEs for nonlinear system identification trade accuracy against inference time, and why compressing the model for the edge is a separate engineering step rather than an afterthought.

FAQ

Can I use SciPy's stiff solvers from torchdiffeq?

torchdiffeq has a scipy_solver method whose solver option maps onto the method argument of scipy.integrate.solve_ivp, so SciPy's stiff methods can be selected there. Step callbacks are not supported for it.

Why does my neural ODE produce NaN gradients in JAX?

If you run tolerances of 1e-8 or tighter, the most likely cause is JAX's default 32-bit precision. diffrax lists NaN gradients, far more solver steps than necessary and a gradient call that changes the forward result as the symptoms. Enable 64-bit.

Is torchode faster than torchdiffeq?

For the forward pass, in its authors' benchmarks, yes. With the adjoint it depends on configuration: solving a separate adjoint per batch instance was much slower on a normalising flow benchmark, while the joint adjoint variant was faster than torchdiffeq and TorchDyn.

Do I need a neural ODE or a neural CDE for irregular sensor data?

If observations arrive at irregular times or with gaps, the controlled form is usually the better fit, because the model is driven by the observed data path rather than integrated from a single initial state. diffrax solves ordinary, stochastic and controlled equations in the same library.