RasmahRasmah

Neural Networks: A Differentiable MLP from Scratch

println("Julia ", VERSION)
Julia 1.12.6

Overview

Rasmah includes a native, dependency-light neural-network core: a multi-layer perceptron, a set of activation functions, weight initializers, a native Adam optimizer, and a training loop. The design is deliberately minimal so the network plays the same role as every other Rasmah quantity: a differentiable program.

A network’s learnable parameters are stored as a single flat vector θ, and the forward pass is the pure function apply(net, θ, x). Training therefore reduces to minimizing a scalar loss loss(θ) with automatic differentiation — the exact same engine that differentiates geometry and the finite-element solver (forward mode by default; reverse mode if one is loaded). This is what lets the same core power implicit neural surfaces, physics-informed networks, and neural ODEs later.

The multi-layer perceptron

MLP(sizes) builds a feed-forward network whose sizes list gives the width of each layer. The learnable parameters are one flat vector, so a network is “just a function of θ”:

using Rasmah
using Random

net = MLP([2, 16, 16, 1])          # 2 inputs → 16 → 16 → 1 output
println("layers  = ", nlayers(net))
println("params  = ", nparams(net)) # weights + biases, flattened
layers  = 3
params  = 337

apply runs the forward pass. It takes a single input vector, or a batched matrix with samples as columns (in × batch):

θ0 = he_init(MersenneTwister(1), [2, 16, 16, 1])
println(apply(net, θ0, [1.0, -1.0]; act=relu, out_act=identity))
[-5.097698975495556]

Activations and initializers

The element-wise activations live as plain functions, and the initializers return a correctly-scaled flat parameter vector for a given sizes:

println(relu(-1.0), " ", sigmoid(0.0), " ", gelu(1.0))   # 0.0 0.5 ~0.84
println(length(xavier_init(MersenneTwister(0), [2, 16, 1])))
println(length(siren_init(MersenneTwister(0), [2, 16, 1]; ω₀=30.0)))
0.0 0.5 0.8411919906082768
65
65

Training with Adam

To learn, write a scalar loss of θ and hand it to train!, which mutates θ in place with Adam and returns the per-epoch loss history. A classic sanity target: regress sin(x) on a handful of points. We record the loss before and after so the convergence is visible:

Random.seed!(0)

model = MLP([1, 8, 1])
θ     = he_init(MersenneTwister(1), [1, 8, 1])

X = reshape(collect(range(-π, π, length=64)), 1, :)   # 1 × 64 (samples as columns)
Y = reshape(sin.(vec(X)), 1, :)

loss(θ) = mse_loss(apply(model, θ, X; out_act=identity), Y)

history = train!(loss, θ; epochs=800, lr=5e-3)
println("loss  ", round(history[1], digits=3), "  →  ", round(history[end], digits=5))
loss  17.914  →  0.02091

The trained network is a differentiable approximation of the function — check it against the analytic value at a fresh point:

ŷ = apply(model, θ, [π/4]; out_act=identity)[1]
println("net(π/4) = ", round(ŷ, digits=4), "   sin(π/4) = ", round(sin(π/4), digits=4))
net(π/4) = 0.7603   sin(π/4) = 0.7071

Beyond the core

This same core, differentiated the same way, is the engine of Rasmah’s neural applications:

  • NeuralSDF / fit_sdf — an implicit signed-distance field backed by an MLP (a SIREN network), so a neural surface plugs into field, gradient, marching_cubes, and CSG exactly like an analytic one.
  • pinn_solve — a physics-informed network that satisfies a PDE residual and boundary conditions as a soft loss.
  • NeuralODE / ude_fit — a neural ordinary differential equation, and a universal differential equation that learns the correction to a known mechanistic model.
  • DeepONet / FNO1D — operator learning, and VAE / DDPM — generative design.

Each is documented in the API reference; they all build on the loss(θ) → train! pattern shown here.

Next steps