# How can I significantly speed up conditional Universal Differential Equation training?

**URL:** <https://discourse.julialang.org/t/how-can-i-significantly-speed-up-conditional-universal-differential-equation-training/138126>\
**Category:** Performance\
**Tags:** machine-learning\
**Created:** [July 12, 2026, 12:34pm UTC](https://discourse.julialang.org/t/how-can-i-significantly-speed-up-conditional-universal-differential-equation-training/138126 "2026-07-12T12:34:29Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![NCCP](https://avatars.discourse-cdn.com/v4/letter/n/c5a1d2/32.png) [@NCCP](https://discourse.julialang.org/u/NCCP)\
**Post date:** [July 12, 2026, 12:34pm UTC](https://discourse.julialang.org/t/how-can-i-significantly-speed-up-conditional-universal-differential-equation-training/138126/1 "2026-07-12T12:34:29Z")

</div>

Hello,

I am training a conditional Universal Differential Equation in Julia. This is a regular UDE with an additional trainable parameter as input that is specific for each particle (individual) in the dataset.

The model has one ODE state, a small shared neural network, and one particle-specific conditional parameter for each of 600 simulated particle trajectories. Training currently takes several hours and training time increases with the number of particles.

The pipeline consists of:

- Jointly optimising the neural-network parameters and 600 conditional parameters using ADAM followed by L-BFGS (this is the bottle-neck)
- Reestimating each conditional parameter separately while keeping the neural network parameters fixed
- Parallelising the multistart runs with `pmap`

I would like to know whether this workflow can be made substantially faster without affecting the outcomes. I would appreciate any feedback that can help speed up this code (hopefully) by orders of magnitude.

Let me know if anything is unclear or required as I’m new to the Julia Discourse.  
Thank you for your time!

Kind regards,  
NCCP  
[dBPM.jl](https://discourse.julialang.org/uploads/short-url/fDeDEzdPs8y0L2yNFKNmZe886dK.jl) (13.2 KB)  
[inputs.jl](https://discourse.julialang.org/uploads/short-url/tPWKhaRKPIJvUODMxe6kTz6qoJK.jl) (7.3 KB)  
[ude.jl](https://discourse.julialang.org/uploads/short-url/qVnNvP6CcrwWEGjXJMKjNlsnoqv.jl) (18.1 KB)

---

<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:** [July 12, 2026, 7:52pm UTC](https://discourse.julialang.org/t/how-can-i-significantly-speed-up-conditional-universal-differential-equation-training/138126/2 "2026-07-12T19:52:33Z")

</div>

Two changes do the heavy lifting here, and neither changes the result.

1. Solve the trajectories as an `EnsembleProblem` with a non-stiff solver (`Tsit5`), threaded. The direct-BPM ODE isn’t stiff, `Tsit5` matches `Rodas4P` to ~1e-4 and each trajectory stays a cheap 1-D solve. (Don’t stack all particles into one big 600-state ODE instead: `Rodas4P` then factorizes a dense 600x600 Jacobian in dual arithmetic every step and actually gets slower. `EnsembleProblem` keeps the solves independent.)

2. Stop differentiating all 600 conditionals through every particle’s solve. With `AutoForwardDiff` over the full `[57 NN weights; 600 β]` vector, every one of the 600 solves drags the whole 657-partial dual set, so cost grows with particle count, training scales O(N^2). But particle i’s ODE only reads `β_i`, so differentiate each trajectory w.r.t. just `(NN weights, β_i)`, 58 partials, constant in N. This gives the identical gradient and scales O(N).

Measured on the exact NN/ODE/tolerances, one Stage-1 gradient:

| N | joint `AutoForwardDiff` + `Rodas4P` (current) |
| --- | --- |
| 100 | 2.6 s |
| 300 | 17.6 s |
| **600** | **70.7 s** |

vs. the structured gradient below at **N=600: 1.6 s single-threaded, 0.8 s on 8 threads** — ~ **45–90×** , same answer.

```julia-auto
using OrdinaryDiffEqTsit5, Lux, ComponentArrays, ForwardDiff, StableRNGs, StaticArrays
using SciMLBase, Optimization, OptimizationOptimisers, OptimizationOptimJL, LineSearches
using Base.Threads

softplus(x) = max(x, 0.0) + log1p(exp(-abs(x)))
const KON, KOFF, NSSB = 5e-3, 0.1, 10.0

rng = StableRNG(42)
U = Chain(Dense(2, 4, swish), Dense(4, 4, swish), Dense(4, 4, swish), Dense(4, 1, softplus))
θ0, snn = Lux.setup(rng, U)

function cude!(du, u, p, t)
    β = softplus(p.cond)
    Û = U(SVector(u[1], β), p.ude, snn)[1] # SVector -> no allocation in the RHS
    du[1] = KON * Û[1] * (NSSB - u[1]) - KOFF * u[1]
    nothing
end
base_prob = ODEProblem{true, SciMLBase.FullSpecialize}(cude!, [0.0], (0.0, 100.0),
                                                       ComponentVector(ude = θ0, cond = 0.0))

function build_ensemble(p̂)
    prob_func = function (prob, a, args...) # portable across SciMLBase versions
        i = a isa Integer ? a : a.sim_id # classic (prob,i,repeat) or new (prob,ctx)
        remake(prob, p = ComponentVector(ude = p̂.ude, cond = p̂.conditional[i]))
    end
    EnsembleProblem(base_prob; prob_func, safetycopy = false)
end

function loss(p̂, (ts, y))
    N = size(y, 2)
    sim = solve(build_ensemble(p̂), Tsit5(), EnsembleThreads();
                trajectories = N, saveat = ts, abstol = 1e-6, reltol = 1e-6)
    s = 0.0
    @inbounds for i in 1:N
        s += sum(abs2, @view(y[:, i]) .- @view(Array(sim.u[i])[1, :]))
    end
    return s / N
end

function loss_i(v, ts, yi)                          
    sol = solve(remake(base_prob; tspan = (ts[1], ts[end]),
                       p = ComponentVector(ude = v.ude, cond = v.cond)),
                Tsit5(); saveat = ts, abstol = 1e-6, reltol = 1e-6)
    return sum(abs2, yi .- @view(Array(sol)[1, :]))
end

function grad!(G, p̂, (ts, y))
    N, nθ = size(y, 2), length(p̂.ude)
    ude_cols = zeros(nθ, N)                           
    @threads for i in 1:N
        v = ComponentVector(ude = p̂.ude, cond = p̂.conditional[i])
        gi = ForwardDiff.gradient(w -> loss_i(w, ts, @view(y[:, i])), v)
        @views ude_cols[:, i] .= gi.ude
        G.conditional[i] = gi.cond
    end
    G.ude .= vec(sum(ude_cols, dims = 2))
    G ./= N
    return G
end

function fit(p0, ts, y; adam_iters = 300, lbfgs_iters = 200)
    data = (ts, y)
    objective(u, p = nothing) = loss(u, data)
    gradient!(G, u, p = nothing) = grad!(G, u, data)
    optf = OptimizationFunction(objective; grad = gradient!) # NoAD: uses our gradient

    res = solve(OptimizationProblem(optf, p0), OptimizationOptimisers.Adam(0.01); maxiters = adam_iters)
    res = solve(OptimizationProblem(optf, res.u),
                Optim.LBFGS(linesearch = LineSearches.BackTracking()); maxiters = lbfgs_iters)
    return res
end

```

Run with `julia -t auto` (or set `JULIA_NUM_THREADS`) so both the ensemble solve and the gradient parallelize over particles.

Two more that compound on top: in Stage 2 you refine all 10 candidate models but only `stage2_params[best_model_idx]` is ever used, refine just the selected one for ~10x; and Stage 2 runs ADAM `maxiters=1000` then LBFGS `maxiters=1000` to fit a _single scalar_ `β` per particle, where a handful of iterations suffices, so drop the ADAM phase (or cap it low) there.

---

<div class="post-metadata">

**Author:** ![NCCP](https://avatars.discourse-cdn.com/v4/letter/n/c5a1d2/32.png) [@NCCP](https://discourse.julialang.org/u/NCCP)\
**Post date:** [July 13, 2026, 3:35pm UTC](https://discourse.julialang.org/t/how-can-i-significantly-speed-up-conditional-universal-differential-equation-training/138126/3 "2026-07-13T15:35:27Z")

</div>

Hi Chris,

Thank’s a lot for your suggestions!  
That did it for me!

Kind regards,  
NCCP