# Improving performance for training Universal Differential Equation

**URL:** <https://discourse.julialang.org/t/improving-performance-for-training-universal-differential-equation/105038>\
**Category:** Machine Learning\
**Created:** [October 16, 2023, 8:44pm UTC](https://discourse.julialang.org/t/improving-performance-for-training-universal-differential-equation/105038 "2023-10-16T20:44:01Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![rkube](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rkube/32/211198_2.png) [@rkube](https://discourse.julialang.org/u/rkube)\
**Post date:** [October 16, 2023, 8:44pm UTC](https://discourse.julialang.org/t/improving-performance-for-training-universal-differential-equation/105038/1 "2023-10-16T20:44:01Z")

</div>

Hi,  
I’m trying to understand the performance bottlenecks in training universal differential equations. Following [this guide](https://docs.sciml.ai/Overview/stable/showcase/missing_physics/#Training), benchmarking the training shows that optimizing the parameters of the Lux network requires a large number of allocations:

```julia
@btime res1 = Optimization.solve(optprob, ADAM(), maxiters = 100) 
417.338 ms (2205218 allocations: 201.85 MiB)

```

I’m going to run a larger number of similar optimization and would like to make reduce time to solution for the optimization problem. My problems are low dimensional and I would need to optimize for a large number of given datasets.  
Does anyone know how to minimize the number of allocations?

I’m thinking about either using StaticArrays, like described [here](https://docs.sciml.ai/DiffEqDocs/stable/tutorials/faster_ode_example/) , but am unsure whether they work with `Lux`. Another thing I’d like to look into is using [SimpleChains](https://github.com/PumasAI/SimpleChains.jl) but I haven’t figured out yet how to make their interface compatible with Optimization.

Does anyone have experience with optimizing for this situation?

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [October 16, 2023, 8:57pm UTC](https://discourse.julialang.org/t/improving-performance-for-training-universal-differential-equation/105038/2 "2023-10-16T20:57:42Z")

</div>

It depends on your model. What does your code look like? Do you have an MWE to play with?

---

<div class="post-metadata">

**Author:** ![rkube](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rkube/32/211198_2.png) [@rkube](https://discourse.julialang.org/u/rkube)\
**Post date:** [October 16, 2023, 9:31pm UTC](https://discourse.julialang.org/t/improving-performance-for-training-universal-differential-equation/105038/3 "2023-10-16T21:31:21Z")

</div>

I’m using [the missing physics guide](https://docs.sciml.ai/Overview/stable/showcase/missing_physics/#Training) as an MWE. Basically I’m wondering if there is a way to keep the number of allocations constant when varying the number of iterations for the optmizer. I’m seeing that they increase linearly with number of steps:

```julia
@btime Optimization.solve(optprob, ADAM(), maxiters=1);
  3.591 ms (21770 allocations: 2.00 MiB)

```

```julia
@btime Optimization.solve(optprob, ADAM(), maxiters=2);
 7.202 ms (43489 allocations: 4.01 MiB)

```

```julia
@btime Optimization.solve(optprob, ADAM(), maxiters=3);
 11.022 ms (65808 allocations: 6.04 MiB)

```

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [October 17, 2023, 7:18am UTC](https://discourse.julialang.org/t/improving-performance-for-training-universal-differential-equation/105038/4 "2023-10-17T07:18:22Z")

</div>

You’d have to do things like cache the ODE solver, which would be difficult with Zygote and require setting it all up with Enzyme.