Two changes do the heavy lifting here, and neither changes the result.
-
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.)
-
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.
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.