Excessive allocations in basic SciMLSensitivity example

moble
Mon, 28 Jul 2025 20:30:27 +0000
General Usage

The first example (copied below) in the SciMLSensitivity docs shows how to optimize the parameters of the Lotka-Volterra system to achieve a steady state. There’s no neural network, or anything, just an ODE system with 2 variables, evaluated at 101 time steps, and 4 parameters to optimize. The optimization is set to go at most 100 iterations. So it’s really not a large problem — certainly not anything like the size of problems I want to solve.

And yet, the allocations are pretty out of control: 233.69 M allocations totaling 11.576 GiB! (There’s no compilation happening.) It reports taking 4,578 function and gradient evaluations (no Hessians). If we count each of those function-gradient evaluations as 10 numbers, evaluated at 101 time steps, 4,578 times, there are still 50x more allocations and 2500x more bytes than I can see any possible need for. I’m sure there’s another factor from stages of the timestepper, and another factor of a few because adjoint sensitivity is hard, but not such large factors. And that’s only if literally every number requires a separate allocation. What’s going on?!

Profiling, it looks like most of the allocations (~95%) are caused by gradient computations in HagerZhang — which correlates very well with where the time is being spent. This example problem still runs in a reasonable amount of time (~30s), but I’m trying a much harder problem (with more evolved variables, more cases, and a NN with many more parameters), and even my simplest case is basically dead in the water (about 1/4 of which seems to be GC).

Is there anything I can do about those allocations?

Full example code (click for more details)
16 posts - 6 participants