Esperimenti con Neural ODEs in Julia

Neural Ordinary Differential Equations (abbreviato Neural ODEs) è un paper che introduce una nuova famiglia di reti neurali in cui alcuni strati nascosti (o anche l'unico strato nei casi più semplici) sono implementati con un risolutore di equazioni differenziali ordinarie.
Questo post mostra un esempio scritto in Julia (in futuro ce ne saranno altri) che utilizza alcune idee descritte nel paper Neural ODEs per risolvere un problema di approssimazione della mappatura tra un input e un output in uno preciso scenario.

Tutti i vari frammenti di codice descritti in questo post richiedono la versione 1.5.3 di Julia e i seguenti package: DifferentialEquations, Flux, DiffEqFlux, Plots.
Per ottenere il codice si veda il paragrafo Download del codice completo in fondo a questo post.
Se si fosse interessati a vedere la soluzione degli stessi problemi in Python con TensorFlow si veda il post Esperimenti con Neural ODEs in Python con TensorFlowDiffEq su questo sito web.

Convenzioni

In questo post le convenzioni adoperate sono le seguenti:

  • $t$ è la variabile indipendente
  • $x$ è la funzione incognita
  • $y$ è la seconda funzione incognita
  • $x$ e $y$ sono da intendersi funzioni di $t$, quindi $x=x(t)$ e $y=y(t)$, ma l'uso di questa notazione compatta, oltre ad avere una maggiore leggibilità a livello matematico rende più agevole la "traduzione" in codice dell'equazione
  • $x'$ è la derivata prima di x rispetto a $t$ e naturalmente $y'$ è la derivata prima di y rispetto a $t$


Esperimento #1: addestrare un sistema di ODE per soddisfare un obiettivo

Un percettrone multistrato (abbreviato MLP) è uno strumento opportuno per imparare una relazione non lineare tra input e output di cui non si conosce la legge.
Ci sono casi invece in cui si ha conoscenza a priori della legge che correla gli input e gli output, ad esempio nella forma di un sistema parametrico di equazioni differenziali: in questa situazione una rete neurale di tipo MLP non consente di utilizzare tale conoscenza mentre una rete di tipo Neural ODEs sì.

Lo scenario di applicazione

Lo scenario di applicazione è il seguente:

    Conoscenza a priori:
    • Un dataset che contenga gli input e gli output
    • Una legge che associ input e output in forma di sistema parametrico di equazioni differenziali
    Obiettivo:
    • Determinare opportuni valori dei parametri affiché il sistema ottenuto sostituendo i parametri formali con i valori determinati approssimi al meglio la mappatura tra input e output.
    Lo stesso scenario è stato applicato nel post DiffEqFlux.jl – A Julia Library for Neural Differential Equations sul blog ufficiale di Julia ove la legge è il sistema di equazioni di Lotka-Volterra che descrive la dinamica della popolazione di prede e predatori.
    In questo post la legge non è volutamente una legge famosa, ma è il sistema parametrico di due equazioni del primo ordine e otto parametri mostrato nel seguente paragrafo.

    Il problema da risolvere

    Sia dato il seguente sistema parametrico di due equazioni differenziali ordinarie con valori iniziali che rappresenta la legge che descrive il comportamento di un ipotetico sistema dinamico: $$ \begin{equation} \begin{cases} x' = a_1x + b_1y + c_1e^{-d_1t} \\ y'= a_2x + b_2y + c_2e^{-d_2t} \\ x(0)=0 \\ y(0)=0 \end{cases} \end{equation} $$ Ovviamente questa è una demo il cui scopo è provare la bontà del metodo, quindi per preparare il dataset fissiamo arbitrariamente gli otto valori dei parametri, ad esempio questi: $$ \left[\begin{matrix}a_1 \\ b_1 \\ c_1 \\ d_1 \\ a_2 \\ b_2 \\ c_2 \\ d_2 \end{matrix} \right] = \left[\begin{matrix}1.11 \\ 2.43 \\ -3.66 \\ 1.37 \\ 2.89 \\ -1.97 \\ 4.58 \\ 2.86 \end{matrix} \right] $$ e sapendo che con tali valori dei parametri la soluzione analitica è la seguente: $$ \begin{equation} \begin{array}{l} x(t) = \\ \;\; -1.38778 \cdot 10^{-17} \; e^{-8.99002 t} - \\ \;\; 2.77556 \cdot 10^{-17} \; e^{-7.50002 t} + \\ \;\; 3.28757 \; e^{-3.49501 t} - \\ \;\; 3.18949 \; e^{-2.86 t} + \\ \;\; 0.258028 \; e^{-1.37 t} - \\ \;\; 0.356108 \; e^{2.63501 t} + \\ \;\; 4.44089 \cdot 10^{-16} \; e^{3.27002 t} + \\ \;\; 1.11022 \cdot 10^{-16} \; e^{4.76002 t} \\ \\ y(t) = \\ \;\; -6.23016 \; e^{-3.49501 t} + \\ \;\; 5.21081 \; e^{-2.86 t} + \\ \;\; 1.24284 \; e^{-1.37 t} - \\ \;\; 0.223485 \; e^{2.63501 t} + \\ \;\; 2.77556 \cdot 10^{-17} \; e^{4.76002 t} \\ \end{array} \end{equation} $$ (verificabile online tramite Wolfram Alpha) si è in grado di preparare il dataset: l'input è un intervallo discretizzato del tempo da $0$ a $1.5$ passo $0.01$, mentre l'output è costituito dalle soluzioni analitiche $x=x(t)$ e $y=y(t)$ per ogni $t$ appartenente all'input.
    Una volta preparato il dataset ci si dimentichi dei valori dei parametri e della soluzione analitica e ci si pone il problema di come addestrare una rete neurale per determinare un opportuno set di valori per gli otto parametri per approssimare al meglio la mappatura non lineare tra input e output del dataset.

    L'implementazione della soluzione

    Il sistema parametrico di equazioni differenziali è già scritto in forma esplicita e in Julia con DifferentialEquations si implementa così:

    function parametric_ode_system!(du,u,p,t)
      x, y = u
      a1, b1, c1, d1, a2, b2, c2, d2 = p
      du[1] = dx = a1*x + b1*y + c1*exp(-d1*t)
      du[2] = dy = a2*x + b2*y + c2*exp(-d2*t)
    end
    I setting sono:
    • Input: $t \in [0, 1.5]$ passo di discretizzazione $0.01$
    • Condizioni al contorno: $x(0)=0; y(0)=0$
    • Valori iniziali del parametri qualsiasi; per comodità qui si impostano tutti uguali a $1$
    In codice:

    tbegin=0.0
    tend=1.5
    tstep=0.01
    trange = tbegin:tstep:tend
    u0 = [0.0,0.0]
    tspan = (tbegin,tend)
    p = ones(8)
    La rete neurale è costituita da un solo strato di tipo ODE solver:

    prob = ODEProblem(parametric_ode_system!, u0, tspan, p)
    
    function net()
        solve(prob, Tsit5(), p=p, saveat=trange)
    end

    Le impostazioni dell'addrestramento sono:
    • Ottimizzatore: ADAM
    • Learning rate: $0.05$
    • Numero di epoche: $1000$
    • Funzione di loss: somma dei quadrati delle differenze
    In codice:

    epochs = 1000
    learning_rate = 0.05
    data = Iterators.repeated((), epochs)
    opt = ADAM(learning_rate)
    callback_func = function ()
      #.......
    end
    fparams = Flux.params(p)
    Flux.train!(loss_func, fparams, data, opt, cb=callback_func)

    Qui di seguito il codice completo:

    using Flux, DiffEqFlux, DifferentialEquations, Plots
    
    function parametric_ode_system!(du,u,p,t)
      x, y = u
      a1, b1, c1, d1, a2, b2, c2, d2 = p
      du[1] = dx = a1*x + b1*y + c1*exp(-d1*t)
      du[2] = dy = a2*x + b2*y + c2*exp(-d2*t)
    end
    
    true_params = [1.11, 2.43, -3.66, 1.37, 2.89, -1.97, 4.58, 2.86]
    
    an_sol_x(t) =
      -1.38778e-17 * exp(-8.99002 * t) -
      2.77556e-17 * exp(-7.50002 * t) +
      3.28757 * exp(-3.49501 * t) -
      3.18949 * exp(-2.86 * t) +
      0.258028 * exp(-1.37 * t) -
      0.356108 * exp(2.63501 * t) +
      4.44089e-16 * exp(3.27002 * t) +
      1.11022e-16 * exp(4.76002 * t)
    an_sol_y(t) =
      -6.23016 * exp(-3.49501 * t) +
      5.21081 * exp(-2.86 * t) +
      1.24284 * exp(-1.37 * t) -
      0.223485 * exp(2.63501 * t) +
      2.77556e-17 * exp(4.76002 * t)
    
    tbegin=0.0
    tend=1.5
    tstep=0.01
    trange = tbegin:tstep:tend
    u0 = [0.0,0.0]
    tspan = (tbegin,tend)
    p = ones(8)
    
    prob = ODEProblem(parametric_ode_system!, u0, tspan, p)
    
    function net()
        solve(prob, Tsit5(), p=p, saveat=trange)
    end
    
    dataset_outs = [an_sol_x.(trange), an_sol_y.(trange)]
    function loss_func()
      pred = net()
      sum(abs2, dataset_outs[1] .- pred[1,:]) +
      sum(abs2, dataset_outs[2] .- pred[2,:])
    end
    
    epochs = 1000
    learning_rate = 0.05
    data = Iterators.repeated((), epochs)
    opt = ADAM(learning_rate)
    callback_func = function ()
      println("loss: ", loss_func())
    end
    fparams = Flux.params(p)
    Flux.train!(loss_func, fparams, data, opt, cb=callback_func)
    
    predict_prob = ODEProblem(parametric_ode_system!, u0, tspan, p)
    predict_sol = solve(prob, Tsit5(), saveat=trange)
    x_predict_sol = [u[1] for u in predict_sol.u]
    y_predict_sol = [u[2] for u in predict_sol.u]
    
    println("Learned parameters:", p)
    
    plot(trange, dataset_outs[1],
        linewidth=2, ls=:dash,
        title="Neural ODEs to fit params",
        xaxis="t",
        label="dataset x(t)",
        legend=true)
    plot!(trange, dataset_outs[2],
        linewidth=2, ls=:dash,
        label="dataset y(t)")
    plot!(predict_sol.t, x_predict_sol,
        linewidth=1,
        label="predicted x(t)")
    plot!(predict_sol.t, y_predict_sol,
        linewidth=1,
        label="predicted y(t)")
    
    Qui il link al codice su GitHub.

    I valori dei parametri ottenuti al termine dell'addrestramento sono i seguenti: $$ \left[\begin{matrix}a_1 \\ b_1 \\ c_1 \\ d_1 \\ a_2 \\ b_2 \\ c_2 \\ d_2 \end{matrix} \right] = \left[\begin{matrix}1.7302980833638142 \\ 1.2823312512074032 \\ -1.6866178290795755 \\ 0.41974163099782325 \\ 1.223075467559363 \\ 0.9410722500584323 \\ 0.18890958911958686 \\ 1.7462909509457183 \end{matrix} \right] $$ che sono ovviamente diversi dai valori dei parametri utilizzati per generare il dataset (passando tramite la soluzione analitica); però il sistema di equazioni ottenuto sostituendo i parametri formali con tali valori ottenuti dall'addestramento della rete neurale e risolvendolo numericamente questo ultimo sistema si ottiene una soluzione numerica che approssima abbastanza bene il dataset, come mostrato dalla figura:

    Comparazione della soluzione numerica che approssima il dataset.

    Differenti strategie di addestramento consentono di ottenere diversi valori dei parametri e quindi una diversa soluzione numerica del sistema con una differente accuratezza.

    Citazioni

    @articleDBLP:journals/corr/abs-1806-07366,
      author    = {Tian Qi Chen and
                   Yulia Rubanova and
                   Jesse Bettencourt and
                   David Duvenaud},
      title     = {Neural Ordinary Differential Equations},
      journal   = {CoRR},
      volume    = {abs/1806.07366},
      year      = {2018},
      url       = {http://arxiv.org/abs/1806.07366},
      archivePrefix = {arXiv},
      eprint    = {1806.07366},
      timestamp = {Mon, 22 Jul 2019 14:09:23 +0200},
      biburl    = {https://dblp.org/rec/journals/corr/abs-1806-07366.bib},
      bibsource = {dblp computer science bibliography, https://dblp.org}
    
    @articleDBLP:journals/corr/abs-1902-02376,
      author    = {Christopher Rackauckas and
                   Mike Innes and
                   Yingbo Ma and
                   Jesse Bettencourt and
                   Lyndon White and
                   Vaibhav Dixit},
      title     = {DiffEqFlux.jl - {A} Julia Library for Neural Differential Equations},
      journal   = {CoRR},
      volume    = {abs/1902.02376},
      year      = {2019},
      url       = {https://arxiv.org/abs/1902.02376},
      archivePrefix = {arXiv},
      eprint    = {1902.02376},
      timestamp = {Tue, 21 May 2019 18:03:36 +0200},
      biburl    = {https://dblp.org/rec/bib/journals/corr/abs-1902-02376},
      bibsource = {dblp computer science bibliography, https://dblp.org}
    

    Download del codice completo

    Il codice completo è disponibile su GitHub.
    Questo materiale è distribuito su licenza MIT; sentiti libero di usare, condividere, "forkare" e adattare tale materiale come credi.
    Sentiti anche libero di pubblicare pull-request e bug-report su questo repository di GitHub oppure di contattarmi sui miei canali social disponibili nell'angolo in alto a destra di questa pagina.