Universal Differential Equations 1: Learning a Missing Term

A Shared Equinox Correction Inside a Mechanistic Network Model

Try this notebook interactively:

Download .ipynb Download .qmd Open in Colab

Introduction

A universal differential equation (UDE) combines known mechanisms with a trainable function. Instead of asking a neural network to reproduce the complete dynamics, we give it a narrower job: correct a specific part of a mechanistic model that we know is incomplete.

This is the classical residual, or “missing physics,” UDE construction: a known mechanistic vector field is retained while a universal approximator learns only an unknown or misspecified contribution (Rackauckas et al., 2020). In neuroscience, this hybrid formulation provides a bridge between interpretable neural dynamics and flexible data-driven models (El-Gazzar and van Gerven, 2025). Unlike a fully learned Neural ODE, where a neural network parameterizes the complete vector field (Chen et al., 2018), the MLP here augments equations whose known structure remains explicit.

This tutorial constructs a reproducible stochastic four-region FitzHugh–Nagumo network (FitzHugh, 1961; Nagumo et al., 1962). The synthetic teacher uses the complete cubic voltage damping, while the incomplete model contains only half of it. One Equinox multilayer perceptron (MLP) supplies the missing local correction at every region:

\[ \begin{aligned} \mathrm dV_i &= \left[V_i - \alpha\frac{V_i^3}{3} - W_i + I + C_i + g_\theta(V_i, W_i)\right]\mathrm dt + \sigma\,\mathrm dB_i, \\ \mathrm dW_i &= \frac{V_i + a - bW_i}{\tau}\,\mathrm dt, \end{aligned} \]

where \(\alpha=1/2\) in the incomplete model, \(C_i\) is the known network input, and \(B_i\) is an independent Brownian motion for each region. The target correction remains

\[ g^*(V_i) = -(1-\alpha)\frac{V_i^3}{3}. \]

The same \(g_\theta\) is evaluated independently at all four nodes. Their different initial conditions and network inputs expose it to several trajectories through state space, but there is still only one set of trainable MLP weights.

NoteWhere this tutorial sits

This page covers the mechanics end to end: marking a module trainable, getting gradients through the solver, fitting, and checking what was learned. The example is deliberately synthetic so that the correction the MLP should recover is known exactly.

Universal Differential Equations 2 picks those mechanics up and asks the modelling questions instead: where a correction belongs, what to condition it on, how to make a known limit exact by construction rather than by fitting, and what counts as evidence that the right term was learned. Read this one first.

NoteWhy retain half of the cubic term?

Removing all cubic damping makes the incomplete FitzHugh–Nagumo system unbounded before an initially untrained MLP can correct it. Retaining a known stabilizing component is both numerically safer and a more realistic UDE setup: the mechanistic model is approximate rather than structurally unstable.

import equinox as eqx
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np
import optax
import time

plt.rcParams.update({"figure.dpi": 200, "savefig.dpi": 200})
jax.config.update("jax_enable_x64", True)

from tvboptim.experimental.network_dynamics import (
    Bunch,
    DenseGraph,
    Network,
    prepare,
)
from tvboptim.experimental.network_dynamics.coupling import LinearCoupling
from tvboptim.experimental.network_dynamics.dynamics import AbstractDynamics
from tvboptim.experimental.network_dynamics.noise import AdditiveNoise
from tvboptim.experimental.network_dynamics.solvers import Heun
from tvboptim.optim import OptaxOptimizer
from tvboptim.types import (
    EquinoxParameter,
    combine_state,
    partition_state,
)

Complete and Incomplete Dynamics

The teacher contains the complete cubic term. It is used only to generate a small synthetic data set whose ground truth is known.

class TeacherFitzHughNagumo(AbstractDynamics):
    STATE_NAMES = ("V", "W")
    INITIAL_STATE = (-1.0, -0.5)
    DEFAULT_PARAMS = Bunch(a=0.7, b=0.8, tau=12.5, I=0.5)
    COUPLING_INPUTS = {"structural": 1}

    def dynamics(self, t, state, params, coupling, external):
        del t, external
        V, W = state
        network_input = coupling.structural[0]
        dV = V - V**3 / 3.0 - W + params.I + network_input
        dW = (V + params.a - params.b * W) / params.tau
        return jnp.stack((dV, dW))

The UDE retains half of the cubic damping and calls the same correction module once per node. TVB-Optim presents the state to dynamics in state-major form, [states, nodes], so both V and W below have shape [nodes]. The Equinox MLP, however, describes the correction at one node: with in_size=2, it expects one feature vector [V_i, W_i]. Stacking along the last axis therefore builds a [nodes, 2] array whose rows are the inputs for the individual nodes. For this two-state example, state.T would have the same values; constructing node_features explicitly documents which states the learned term sees and in which order, and generalizes naturally to selected or derived features.

jax.vmap then applies the same MLP to every row of node_features. This is shared-weight evaluation, not four copies of the MLP. Because the MLP was created with out_size=1, one evaluation returns shape [1] rather than a JAX scalar. The mapped result consequently has shape [nodes, 1], and [:, 0] removes that singleton output dimension to obtain the [nodes] correction required by base_drift. Without this indexing, adding [nodes, 1] to [nodes] would trigger broadcasting and produce an unintended two-dimensional array.

class UDEFitzHughNagumo(AbstractDynamics):
    STATE_NAMES = ("V", "W")
    INITIAL_STATE = (-1.0, -0.5)
    DEFAULT_PARAMS = Bunch(
        a=0.7,
        b=0.8,
        tau=12.5,
        I=0.5,
        cubic_fraction=0.5,
        correction=None,
    )
    COUPLING_INPUTS = {"structural": 1}

    def dynamics(self, t, state, params, coupling, external):
        del t, external
        V, W = state
        network_input = coupling.structural[0]

        node_features = jnp.stack((V, W), axis=-1)  # [nodes, 2]
        learned_term = jax.vmap(params.correction)(node_features)[:, 0]

        base_drift = (
            V
            - params.cubic_fraction * V**3 / 3.0
            - W
            + params.I
            + network_input
        )
        dW = (V + params.a - params.b * W) / params.tau
        return jnp.stack((base_drift + learned_term, dW))

Four Coupled Regions, One Shared Correction

We use a small directed graph. TVB-Optim stores dense connectivity as weights[target, source], so each row lists the incoming connection strengths for one target region. The graph is asymmetric to give the four otherwise identical regions different network contexts.

weights = jnp.array(
    [
        [0.0, 0.8, 0.0, 0.2],
        [0.1, 0.0, 0.7, 0.0],
        [0.4, 0.0, 0.0, 0.6],
        [0.0, 0.3, 0.2, 0.0],
    ]
)

initial_state = jnp.array(
    [
        [-1.2, -0.8, 0.1, 0.7],  # V
        [-0.6, -0.2, 0.3, 0.5],  # W
    ]
)

DT = 0.05
T1 = 15.0
COUPLING_GAIN = 0.12
NOISE_SIGMA = 0.08
TRAIN_NOISE_KEY = jax.random.key(11)
N_NODES = weights.shape[0]
REGION_LABELS = tuple(f"R{node + 1}" for node in range(N_NODES))

Figure 1: Four coupled regions with one shared local correction. Directed edges carry the mechanistic network input. Every node evaluates the same MLP \(g_\theta(V,W)\) with shared weights; the MLP is not replicated or fitted independently per region.

Generate Reproducible Noisy Training Data

Both networks use the same graph, coupling, initial conditions, solver, and mechanistic parameters. A small additive process noise acts only on voltage. Their only intentional structural difference is the missing half of the cubic voltage damping.

def prepare_network(dynamics, noise_key=TRAIN_NOISE_KEY, t1=T1):
    network = Network(
        dynamics=dynamics,
        coupling={
            "structural": LinearCoupling(source="V", G=COUPLING_GAIN)
        },
        graph=DenseGraph(weights),
        noise=AdditiveNoise(
            sigma=NOISE_SIGMA,
            apply_to="V",
            key=noise_key,
        ),
    )
    solve_fn, config = prepare(
        network,
        Heun(),
        t0=0.0,
        t1=t1,
        dt=DT,
    )
    config.initial_state.dynamics = initial_state
    return solve_fn, config


teacher_solve, teacher_config = prepare_network(TeacherFitzHughNagumo())
target_solution = teacher_solve(teacher_config)

assert target_solution.variable_names == ("V", "W")
assert target_solution.ys.shape == (int(T1 / DT), 2, N_NODES)
assert jnp.all(jnp.isfinite(target_solution.ys))

The teacher and UDE use the same fixed random key. This common random numbers design makes the example reproducible and prevents the optimizer from comparing different random forcing on every evaluation. The noise realization is part of the synthetic training data; only the MLP weights change during fitting.

This is a software demonstration with synthetic observations, not evidence that the learned term is a biologically identified mechanism. Its advantage is that we know exactly what the correction should recover.

Add an Equinox Parameter

We initialize a small MLP and set its final layer to zero. The first incomplete rollout therefore contains no learned correction and exposes the mechanistic model’s misspecification directly.

def make_correction_module(key, in_size, out_size, width_size=16, depth=2):
    """Small MLP with a zeroed final layer, so training starts from no correction.

    The zeroed final layer makes the correction identically zero at step zero, so
    the first rollout is the mechanistic model on its own.
    """
    module = eqx.nn.MLP(
        in_size=in_size,
        out_size=out_size,
        width_size=width_size,
        depth=depth,
        activation=jax.nn.tanh,
        key=key,
    )
    final_layer = module.layers[-1]
    return eqx.tree_at(
        lambda model: (model.layers[-1].weight, model.layers[-1].bias),
        module,
        (jnp.zeros_like(final_layer.weight), jnp.zeros_like(final_layer.bias)),
    )

correction = EquinoxParameter(
    make_correction_module(jax.random.key(7), in_size=2, out_size=1)
)
ude_solve, ude_config = prepare_network(
    UDEFitzHughNagumo(correction=correction)
)

print(f"Correction wrapper: {type(ude_config.dynamics.correction).__name__}")
print(f"Shared module: {type(ude_config.dynamics.correction.module).__name__}")

Correction wrapper: EquinoxParameter
Shared module: MLP

EquinoxParameter marks the MLP’s inexact array leaves as trainable. Its activation functions, architecture, and other non-array leaves remain static. The wrapper itself is callable, so the dynamics does not need to unwrap it.

Because an Equinox module contains callable leaves, a full configuration should cross the compilation boundary through eqx.filter_jit:

compiled_solve = eqx.filter_jit(ude_solve)

initial_solution = ude_solve(ude_config)
compiled_initial = compiled_solve(ude_config)

assert jnp.allclose(compiled_initial.ys, initial_solution.ys)
assert jnp.all(jnp.isfinite(initial_solution.ys))

Figure 2: Effect of the missing nonlinear damping under shared random forcing. The complete teacher trajectory (black dashed) and the stable but misspecified mechanistic UDE before training (colored) diverge at every region. Different initial states and directed coupling give the shared correction several dynamical contexts to learn from.

Verify Gradient Flow

Partitioning separates the MLP arrays from the fixed graph, solver state, mechanistic parameters, and callable activation leaves. A direct gradient through the rollout confirms that every trainable leaf receives finite signal.

target_V = target_solution.ys[:, 0, :]

def loss(config):
    predicted_V = ude_solve(config).ys[:, 0, :]
    return jnp.mean((predicted_V - target_V) ** 2)


diff_config, static_config = partition_state(ude_config)

def partitioned_loss(diff):
    return loss(combine_state(diff, static_config))


gradient = jax.grad(partitioned_loss)(diff_config)
gradient_leaves = jax.tree.leaves(gradient)
gradient_norm = jnp.sqrt(sum(jnp.sum(leaf**2) for leaf in gradient_leaves))

assert gradient_leaves
assert all(jnp.all(jnp.isfinite(leaf)) for leaf in gradient_leaves)
assert gradient_norm > 0.0

print(f"Trainable array leaves: {len(gradient_leaves)}")
print(f"Initial gradient norm: {float(gradient_norm):.3e}")

Trainable array leaves: 6
Initial gradient norm: 1.460e+01

Fit the Shared Neural Correction

OptaxOptimizer uses the same API for ordinary Parameter values and EquinoxParameter modules. We record one loss value per ten-step optimizer chunk to keep the notebook output and Python overhead small.

compiled_loss = eqx.filter_jit(loss)
initial_loss = float(compiled_loss(ude_config))
loss_steps = [0]
loss_values = [initial_loss]

def record_loss(step, diff, static, fitting_data, aux, loss_value, grads):
    del fitting_data, aux, grads
    loss_steps.append(int(step) + 1)
    loss_values.append(float(loss_value))
    return False, diff, static

optimizer = OptaxOptimizer(
    loss,
    optax.adam(learning_rate=3e-3),
    callback=record_loss,
)
fitted_config, _ = optimizer.run(
    ude_config,
    max_steps=500,
    chunk_size=10,
)

final_loss = float(compiled_loss(fitted_config))
loss_values[-1] = final_loss
fitted_solution = compiled_solve(fitted_config)

assert jnp.all(jnp.isfinite(fitted_solution.ys))
assert final_loss < initial_loss * 0.05

Figure 3: Training the shared correction. Left: trajectory loss decreases as gradients pass through the complete network rollout. Right: after fitting, one shared MLP reproduces the voltage trajectories at all four coupled regions.

What Did the MLP Learn?

The target residual depends only on voltage, although the MLP receives both V and W. We evaluate it at every state visited by the teacher and compare it with the known missing half-cubic term. We also differentiate the learned correction with respect to W at those same states. This partial derivative holds V fixed, so the true missing term has zero W sensitivity even though V and W are correlated along a trajectory.

V_observed = target_solution.ys[:, 0, :].reshape(-1)
W_observed = target_solution.ys[:, 1, :].reshape(-1)
observed_features = jnp.stack((V_observed, W_observed), axis=-1)

learned_residual = jax.vmap(fitted_config.dynamics.correction)(
    observed_features
)[:, 0]
true_residual = -(1.0 - fitted_config.dynamics.cubic_fraction) * V_observed**3 / 3.0
residual_correlation = np.corrcoef(
    np.asarray(learned_residual), np.asarray(true_residual)
)[0, 1]

def correction_scalar(features):
    return fitted_config.dynamics.correction(features)[0]


correction_gradient = jax.grad(correction_scalar)
w_sensitivity = jax.vmap(correction_gradient)(observed_features)[:, 1]
w_sensitivity_rms = jnp.sqrt(jnp.mean(w_sensitivity**2))

Figure 4: Learned missing dynamics and dependence on the slow state. Left: each point evaluates the shared MLP at a state visited during training; color identifies the region that supplied the state, and the black curve is the known missing half-cubic term. Right: the local partial derivative of the correction with respect to \(W\) at the same observed states. The true residual is independent of \(W\), so its sensitivity is zero (black dashed line). Nonzero sensitivity indicates that the MLP uses correlated \(W\) information to reproduce the trajectories. Both diagnostics apply only along the observed state manifold.

The nonzero sensitivity shows that the fitted MLP does use W, despite the true missing term depending only on V. This is compatible with an accurate trajectory fit because the observed states occupy a narrow, correlated part of the phase plane. Recovering the intended mechanism would require stronger information, such as restricting the correction’s inputs, sampling more of the phase plane, or regularizing unwanted state dependence.

That is a general diagnostic, not a quirk of this example. Whenever the true term depends on fewer variables than the network is handed, the surplus dependence is measurable, and it is the cheapest sign that a good trajectory fit is not the same thing as a recovered mechanism. Tutorial 2 measures the same thing a different way, as the width of a loop in the learned residual.

Validation Initial Conditions

The fitted module has seen four related trajectories, not the entire phase plane. We therefore rerun both models from new node states without changing the MLP or graph, and replace the training key with a new shared noise key. This checks interpolation to nearby trajectories and another random realization; it is not a claim of global extrapolation.

validation_initial_state = jnp.array(
    [
        [-1.05, -0.45, 0.35, 0.95],
        [-0.45, 0.05, 0.40, 0.70],
    ]
)

validation_teacher_config = teacher_config.copy()
validation_teacher_config.initial_state.dynamics = validation_initial_state
validation_fitted_config = fitted_config.copy()
validation_fitted_config.initial_state.dynamics = validation_initial_state

validation_noise_key = jax.random.key(29)
validation_teacher_config.noise.key = validation_noise_key
validation_fitted_config.noise.key = validation_noise_key

validation_target = teacher_solve(validation_teacher_config)
validation_prediction = compiled_solve(validation_fitted_config)
validation_mse = jnp.mean(
    (validation_prediction.ys[:, 0, :] - validation_target.ys[:, 0, :]) ** 2
)

assert jnp.isfinite(validation_mse)
assert validation_mse < initial_loss * 0.1

Figure 5: Validation stochastic rollout. The fitted shared correction is reused without retraining from four unseen initial states and an unseen noise realization. Teacher and UDE receive the same random forcing. Dashed lines show the complete teacher and solid lines show the UDE prediction.

Explore Network Capacity and Cost

The width of 16 used above is a deliberately modest default, not a previously optimized hyperparameter. We can make that choice more systematic by holding the depth, optimizer, step count, training trajectory, and validation trajectory fixed while varying only the hidden width.

Here we use a simple tolerance rule: choose the smallest network whose validation MSE is within 20% of the best value in the sweep. Unlike selecting the absolute minimum, this favors a smaller model when additional capacity provides only a marginal improvement.

WIDTHS = (2, 4, 8, 16, 32)
SWEEP_DEPTH = 2
SWEEP_STEPS = 500
KNEE_TOLERANCE = 0.20

These are the user-facing sweep controls. The folded implementation below fits one otherwise identical model per width, records first-fit time, and applies the tolerance rule.

def parameter_count(module):
    arrays = jax.tree.leaves(eqx.filter(module, eqx.is_inexact_array))
    return sum(array.size for array in arrays)


def block_until_ready(tree):
    for leaf in jax.tree.leaves(tree):
        if hasattr(leaf, "block_until_ready"):
            leaf.block_until_ready()


capacity_results = []
for width in WIDTHS:
    module_key = jax.random.fold_in(jax.random.key(101), width)
    candidate_correction = EquinoxParameter(
        make_correction_module(
            module_key,
            in_size=2,
            out_size=1,
            width_size=width,
            depth=SWEEP_DEPTH,
        )
    )
    candidate_solve, candidate_config = prepare_network(
        UDEFitzHughNagumo(correction=candidate_correction)
    )

    def candidate_loss(config):
        predicted_V = candidate_solve(config).ys[:, 0, :]
        return jnp.mean((predicted_V - target_V) ** 2)

    candidate_optimizer = OptaxOptimizer(
        candidate_loss,
        optax.adam(learning_rate=3e-3),
    )

    start = time.perf_counter()
    candidate_fitted, _ = candidate_optimizer.run(
        candidate_config,
        max_steps=SWEEP_STEPS,
        chunk_size=50,
    )
    block_until_ready(candidate_fitted)
    fit_seconds = time.perf_counter() - start

    candidate_train_mse = float(candidate_loss(candidate_fitted))
    candidate_validation_config = candidate_fitted.copy()
    candidate_validation_config.initial_state.dynamics = validation_initial_state
    candidate_validation_config.noise.key = validation_noise_key
    candidate_validation = candidate_solve(candidate_validation_config)
    candidate_validation_mse = float(
        jnp.mean(
            (
                candidate_validation.ys[:, 0, :]
                - validation_target.ys[:, 0, :]
            )
            ** 2
        )
    )

    capacity_results.append(
        {
            "width": width,
            "parameters": parameter_count(candidate_correction.module),
            "train_mse": candidate_train_mse,
            "validation_mse": candidate_validation_mse,
            "fit_seconds": fit_seconds,
        }
    )

assert all(
    np.isfinite(result[key])
    for result in capacity_results
    for key in ("train_mse", "validation_mse", "fit_seconds")
)

best_validation_mse = min(
    result["validation_mse"] for result in capacity_results
)
knee_threshold = (1.0 + KNEE_TOLERANCE) * best_validation_mse
selected_capacity = next(
    result
    for result in capacity_results
    if result["validation_mse"] <= knee_threshold
)

Figure 6: Capacity and computational cost. Left: training and validation errors after 500 optimizer steps. The selected point is the smallest model within 20% of the best validation MSE. Right: first-fit wall time, including JAX compilation, on the machine that rendered this page. Width labels appear beside the validation points.

In this run the rule selects width 16, with 337 trainable parameters. Increasing the width to 32 raises the count to 1,185 and takes longer, but does not improve validation error. That is the kind of diminishing-return knee the sweep is designed to expose.

The timing includes tracing and compilation, so it describes the complete first fit a notebook user experiences rather than steady-state kernel throughput. It is hardware- and software-dependent and should not be treated as a portable benchmark. Compilation caches and system load can also make small timing differences non-monotonic.

This single sweep gives an objective rule for this example, but not uncertainty on the architecture choice. For a scientific analysis, repeat the comparison across initialization and noise seeds, then evaluate the selected architecture once on a final test set that was not used for fitting or model selection.

Train One Shared Correction Across Signal Segments

The fit above learns from one 15-unit trajectory. A longer recording exposes the same correction to more of the phase plane, but differentiating one continuous rollout also lengthens the gradient path. Here we instead construct consecutive fixed-length teacher segments and average their losses with jax.vmap. Every segment evaluates the same capacity-selected EquinoxParameter, so their gradients contribute to one shared update. We use shorter five-unit segments to make the state-coverage difference between one and four segments explicit.

N_TRAIN_SEGMENTS = 4
N_HELD_OUT_SEGMENTS = 2
SEGMENT_T1 = 5.0
SEGMENT_FIT_STEPS = 2_000
SELECTED_WIDTH = selected_capacity["width"]

def make_segment_dataset(start_state, n_segments, seed):
    segment_keys = jax.random.split(jax.random.key(seed), n_segments)
    segment_initial_states = []
    segment_targets = []
    current_state = start_state

    for segment_key in segment_keys:
        local_config = segment_teacher_config.copy()
        local_config.initial_state.dynamics = current_state
        local_config.noise.key = segment_key
        segment_solution = segment_teacher_solve(local_config)

        segment_initial_states.append(current_state)
        segment_targets.append(segment_solution.ys)
        current_state = segment_solution.ys[-1]

    return {
        "initial_state": jnp.stack(segment_initial_states),
        "target": jnp.stack(segment_targets),
        "noise_key": jnp.stack(segment_keys),
    }


segment_teacher_solve, segment_teacher_config = prepare_network(
    TeacherFitzHughNagumo(), t1=SEGMENT_T1
)
all_signal_segments = make_segment_dataset(
    initial_state,
    N_TRAIN_SEGMENTS + N_HELD_OUT_SEGMENTS,
    seed=41,
)
training_segments = jax.tree.map(
    lambda value: value[:N_TRAIN_SEGMENTS], all_signal_segments
)
held_out_segments = jax.tree.map(
    lambda value: value[N_TRAIN_SEGMENTS:], all_signal_segments
)

The complete synthetic state is known at every segment boundary. With partially observed recordings, those initial states would need to be estimated, encoded, or preceded by a burn-in; segmenting must not silently invent latent states.

def segment_batch_loss(config, dataset):
    def one_segment(initial_state, target, noise_key):
        local_config = eqx.tree_at(
            lambda c: (c.initial_state.dynamics, c.noise.key),
            config,
            (initial_state, noise_key),
        )
        prediction = segment_solve(local_config).ys[:, 0, :]
        return jnp.mean((prediction - target[:, 0, :]) ** 2)

    losses = jax.vmap(one_segment)(
        dataset["initial_state"],
        dataset["target"],
        dataset["noise_key"],
    )
    return jnp.mean(losses)


single_training_segment = jax.tree.map(
    lambda value: value[:1], training_segments
)

def single_segment_loss(config):
    return segment_batch_loss(config, single_training_segment)


def multi_segment_loss(config):
    return segment_batch_loss(config, training_segments)


def held_out_segment_loss(config):
    return segment_batch_loss(config, held_out_segments)

Both fits start from identical MLP weights and use the same optimizer settings. Only the number of training segments differs.

segment_module_key = jax.random.key(151)
single_correction = EquinoxParameter(
    make_correction_module(
        segment_module_key, in_size=2, out_size=1, width_size=SELECTED_WIDTH
    )
)
multi_correction = EquinoxParameter(
    make_correction_module(
        segment_module_key, in_size=2, out_size=1, width_size=SELECTED_WIDTH
    )
)

segment_solve, single_segment_config = prepare_network(
    UDEFitzHughNagumo(correction=single_correction), t1=SEGMENT_T1
)
_, multi_segment_config = prepare_network(
    UDEFitzHughNagumo(correction=multi_correction), t1=SEGMENT_T1
)

single_segment_fitted, _ = OptaxOptimizer(
    single_segment_loss, optax.adam(3e-3)
).run(single_segment_config, max_steps=SEGMENT_FIT_STEPS, chunk_size=50)
multi_segment_fitted, _ = OptaxOptimizer(
    multi_segment_loss, optax.adam(3e-3)
).run(multi_segment_config, max_steps=SEGMENT_FIT_STEPS, chunk_size=50)

Figure 7: One shared correction trained across signal segments. Left: learned corrections evaluated along the complete segmented signal and averaged within voltage bins across times and nodes; the black curve is the true residual. This uses the same voltage range as Figure 4, whereas the temporal holdout alone occupies only its negative-voltage branch. Right: marker shape distinguishes training from held-out trajectory error, while color consistently identifies the one- and four-segment fits; the lower panel shows their absolute train–validation gaps. Both models use the capacity-selected architecture, identical initial weights, and 2,000 updates. The single-segment model sees 5 time units; the multi-segment model averages the loss over four consecutive segments (20 time units) at every update. Both are evaluated on the final two segments of the same longer signal, which were not used for fitting.

The one-segment model nearly interpolates its short training window but fails on the later signal. Training across four segments reduces the held-out error and the train–validation gap by roughly three orders of magnitude. The improvement comes from exposing the shared correction to more of the trajectory, not from stochastic mini-batch regularization: every update uses all four segments.

Segment boundaries also stop state gradients, so this full-batch multiple-shooting objective is distinct from both one continuous long rollout and the solver’s grad_horizon option.

WarningInterpretation and identifiability

Trajectory agreement does not prove that the MLP discovered a unique physical law. A flexible correction can compensate for errors in fixed parameters, coupling, observations, or initial conditions, all of which are held at their known synthetic values here precisely so that the residual has an interpretable target.

Tutorial 2 treats this as the central question rather than a closing caveat: it works in a case where the missing term is known in closed form, and sets out what to measure, what to hold out, and which comparisons to run before crediting a network with a discovery.

NoteFrom this controlled example to noisy observations

Common random numbers isolate model error in this synthetic experiment because the latent forcing is known. With experimental observations, that forcing is usually unknown. A practical loss may then average over several simulated noise realizations or compare robust summaries rather than matching one trajectory point by point.

Summary

This tutorial demonstrated the complete UDE workflow in TVB-Optim:

1. A mechanistic network model supplied stable, interpretable dynamics and explicit inter-region coupling.
2. One EquinoxParameter marked a shared MLP as trainable inside config.dynamics.
3. eqx.filter_jit compiled a full configuration containing callable Equinox leaves.
4. Gradients propagated through the network solver to every MLP array leaf.
5. OptaxOptimizer fitted the correction without a special neural-network optimization path.
6. Multiple nodes jointly trained the same local correction under fixed random forcing, which was then evaluated on validation initial conditions and a new noise realization.
7. A capacity sweep compared validation error, parameter count, and indicative first-fit runtime instead of assuming that a larger MLP is better.
8. A vectorized multi-segment loss aggregated gradients from a longer signal into one shared correction and substantially improved temporal validation.

The essential modeling choice is separation of responsibilities: known coupling and slow recovery dynamics remain mechanistic, while the MLP learns only the local voltage residual placed explicitly in the equations.

Where those choices themselves come from, and how to tell a useful correction from a merely well-fitted one, is the subject of Universal Differential Equations 2.

References

• Rackauckas, C., Ma, Y., Martensen, J., et al. (2020). Universal Differential Equations for Scientific Machine Learning. arXiv:2001.04385.
• El-Gazzar, A., and van Gerven, M. (2025). Universal differential equations as a unifying modeling language for neuroscience. Frontiers in Computational Neuroscience, 19, 1677930.
• Chen, R. T. Q., Rubanova, Y., Bettencourt, J., and Duvenaud, D. K. (2018). Neural Ordinary Differential Equations. Advances in Neural Information Processing Systems, 31.
• FitzHugh, R. (1961). Impulses and physiological states in theoretical models of nerve membrane. Biophysical Journal, 1(6), 445–466.
• Nagumo, J., Arimoto, S., and Yoshizawa, S. (1962). An active pulse transmission line simulating nerve axon. Proceedings of the IRE, 50(10), 2061–2070.