# Adding external parameters to neural network

**URL:** <https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968>\
**Category:** Machine Learning\
**Tags:** optimization, neural-network, lux\
**Created:** [April 15, 2024, 3:05pm UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968 "2024-04-15T15:05:50Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![KianH](https://avatars.discourse-cdn.com/v4/letter/k/c6cbf5/32.png) [@KianH](https://discourse.julialang.org/u/KianH)\
**Post date:** [April 15, 2024, 3:05pm UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968/1 "2024-04-15T15:05:50Z")

</div>

Is there any way in Lux.jl or Flux.jl to create a function for a UDE and add some parameter, for instance the parameter `c` in the function `rhs!(...)` to a neural network so that the parameter `c` is also estimated through the neural network?

```julia
function rhs!(du, u, p, t)
    
    û = dudt2(u, p, st)[1]
    du[1] = û[1] + log.(c./u[1])
    du[2] = û[2] 
    du[3] = û[3]

end

```

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 15, 2024, 3:12pm UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968/2 "2024-04-15T15:12:20Z")

</div>

yes you can. Set it up like [https://lux.csail.mit.edu/stable/tutorials/intermediate/1\_NeuralODE#Define-the-Neural-ODE-Layer](https://lux.csail.mit.edu/stable/tutorials/intermediate/1_NeuralODE#Define-the-Neural-ODE-Layer). Here is a pseudocode

```julia
# `c` will have to be an array for it to be trainable
my_custom_model = @compact(; model, c = [2.0]) do x, ps
   # Note that state handling is automatic with `@compact`
    function rhs!(du, u, p, t)
       û = model(u, p.model) # Parameters accessible via p.<fieldname> 
       du[1] = û[1] + log.(p.c[1]./u[1]) # p.c[1] since we set it up as an array
       du[2] = û[2] 
       du[3] = û[3]
    end

    prob = ....
    return solve(prob, ....)
end

ps, st = Lux.setup(rng, my_custom_model)
# ps.c and ps.model are automatically populated

my_custom_model(x_input, ps, st)

```

There is a new tutorial that is going to be merged this week, which demonstrates how to do UDEs with a clean syntax [Solving Optimal Control Problems with Symbolic Universal Differential Equations | Lux.jl Documentation](https://lux.csail.mit.edu/previews/PR585/tutorials/advanced/2_SymbolicOptimalControl#Training-a-Neural-Network-based-UDE)

---

<div class="post-metadata">

**Author:** ![KianH](https://avatars.discourse-cdn.com/v4/letter/k/c6cbf5/32.png) [@KianH](https://discourse.julialang.org/u/KianH)\
**Post date:** [April 16, 2024, 9:29am UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968/3 "2024-04-16T09:29:47Z")

</div>

Thank you for the pseudocode.

I have some follow-up questions: I am getting an error when calling:

`my_custom_model(x_input, p, st)`

That states `type NullParameters has no field model`. Note that in my case:

```julia
 model = Lux.Chain(Lux.Dense(3, 8, swish), Lux.Dense(8, 3))

```

I also have a function that updates the network parameter values as shown below:

```julia
function predict_node(θ)
    prob = ODEProblem{true}(my_custom_model, u0, tspan)
    _prob = remake(prob, p = θ)
    Array(solve(_prob, Vern7(), saveat = tsteps, sensealg = InterpolatingAdjoint(; autojacvec = ZygoteVJP())))
end

```

I am wondering whether it is possible to output the extended model in the function `my_custom_model` rather than the solution to the model?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 16, 2024, 2:11pm UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968/4 "2024-04-16T14:11:02Z")

</div>

The custom model code I shared already encapsulates the odeproblem so you don’t need to define the predict function ( you effectively created a neural ode inside a neural ode iiuc)

The null params error is because you did not specify parameters p to the odeproblem struct

---

<div class="post-metadata">

**Author:** ![KianH](https://avatars.discourse-cdn.com/v4/letter/k/c6cbf5/32.png) [@KianH](https://discourse.julialang.org/u/KianH)\
**Post date:** [April 24, 2024, 2:15pm UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968/5 "2024-04-24T14:15:13Z")

</div>

It seems that when I give structure to the RHS the performance becomes significantly worse rather than using a pure neural ODE. It has a lot of difficulties estimating the parameter `c` and is largely dependent on the initialization. Is there a way to train the network and the estimation of the parameter `c` separately?

---

<div class="post-metadata">

**Author:** ![avikpal](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/avikpal/32/6550_2.png) [@avikpal](https://discourse.julialang.org/u/avikpal)\
**Post date:** [April 24, 2024, 3:03pm UTC](https://discourse.julialang.org/t/adding-external-parameters-to-neural-network/112968/6 "2024-04-24T15:03:41Z")

</div>

For this problem, you can make `c` a scalar, so that it doesn’t get treated as a trainable parameters.

Unfortunately there isn’t a nice way to mark it as “fixed” any other way. I had created an issue on this [Allow "const" arrays as inputs to `@compact` · Issue #588 · LuxDL/Lux.jl · GitHub](https://github.com/LuxDL/Lux.jl/issues/588) but haven’t gotten around to fixing it yet