Neural networks

Overview

Rasmah includes a small, native neural-network core — no external deep-learning dependency — built around one idea: a network's learnable parameters are a single flat vector $\theta$, and the network is just a differentiable function $f(\theta, x)$ of that vector and the input.

That design choice is what makes the rest of Rasmah possible. Training is then a plain optimization problem, "minimize a scalar loss $\ell(\theta)$", solved with the same gradient machinery Rasmah uses everywhere else. The network can also be wrapped up as geometry (a neural signed distance field) or as a PDE solver (a physics-informed neural network) — the subjects of the next two chapters.

The multi-layer perceptron

The workhorse is the feed-forward multi-layer perceptron (MLP), specified only by its layer sizes:

Rasmah.Neural.MLP — Type

MLP(sizes)

A feed-forward network with layer sizes sizes = [in, h1, …, hk, out]. The learnable parameters are a flat vector of length nparams, initialised by xavier_init / he_init / siren_init and evaluated with apply.

m = MLP([2, 8, 8, 1])
nlayers(m)

Rasmah.Neural.nlayers — Function

nlayers(m::MLP) -> Int

Number of affine layers (length(m.sizes) - 1).

Rasmah.Neural.nparams — Function

nparams(m::MLP) -> Int

Total number of scalar parameters (weights + biases).

nparams(m)

The forward pass is a pure function of the parameters and the input:

Rasmah.Neural.apply — Function

apply(m::MLP, θ, x; act=relu, out_act=identity)

Forward pass for a single input vector x (length m.sizes[1]). Hidden layers use act; the output layer uses out_act. Returns a vector of length m.sizes[end].

apply(m::MLP, θ, X::AbstractMatrix; act=relu, out_act=identity)

Batched forward pass: X is in × batch (samples as columns), returns out × batch.

apply(op::DeepONet, θ, u, y; act=tanh) -> Real

Evaluate the operator G(u)(y): ⟨branch(u; θ), trunk(y; θ)⟩ + b.

apply(net::PODDeepONet, θ, u; act=tanh) -> Vector

Evaluate the operator at input u: the branch produces the POD coefficients and V lifts them to the full n-dimensional field.

apply(fno::FNO1D, θ, u; act=gelu) -> Vector

Apply the FNO to a length-n grid signal u.

apply(fno::FNO2D, θ, u; act=gelu) -> Matrix

Apply the 2D FNO to an n × n grid field u.

apply(fno::FNO3D, θ, u; act=gelu) -> Array

Apply the 3D FNO to an n × n × n grid field u.

apply(lf::LatentField, θ, z; act=relu, out_act=sigmoid) -> Vector

Decode a latent code z into a design field.

θ = he_init(MersenneTwister(0), m.sizes)
apply(m, θ, [1.0, -1.0])

1-element Vector{Float64}:
 0.08196975428302172

Activations

Each hidden layer applies an element-wise nonlinearity. Rasmah provides the common ones:

Rasmah.Neural.relu — Function

relu(x)

Rectified linear unit: max(0, x).

relu(-1.0), relu(2.0)

(0.0, 2.0)

Rasmah.Neural.sigmoid — Function

sigmoid(x)

Logistic sigmoid 1 / (1 + exp(-x)).

round(sigmoid(0.0); digits=4)

0.5

Rasmah.Neural.leaky_relu — Function

leaky_relu(x; α=0.01)

Leaky rectified linear unit: x for x > 0, else α * x.

leaky_relu(-1.0)

-0.01

Rasmah.Neural.gelu — Function

gelu(x)

Gaussian error linear unit (tanh approximation), 0.5 x (1 + tanh(√(2/π)(x + 0.044715 x³))).

round(gelu(0.0); digits=4)

0.0

Rasmah.Neural.swish — Function

swish(x)

Swish / SiLU activation x * sigmoid(x) (Ramachandran et al. 2017).

Initializers

Before training, the parameters need a starting point. The three initializers below return a flat parameter vector with a scheme matched to the activation:

Rasmah.Neural.xavier_init — Function

xavier_init(rng, sizes; gain=1.0, T=Float64) -> Vector{T}

Xavier/Glorot initialization: weights ~ N(0, gain² · 2/(nin+nout)), zero biases.

Rasmah.Neural.he_init — Function

he_init(rng, sizes; gain=√2, T=Float64) -> Vector{T}

Kaiming/He initialization (ReLU): weights ~ N(0, gain² / nin), zero biases.

he_init(MersenneTwister(0), [2, 8, 1]) |> length

Rasmah.Neural.siren_init — Function

siren_init(rng, sizes; ω₀=30.0, T=Float64) -> Vector{T}

SIREN initialization for periodic (sin) networks: first-layer weights ~ U(-1/nin, 1/nin), hidden-layer weights ~ U(-√(6/nin)/ω₀, √(6/nin)/ω₀), zero biases.

Losses

Rasmah.Neural.mse_loss — Function

mse_loss(ŷ, y)

Mean squared error mean(|ŷ - y|²).

mse_loss([1.0, 2.0], [1.5, 1.5])

0.25

Rasmah.Neural.mae_loss — Function

mae_loss(ŷ, y)

Mean absolute error mean(|ŷ - y|).

mae_loss([1.0, 2.0], [1.5, 1.5])

0.5

Training

Training differentiates the loss with an AD backend. neural_backend picks a reverse-mode backend automatically when one is loaded:

Rasmah.Neural.neural_backend — Function

neural_backend() -> ADTypes.AbstractADType

The default AD backend for neural-network training. Prefers reverse-mode: Mooncake if loaded, else Enzyme if loaded, else ForwardDiff (forward mode, for correctness when neither reverse backend is present). Load one of the reverse backends (using Mooncake / using Enzyme) to activate it; the choice is swappable through DifferentiationInterface by passing an explicit backend to train!/fit_sdf.

Training is Adam — an adaptive per-parameter step size that is the de-facto standard:

Rasmah.Neural.Adam — Type

Adam([T], n::Integer)

Adam optimizer state (first/second moment buffers m, v and step counter t) for an n-parameter vector. Adam(n) uses Float64 moments; Adam(T, n) uses element type T (e.g. Float32 so the moments stay in single precision).

Rasmah.Neural.adam_step! — Function

adam_step!(opt::Adam, θ, g; lr=1e-3, β1=0.9, β2=0.999, ϵ=1e-8, weight_decay=0.0) -> θ

One Adam update of θ in place along the gradient g. weight_decay > 0 gives AdamW (decoupled weight decay θ -= lr·wd·θ). Returns the updated θ.

Rasmah.Neural.train! — Function

train!(loss, θ; backend=ad_backend(), epochs=1000, lr=1e-3, …) -> Vector{Float64}

Minimize the scalar loss(θ) over θ (mutated in place) with Adam, returning the per-epoch loss history. Stops early when loss <= atol (for atol > 0).

Keyword arguments forward to adam_step! (lr, β1, β2, ϵ, weight_decay); verbose/verbose_every control per-epoch logging.

A worked example: fitting a sine

Put the pieces together: train an MLP to reproduce $\sin$ on an interval.

xs = reshape(collect(range(-3, 3, length=100)), 1, :)
ys = sin.(xs)

fit = MLP([1, 16, 1])
θfit = he_init(MersenneTwister(0), fit.sizes)
loss(θ) = mse_loss(apply(fit, θ, xs), ys)
history = train!(loss, θfit; epochs=500, lr=1e-2, atol=1e-4)

500-element Vector{Float64}:
 4.199864493983416
 3.456003952307347
 2.809201032627978
 2.2576037786077037
 1.7982348120915423
 1.4264341468940043
 1.1352722025488469
 0.9152425863059293
 0.754385561738421
 0.6392136186490337
 ⋮
 0.012416882310159338
 0.012366306115453364
 0.012315920433601732
 0.012264952411870644
 0.01221388392271712
 0.012162998321742235
 0.012112262174802631
 0.012061733351522962
 0.012011427872426183

history[end]

0.012011427872426183

The final loss is the mean-squared error between the network and $\sin$ over the 100 sample points — a few thousandths after 500 epochs.

Neural ODEs and universal differential equations

A network can also model dynamics. A neural ordinary differential equation is a network whose output is the time derivative $du/dt = \text{net}([u; t])$ — a continuous-depth model solved by a native Runge–Kutta integrator:

Rasmah.Neural.solve_ode — Function

solve_ode(f, u0, tspan, p=nothing; method=RK4(), steps=100) -> ODESolution

Solve du/dt = f(u, p, t) from tspan[1] to tspan[2] with the given Runge–Kutta method and a fixed number of steps. Returns the full trajectory.

For the implicit (stiff) methods BackwardEuler / ImplicitMidpoint, the keyword arguments jac (a Jacobian function jac(u, p, t) -> J, defaulting to AD), linsolve (:direct cholesky, :lu, :amg, :gmres), newtontol, and maxnewton control the Newton solves; the sparse preconditioner's symbolic analysis / hierarchy is built once and reused across the time series.

sol = solve_ode((u, p, t) -> -u, [1.0], (0.0, 1.0); method = RK4(), steps = 10)
sol.u[end]

1-element Vector{Float64}:
 0.3678797744124984

Exponential decay $du/dt = -u$ has the exact solution $e^{-1} \approx 0.36788$; the 10-step RK4 integrator lands on it to seven digits.

Wrapping a network as the vector field is NeuralODE, and fitting one to an observed trajectory is neural_ode_fit:

Rasmah.Neural.NeuralODE — Type

NeuralODE(net, act, sensealg=DiscreteAdjoint())

A neural ODE vector field du/dt = net([u; t]; act), differentiated through its solution with the sensitivity algorithm sensealg. Calling ode(θ, u0, t0, t1; method, steps) integrates the ODE from (t0, u0) to t1 and returns the final state.

Rasmah.Neural.neural_ode_fit — Function

neural_ode_fit(net, θ0, ts, us; …) -> NeuralODEResult

Fit a NeuralODE to an observed trajectory (ts, us) — ts are the time samples and us an n × m matrix of states (columns are time samples; us[:, 1] is the initial condition).

Keywords: act (hidden activation, default tanh), sensealg (sensitivity algorithm, default DiscreteAdjoint), method, steps (integration resolution), epochs, lr, seed, backend (used by DiscreteAdjoint).

A universal differential equation (UDE) augments a known mechanistic model with a learned neural correction, $du/dt = f_{\text{known}}(u, t) + \text{net}([u; t])$ — so the network only has to learn what the model is missing:

Rasmah.Neural.ude_fit — Function

ude_fit(net, θ0, known, ts, us; …) -> NeuralODEResult

Fit a universal differential equation du/dt = known(u, t) + net([u; t]) to the observed trajectory (ts, us) (same convention as neural_ode_fit). known(u, t) is the mechanistic part; the network learns the correction.

Keywords: act, sensealg, method, steps, epochs, lr, seed, backend.

Next steps

A network can do more than fit curves: Neural signed distance fields wrap one up as a geometry, and Physics-informed neural networks use one to solve differential equations.