Universal Differential Equations 2: Correcting a Neural Mass Model

A Conditioned, Anchored Correction for Conductance-Based Synapses

Try this notebook interactively:

Download .ipynb Download .qmd Open in Colab

Introduction

Next-generation neural mass models are exact reductions of spiking populations, but only under stated assumptions (Montbrió et al., 2015; Coombes and Byrne, 2019). All-to-all connectivity, Lorentzian heterogeneity, instantaneous synapses and current-based input are each a modelling decision that a real population may violate.

When your population violates one, there are two options: re-derive the reduction, or keep the derived model and learn the correction. This tutorial does the second, using the universal differential equation (UDE) construction (Rackauckas et al., 2020; El-Gazzar and van Gerven, 2025).

Universal Differential Equations 1 introduced the mechanics: an EquinoxParameter holding a shared module, partition_state and combine_state, gradients through the solver, and OptaxOptimizer. All of that is used here without re-explanation. This tutorial is about the modelling decisions that make such a correction scientifically useful rather than merely well fitted:

1. keep the mechanistic vector field in full and learn only the residual;
2. condition it on the parameter that measures how badly the assumption is violated;
3. anchor it so it vanishes exactly where the base model is already exact;
4. hold out conditions, not time points.

We work in a case where the correct answer is known in closed form, so every claim the method makes can be checked.

The violated assumption

The Montbrió-Pazó-Roxin (MPR) mean field assumes current-based synaptic input: each spike injects the same current regardless of the postsynaptic voltage. Real synapses are conductance-based: they open channels, and the resulting current is proportional to the driving force \((v_{\text{syn}} - V)\), so it shrinks as the membrane approaches the reversal potential and reverses beyond it.

Both mean fields ship with TVB-Optim, as MontbrioPazoRoxin (MPR) and CoombesByrne2D (CB). Writing \(g = \kappa\pi r\) for the population conductance and setting \(\tau = 1\):

\[ \begin{aligned} \text{current-based (MPR):}\quad \dot r &= \frac{\Delta}{\pi} + 2Vr, & \dot V &= V^2 - (\pi r)^2 + \eta + Jr + I(t), \\[4pt] \text{conductance-based (CB):}\quad \dot r &= \frac{\Delta}{\pi} + 2Vr - gr, & \dot V &= V^2 - (\pi r)^2 + \eta + (v_{\text{syn}} - V)g + I(t). \end{aligned} \]

Expanding the conductance term in \(\dot V\) gives \((v_{\text{syn}} - V)g = \kappa\pi v_{\text{syn}} r - \kappa\pi Vr\), which splits into a current-based part with \(J = \kappa\pi v_{\text{syn}}\) and a remainder. The \(\dot r\) equation admits no such split: the current-based reduction carries no synaptic term in the rate equation at all, so the whole \(-gr\) contribution is missing rather than partly accounted for. The exact difference between the two vector fields is

\[ \delta_r(r, V, \kappa) = -\kappa\pi r^2, \qquad \delta_V(r, V, \kappa) = -\kappa\pi Vr . \]

This is the target. It depends on the state, it depends on \(\kappa\), and it vanishes identically at \(\kappa = 0\). Those three properties are what the construction below is designed to respect.

ImportantThe baseline is given its best possible \(J\)

\(J = \kappa\pi v_{\text{syn}}\) is neither a free parameter nor a constant held fixed across conditions. It is the exact current-based coefficient obtained by expanding the conductance term, recomputed at every \(\kappa\). The derived model is therefore handed the best current-based approximation that exists, rather than a deliberately mis-set one, and what the correction has to learn is a term that no choice of \(J\) could have supplied. This is the difference between a correction that captures missing structure and one that is quietly compensating for a badly chosen constant.

Both vector fields are written out below rather than imported from MontbrioPazoRoxin and CoombesByrne2D, so that the drive and the correction stay visible at the point in the equations where they enter. They are otherwise those of the shipped models, with \(\tau = 1\).

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

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

from tvboptim.experimental.network_dynamics import Bunch, prepare
from tvboptim.experimental.network_dynamics.dynamics import AbstractDynamics
from tvboptim.experimental.network_dynamics.external_input import SineInput
from tvboptim.experimental.network_dynamics.solvers import RungeKutta4
from tvboptim.optim import OptaxOptimizer
from tvboptim.types import (
    EquinoxParameter,
    combine_state,
    partition_state,
)

Population Parameters and Drive

One excitatory population, driven by a slow sinusoid so that the trajectory visits a wide range of firing rates instead of sitting at a fixed point, with a reversal potential above the operating voltage so that the driving force matters.

DELTA = 1.0        # Lorentzian half-width of the excitability distribution
ETA = -2.0         # mean background excitability
V_SYN = 2.0        # synaptic reversal potential
DRIVE_AMPLITUDE = 3.0
DRIVE_FREQUENCY = 0.15

DT = 0.05
T1 = 25.0
INITIAL_RV = (0.1, -2.0)

TRAIN_KAPPA = jnp.array([0.25, 0.5, 1.0, 1.5])
HELD_OUT_KAPPA = jnp.array([0.0, 0.75, 2.0])

\(\kappa\) measures how strongly the conductance mechanism acts, so it is the natural coordinate for “how badly is the assumption violated”. We fit four values and hold out three: the anchor \(\kappa = 0\), an interpolation point inside the training range, and an extrapolation point beyond it.

Figure 1: The assumption that breaks, and what it costs. Left: synaptic input as a function of membrane voltage at fixed firing rate. The current-based model injects the same current at every voltage (flat), while a conductance-based synapse scales with the driving force and reverses at \(v_{\text{syn}}\) (sloped). Middle and right: the exact difference between the two vector fields, \(\delta_r = -\kappa\pi r^2\) and \(\delta_V = -\kappa\pi Vr\), over the fitted conditions. Both grow with \(\kappa\) and vanish identically at \(\kappa = 0\). Note that the two missing terms have different shapes: one is a quadratic damping of the rate, the other changes sign with the membrane potential.

Simulating Every Condition at Once

We need trajectories at four values of \(\kappa\), and one correction that learns from all of them. The obvious way to get them is a Python loop with one solve per condition. There is a better one.

A prepare() call with n_nodes=4 integrates four uncoupled populations in a single rollout, and model parameters broadcast over that node axis. So assigning an array to config.dynamics.kappa gives each node its own condition:

config.dynamics.kappa = jnp.array([0.25, 0.5, 1.0, 1.5])
#                                   node 0  node 1  node 2  node 3

The node axis normally holds brain regions. This model has no graph and no coupling, so it is free to hold conditions instead. One rollout returns all four trajectories, one loss averages over them, and the shared correction receives gradients from every condition in a single update. Adding a condition means lengthening an array rather than writing a loop.

Three models, and what separates them

Three vector fields appear from here on, and the argument depends on keeping them apart.

In the text	Class	Vector field	Role
reference	ConductancePopulation	conductance-based (CB)	Stands in for the real population. The data to be reproduced, and the only thing any error is measured against.
derived model	CurrentBasedPopulation	current-based (MPR)	What the reduction hands you once its synaptic assumption is violated. The baseline to beat.
corrected model	CorrectedPopulation	current-based \(+\ \delta f_\theta\)	The UDE, defined further below.

The derived and corrected models differ in exactly one thing: the corrected one adds \(\delta f_\theta\) to the same drift. Same equations, same parameters, same solver, same drive, same \(J = \kappa\pi v_{\text{syn}}\). Any difference between their trajectories is the learned term and nothing else. The reference is drawn black and dashed in every figure that follows.

def current_based_drift(r, V, params, drive):
    """Vector field of the current-based mean field with J = kappa * pi * v_syn."""
    J = params.kappa * jnp.pi * params.v_syn
    dr = params.Delta / jnp.pi + 2.0 * V * r
    dV = V**2 - (jnp.pi * r) ** 2 + params.eta + J * r + drive
    return jnp.stack((dr, dV))


BASE_PARAMS = dict(Delta=DELTA, eta=ETA, v_syn=V_SYN, kappa=1.0)


class Population(AbstractDynamics):
    """What the three models share. They differ only in their vector field."""

    STATE_NAMES = ("r", "V")
    INITIAL_STATE = INITIAL_RV
    DEFAULT_PARAMS = Bunch(**BASE_PARAMS)
    EXTERNAL_INPUTS = {"drive": 1}


class ConductancePopulation(Population):
    """Reference model: conductance-based synapse with an explicit driving force."""

    def dynamics(self, t, state, params, coupling, external):
        del t, coupling
        r, V = state
        g = params.kappa * jnp.pi * r
        dr = params.Delta / jnp.pi + 2.0 * V * r - g * r
        dV = (
            V**2
            - (jnp.pi * r) ** 2
            + params.eta
            + (params.v_syn - V) * g
            + external.drive[0]
        )
        return jnp.stack((dr, dV))


class CurrentBasedPopulation(Population):
    """Uncorrected model: the mean field you would have derived."""

    def dynamics(self, t, state, params, coupling, external):
        del t, coupling
        r, V = state
        return current_based_drift(r, V, params, external.drive[0])

def build(dynamics, kappa, dt=DT):
    solve_fn, config = prepare(
        dynamics,
        RungeKutta4(),
        t0=0.0,
        t1=T1,
        dt=dt,
        n_nodes=kappa.shape[0],
        externals={
            "drive": SineInput(
                frequency=DRIVE_FREQUENCY,
                amplitude=DRIVE_AMPLITUDE,
            )
        },
    )
    config.dynamics.kappa = kappa
    return solve_fn, config


def simulate(dynamics, kappa, dt=DT):
    """Build a configuration at these conditions and run it."""
    solve_fn, config = build(dynamics, kappa, dt=dt)
    return solve_fn(config)


target_solution = simulate(ConductancePopulation(), TRAIN_KAPPA)
uncorrected_solution = simulate(CurrentBasedPopulation(), TRAIN_KAPPA)

assert target_solution.variable_names == ("r", "V")
assert target_solution.ys.shape == (int(T1 / DT), 2, TRAIN_KAPPA.shape[0])
assert jnp.all(jnp.isfinite(target_solution.ys))

The assert statements here and throughout are not part of the workflow. They are executable claims about the code: this page fails to build if any of them stops holding, so nothing the text asserts can quietly drift away from what the code does.

NoteChoosing the time step

This vector field is smooth but strongly nonlinear, so the reference data should be integrator-limited by a comfortable margin rather than marginally resolved. RungeKutta4 at the step used here agrees with a ten-times finer reference to well below the size of the correction we are trying to learn. Check that margin before trusting a learned residual: discretization error the network can see is discretization error the network will absorb.

fine_solution = simulate(ConductancePopulation(), TRAIN_KAPPA, dt=DT / 10)
discretization_error = float(
    jnp.max(jnp.abs(target_solution.ys - fine_solution.ys[9::10]))
)

print(f"Max deviation from a 10x finer solve: {discretization_error:.2e}")

Max deviation from a 10x finer solve: 1.11e-05

Figure 2: How the derived model fails. Firing rate under the same drive at each fitted condition: current-based (solid) against the conductance-based reference (dashed). Both respond more strongly as \(\kappa\) grows, since the effective synaptic drive \(J = \kappa\pi v_{\text{syn}}\) grows with it. The reference is held back by the shrinking driving force \((v_{\text{syn}} - V)\) and by the rate damping, and rises only modestly. The current-based model has neither brake and runs away, so the gap widens with every condition. By \(\kappa = 1.5\) the failure is qualitative rather than quantitative: a second peak appears within each drive cycle. Rightmost panel: the same discrepancy on a log axis, showing that it is already present at \(\kappa = 0.25\), where the two trajectories look identical on the shared scale, and that it spans orders of magnitude across the fitted conditions.

A Conditioned, Anchored Correction

Two design decisions do the scientific work.

Condition on the violation. The correction receives \(\kappa\) alongside the state. Without it, one network would have to represent four different residuals using only the trajectories to tell them apart, and nothing would transfer to an unseen condition. With it, the network learns a family of corrections indexed by how badly the assumption is violated.

Anchor by construction. At \(\kappa = 0\) the derived model is exact and the correction must be exactly zero. Rather than hoping training discovers this, or adding a penalty that only encourages it, we build it into the functional form:

\[ \delta f_\theta(y, \kappa) = g_\theta(y, \kappa) - g_\theta(y, 0). \]

Two evaluations of the same module, subtracted. This is identically zero at \(\kappa = 0\) for every value of the weights, before training, during training, and after. The known limit is a property of the model, not an outcome of the fit.

class CorrectedPopulation(Population):
    """Current-based mean field plus a conditioned, anchored learned residual."""

    DEFAULT_PARAMS = Bunch(**BASE_PARAMS, correction=None)

    def dynamics(self, t, state, params, coupling, external):
        del t, coupling
        r, V = state
        base_drift = current_based_drift(r, V, params, external.drive[0])

        kappa = jnp.broadcast_to(params.kappa, r.shape)
        here = jnp.stack((r, V, kappa), axis=-1)              # [nodes, 3]
        anchor = jnp.stack((r, V, jnp.zeros_like(kappa)), axis=-1)

        correction = (
            jax.vmap(params.correction)(here)
            - jax.vmap(params.correction)(anchor)
        )                                                     # [nodes, 2]
        return base_drift + correction.T

jax.vmap applies the same module to every node, so all conditions share one set of weights. The module has out_size=2 because both equations need correcting, and .T converts the [nodes, 2] result to the [2, nodes] state-major layout the solver expects.

The module built below also starts with a zeroed final layer, so the correction is zero everywhere at step zero and training begins from the derived model itself. Keep the two mechanisms apart. Zero initialization says where the optimizer starts, and the first gradient step destroys it. The anchor says what the functional form is, and no gradient step can touch it.

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(0), in_size=3, out_size=2)
)
ude_solve, ude_config = build(
    CorrectedPopulation(correction=correction), TRAIN_KAPPA
)

The Anchor Holds Exactly

Before fitting anything, check the construction. Simulated at \(\kappa = 0\), the corrected model must reproduce the conductance-based reference bit for bit, because at that condition the two models are the same equations and the correction is structurally zero.

We check it with untrained random weights, not with the zero-initialized module, so that the result cannot be mistaken for a consequence of the initialization.

random_module = eqx.nn.MLP(
    in_size=3, out_size=2, width_size=16, depth=2,
    activation=jax.nn.tanh, key=jax.random.key(99),
)

anchor_reference = simulate(ConductancePopulation(), jnp.array([0.0]))
anchor_ude = simulate(
    CorrectedPopulation(correction=EquinoxParameter(random_module)),
    jnp.array([0.0]),
)
anchor_deviation = float(jnp.max(jnp.abs(anchor_ude.ys - anchor_reference.ys)))

assert anchor_deviation == 0.0
print(f"Random-weight correction, max deviation at kappa = 0: {anchor_deviation:.1e}")

Random-weight correction, max deviation at kappa = 0: 0.0e+00

This is an architectural unit test, not a fitting result. A penalty term that merely encourages the correction to vanish at the known limit could not make this assertion, and would leave the limit accurate only to the tolerance the optimizer happened to reach.

Figure 3 shows the shape of that guarantee: the untrained correction is wrong everywhere it is free to be wrong, and exactly zero where it is not.

Figure 3: What anchoring guarantees before any training. The correction produced by the untrained module at one representative state, as a function of the conditioning variable, against the exact residual that fitting has to recover. Away from the anchor the untrained correction is small and points the wrong way in both components, so training has real work to do. At \(\kappa = 0\) it is exactly zero rather than approximately zero, because the construction subtracts the module from itself there, and it stays exactly zero at every step of training.

Gradient Flow Through the Solver

Before spending compute on a fit, confirm that gradients reach the module through the entire rollout. A silently zero or non-finite gradient is far cheaper to find here than after four thousand steps. What partition_state and combine_state separate, and why an Equinox module needs them, is covered in tutorial 1.

target_ys = target_solution.ys

def loss(config):
    return jnp.mean((ude_solve(config).ys - target_ys) ** 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.075e-01

Fit One Correction Across All Conditions

Every optimizer step uses all four conditions. This is full-batch training over the condition axis, not stochastic mini-batching. It runs on a CPU without special hardware.

NoteReading the training cell

chunk_size fuses that many optimizer steps into a single compiled lax.scan, which avoids a Python round-trip per step. It does not change the result, but it does change what the callback sees: the callback fires once per chunk, with the step index and loss of the last step in it, so the recorded loss history below is sampled every hundred steps rather than every step.

The callback signature is fixed by OptaxOptimizer. It receives the step index, the differentiable and static halves of the state, the fitting data, any auxiliary output, the loss and the gradients, and returns (stop, new_diff, new_static). A callback can therefore halt the run early or modify the state mid-optimization; this one only records the loss and returns False.

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

FIT_STEPS = 4_000

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

final_loss = float(compiled_loss(fitted_config))
loss_values[-1] = final_loss
fitted_solution = eqx.filter_jit(ude_solve)(fitted_config)

assert jnp.all(jnp.isfinite(fitted_solution.ys))
assert final_loss < initial_loss * 1e-2

Figure 4: One correction, four conditions. Top: firing rate after fitting, one panel per condition. The vertical axis stops just above \(r = 1\), where Figure 2 needed more than twice that range to contain the derived model: the overshoot is gone. Bottom left: trajectory loss over the shared update, aggregated across all fitted values of \(\kappa\). Bottom right: what is left over, corrected minus reference, on one scale for all conditions. The remaining error is far smaller than the uncorrected mismatch, but it is not noise. It concentrates at the peaks of each drive cycle, which is where the missing term \(-\kappa\pi r^2\) is largest. The fit is close everywhere and exact nowhere.

TipIf the fit does not converge

This problem is small enough that a direct fit over the whole rollout works. Longer horizons, stiffer models and chaotic regimes often need help, and three options compose with everything above without changing the model: grad_horizon on the native solvers truncates how far gradients propagate backwards; a multi-segment loss, shown at the end of tutorial 1, cuts the trajectory into shorter shooting windows; and prediction-error feedback adds a term \(K(\hat y(t) - y(t))\) to the observed equations during training, annealing \(K\) to zero so that the delivered model is autonomous.

Held-Out Conditions

The interesting question is not whether the fitted conditions are reproduced, but whether the family was learned. We reuse the trained module unchanged at three conditions it never saw: the anchor, a value inside the fitted range, and a value beyond it.

Note what is held out. The reflex is to split time points, train on the first half of each trajectory and test on the second. That measures almost nothing here, and the reason is worth spelling out: under a periodic drive the second half revisits states the first half already covered, so every test point has a near-duplicate in the training set. Held-out conditions do not, and further down that gap is measured rather than claimed.

Three simulations are needed at the held-out conditions: the reference to compare against, the derived model as the baseline, and the corrected model under test.

held_out_reference = simulate(ConductancePopulation(), HELD_OUT_KAPPA)
held_out_uncorrected = simulate(CurrentBasedPopulation(), HELD_OUT_KAPPA)
held_out_prediction = simulate(
    CorrectedPopulation(correction=fitted_config.dynamics.correction),
    HELD_OUT_KAPPA,
)

Only the third line involves anything learned, and even there nothing is fitted: fitted_config.dynamics.correction is the module trained at the other four conditions, dropped unchanged into a fresh configuration and simulated. The first two lines are the same two mechanistic models from the start of the tutorial, run at the new \(\kappa\) values.

Now the promised measurement. The two ways of holding data out differ in how far a test point sits from the training set, in the space the correction actually sees, \((r, V, \kappa)\), with each coordinate scaled by its range over the fitted data.

def visited_states(solution, kappa_values, time_slice):
    """States the solution passes through, as [n, 3] points in (r, V, kappa)."""
    r = np.asarray(solution.ys[time_slice, 0, :])
    V = np.asarray(solution.ys[time_slice, 1, :])
    kappa = np.broadcast_to(np.asarray(kappa_values), r.shape)
    return np.stack((r.ravel(), V.ravel(), kappa.ravel()), axis=-1)


half = target_solution.ys.shape[0] // 2
first_half = visited_states(target_solution, TRAIN_KAPPA, slice(0, half))
coordinate_scale = first_half.max(axis=0) - first_half.min(axis=0)


def distance_to_nearest(query, reference):
    """Median distance from each query point to the closest reference point."""
    offsets = (query[:, None, :] - reference[None, :, :]) / coordinate_scale
    return float(np.median(np.linalg.norm(offsets, axis=-1).min(axis=1)))


split_by_time = visited_states(target_solution, TRAIN_KAPPA, slice(half, None))
split_by_condition = visited_states(held_out_reference, HELD_OUT_KAPPA, slice(None))

print("Median distance from a test point to the nearest training point:")
print(f"  held out in time:      {distance_to_nearest(split_by_time, first_half):.4f}")
print(f"  held out by condition: {distance_to_nearest(split_by_condition, first_half):.4f}")

Median distance from a test point to the nearest training point:
  held out in time:      0.0039
  held out by condition: 0.2011

A time split leaves each test point sitting essentially on top of a training point, so passing it says nothing the training loss did not already say. A condition split moves the test points somewhere the correction has never been evaluated. Only the second measures generalization.

def normalized_rmse(reference, prediction):
    spread = reference.max() - reference.min()
    return float(
        jnp.sqrt(jnp.mean((reference - prediction) ** 2))
        / jnp.where(spread > 0, spread, 1.0)
    )


held_out_rows = [
    {
        "kappa": float(kappa),
        "uncorrected": normalized_rmse(
            held_out_reference.ys[:, 0, index], held_out_uncorrected.ys[:, 0, index]
        ),
        "corrected": normalized_rmse(
            held_out_reference.ys[:, 0, index], held_out_prediction.ys[:, 0, index]
        ),
    }
    for index, kappa in enumerate(np.asarray(HELD_OUT_KAPPA))
]

interpolation = next(row for row in held_out_rows if row["kappa"] == 0.75)

assert jnp.all(jnp.isfinite(held_out_prediction.ys))
assert held_out_rows[0]["corrected"] == 0.0          # the anchor is exact
assert interpolation["corrected"] < 0.3 * interpolation["uncorrected"]

Figure 5: Generalization across conditions. Left: normalized firing-rate error at each held-out \(\kappa\), before and after correction. The anchor is exact by construction and interpolation inside the fitted range is strongly improved. Extrapolation lowers the error number, but that number is misleading, which is what the right panel is for. Right: the extrapolation condition in full. Beyond the fitted range the correction does not degrade gracefully. It fails to hold down the spiking instability of the derived model and produces excursions several times larger than anything the reference reaches, so the trajectory has the wrong shape rather than being a less accurate version of the right one. An averaged error that halves can hide exactly this.

Interpolation inside the fitted range is what the construction was designed to deliver, and it delivers it. Extrapolation is not a weaker version of the same claim, it is a different claim, and here it fails: nothing constrains the network at a \(\kappa\) it never saw, and what it supplies is too little damping to hold down the derived model’s spurious spiking.

The transferable rule is to report the two separately, and to judge neither by an error norm alone. Whenever the base model has an instability that the correction is holding in check, losing that check outside the fitted range is the failure mode to expect, and an averaged error is the statistic least likely to reveal it.

NoteWhy extrapolation fails here, and what would fix it

Three things compound, and none of them is an accident of this fit.

The term is unbounded and the network is not. \(-\kappa\pi r^2\) grows without limit in both \(\kappa\) and \(r\). An eqx.nn.MLP with tanh hidden units cannot: its activations saturate, so its output is bounded for every choice of weights, before and after training. Outside the fitted range the correction can only flatten. Figure 7 shows it happening, the extrapolation points bending away from the exact curve as \(r\) grows.

The missing term is the stabilizing one. Under-supplying \(-\kappa\pi r^2\) does not nudge the trajectory, it removes the damping that keeps the population bounded and lets the derived model’s runaway excitation win. A correction whose job is to stabilize fails loudly, not gracefully.

The failure feeds itself. Once the trajectory leaves the states visited during fitting, the correction is queried further out of distribution, which makes it worse, which pushes the trajectory further out.

What would fix it, partly. Put more of what you know into the functional form, which is the move that made the anchor exact. The next section does that and measures how much it buys.

A Stronger Form for the Same Knowledge

The anchor worked because a fact we knew, that the correction vanishes at \(\kappa = 0\), was built into the functional form instead of being left to the optimizer. We know something stronger about this residual. Both components are \(\kappa\) times a function of the state alone:

\[ \delta_r = \kappa \cdot \left(-\pi r^2\right), \qquad \delta_V = \kappa \cdot \left(-\pi Vr\right). \]

So write the correction that way, with the module seeing only the state:

\[ \delta f_\theta(y, \kappa) = \kappa\, h_\theta(y). \]

The network never sees \(\kappa\). The \(\kappa\) dependence is supplied by multiplication, exactly, at every value including ones that were never fitted. It also vanishes at \(\kappa = 0\) without any subtraction, so this form subsumes the anchor instead of competing with it, and it gives \(h_\theta\) four times as much evidence about one function instead of asking one network to represent a family.

class ProportionalCorrectionPopulation(Population):
    """Base model plus a correction proportional to kappa by construction."""

    DEFAULT_PARAMS = Bunch(**BASE_PARAMS, correction=None)

    def dynamics(self, t, state, params, coupling, external):
        del t, coupling
        r, V = state
        base_drift = current_based_drift(r, V, params, external.drive[0])

        kappa = jnp.broadcast_to(params.kappa, r.shape)
        here = jnp.stack((r, V), axis=-1)                     # [nodes, 2]
        correction = kappa[:, None] * jax.vmap(params.correction)(here)
        return base_drift + correction.T

Nothing else changes: the same base model, the same reference data, the same optimizer, the same number of steps and the same seed. Only the place where the \(\kappa\) dependence lives is different.

proportional_solve, proportional_config = build(
    ProportionalCorrectionPopulation(
        correction=EquinoxParameter(
            make_correction_module(jax.random.key(0), in_size=2, out_size=2)
        )
    ),
    TRAIN_KAPPA,
)


def proportional_loss(config):
    return jnp.mean((proportional_solve(config).ys - target_ys) ** 2)


proportional_fitted, _ = OptaxOptimizer(
    proportional_loss, optax.adam(learning_rate=3e-3)
).run(proportional_config, max_steps=FIT_STEPS, chunk_size=100)

proportional_held_out = simulate(
    ProportionalCorrectionPopulation(
        correction=proportional_fitted.dynamics.correction
    ),
    HELD_OUT_KAPPA,
)

for index, row in enumerate(held_out_rows):
    row["proportional"] = normalized_rmse(
        held_out_reference.ys[:, 0, index], proportional_held_out.ys[:, 0, index]
    )

Normalized firing-rate error at each held-out condition, for the derived model and for both correction forms:

header = f"{'kappa':>6}  {'derived':>9}  {'subtraction':>12}  {'proportional':>13}"
print(header)
print("-" * len(header))
for row in held_out_rows:
    print(
        f"{row['kappa']:>6}  {row['uncorrected']:>9.4f}"
        f"  {row['corrected']:>12.4f}  {row['proportional']:>13.4f}"
    )

 kappa    derived   subtraction   proportional
----------------------------------------------
   0.0     0.0000        0.0000         0.0000
  0.75     0.1790        0.0276         0.0061
   2.0     0.9082        0.6151         0.3747

extrapolation = next(row for row in held_out_rows if row["kappa"] == 2.0)
reference_peak = float(jnp.max(held_out_reference.ys[:, 0, 2]))
proportional_peak = float(jnp.max(proportional_held_out.ys[:, 0, 2]))

assert held_out_rows[0]["proportional"] == 0.0                      # anchor still exact
assert interpolation["proportional"] < 0.5 * interpolation["corrected"]
assert extrapolation["proportional"] < 0.8 * extrapolation["corrected"]
assert proportional_peak > 2.0 * reference_peak                     # and still wrong

Figure 6: What the stronger form buys, and what it does not. The same held-out conditions, with both correction forms drawn against the derived model and the reference. Left, interpolation: the proportional form lies on the reference, while the subtraction form undershoots each peak. Right, extrapolation: all three models still spike where the reference does not. The proportional form roughly halves the spurious excursion but does not remove it, so the trajectory remains the wrong shape. Fixing the \(\kappa\) dependence by construction bought the \(\kappa\) axis, and only that axis.

The anchor survives untouched: \(\kappa\,h_\theta(y)\) is zero at \(\kappa = 0\) for every weight vector, exactly as the subtraction was.

Interpolation improves by a large factor and is now visually exact. That gain is not about extrapolation at all, it is about a better-posed fit: \(h_\theta\) has one function to represent rather than a family, and all four conditions supply evidence about that same function.

Extrapolation improves, and still fails.

CautionWhat the structure bought, and what it did not

Making the \(\kappa\) dependence exact removed \(\kappa\) as an axis of extrapolation. \(\kappa = 2\) is now interpolation in the conditioning variable, and the error at that condition drops accordingly.

The state is still extrapolated. As soon as the trajectory overshoots, \(h_\theta\) is queried at firing rates well beyond anything the fit ever visited, and there the boundedness argument from the previous section applies unchanged: a tanh network saturates, so the damping it can supply flattens out exactly where an unbounded \(-\pi r^2\) is required. The spurious spike is roughly halved rather than removed.

That is the honest shape of the result, and the transferable lesson. Encoding what you know buys the axis you encoded and nothing else. Recovering the rest would mean putting structure into the state dependence as well, a polynomial basis in \(r\) in place of a free network, at which point you are less learning missing physics than fitting a form you had already guessed. Where that line sits is a modelling judgement, not something the optimizer can decide for you.

Did It Learn the Right Term?

Trajectory agreement is necessary, not sufficient. Because we know the exact answer here, we can ask the stronger question directly: evaluate the trained correction at every state the reference visited, and compare it with \(\delta_r = -\kappa\pi r^2\) and \(\delta_V = -\kappa\pi Vr\).

def evaluate_correction(module, r, V, kappa):
    """Same anchored subtraction the dynamics uses, evaluated outside the solver."""
    kappa_column = jnp.full_like(r, kappa)
    here = jnp.stack((r, V, kappa_column), axis=-1)
    anchor = jnp.stack((r, V, jnp.zeros_like(kappa_column)), axis=-1)
    return jax.vmap(module)(here) - jax.vmap(module)(anchor)

def recovery_metrics(learned, true):
    """Shape, scale and size of the mismatch between a learned and a known term.

    Correlation reports shape only and cannot see a uniform rescaling, so it is
    paired here with the least-squares slope through the origin, where 1.0 means
    the right magnitude, and with the RMS error relative to the RMS of the true
    term.
    """
    return {
        "correlation": float(np.corrcoef(learned, true)[0, 1]),
        "slope": float(np.sum(learned * true) / np.sum(true * true)),
        "relative_rmse": float(
            np.sqrt(np.mean((learned - true) ** 2)) / np.sqrt(np.mean(true**2))
        ),
    }


def spurious_state_dependence(r_values, learned_dr, true_dr, n_bins=25):
    """Width of the loop in the left panel of the figure below.

    The exact correction to the rate equation is a function of r alone, so any
    variation of the learned one at equal r is dependence on V that the network
    absorbed. Reported relative to the range of the exact term.
    """
    edges = np.linspace(r_values.min(), r_values.max(), n_bins + 1)
    bin_index = np.clip(np.digitize(r_values, edges) - 1, 0, n_bins - 1)
    spreads = [
        np.ptp(learned_dr[bin_index == b])
        for b in range(n_bins)
        if np.sum(bin_index == b) > 5
    ]
    return float(np.max(spreads) / np.ptp(true_dr))


fitted_module = fitted_config.dynamics.correction
recovery_rows = []
for index, kappa in enumerate(np.asarray(HELD_OUT_KAPPA)):
    if kappa == 0.0:
        continue                      # the true residual is identically zero here
    r = held_out_reference.ys[:, 0, index]
    V = held_out_reference.ys[:, 1, index]
    learned = np.asarray(evaluate_correction(fitted_module, r, V, kappa))
    true = np.asarray(
        jnp.stack((-kappa * jnp.pi * r**2, -kappa * jnp.pi * V * r), axis=-1)
    )
    row = {
        "kappa": float(kappa),
        "r": np.asarray(r),
        "V": np.asarray(V),
        "learned": learned,
        "true": true,
        "loop_width": spurious_state_dependence(
            np.asarray(r), learned[:, 0], true[:, 0]
        ),
    }
    row["metrics"] = {
        name: recovery_metrics(learned[:, component], true[:, component])
        for component, name in enumerate(("r", "V"))
    }
    recovery_rows.append(row)

for row in recovery_rows:
    print(f"kappa = {row['kappa']}")
    for name, metrics in row["metrics"].items():
        print(
            f"  delta_{name}: correlation {metrics['correlation']:.3f}"
            f"   slope {metrics['slope']:.2f}"
            f"   relative RMSE {metrics['relative_rmse']:.2f}"
        )
    print(f"  spurious V-dependence in delta_r: {row['loop_width']:.2f}")

for name in ("r", "V"):
    assert all(row["metrics"][name]["correlation"] > 0.9 for row in recovery_rows)
    assert all(0.5 < row["metrics"][name]["slope"] < 1.3 for row in recovery_rows)

kappa = 0.75
  delta_r: correlation 0.991   slope 1.04   relative RMSE 0.12
  delta_V: correlation 0.978   slope 1.24   relative RMSE 0.40
  spurious V-dependence in delta_r: 0.28
kappa = 2.0
  delta_r: correlation 0.994   slope 0.85   relative RMSE 0.16
  delta_V: correlation 0.997   slope 0.78   relative RMSE 0.23
  spurious V-dependence in delta_r: 0.14

Correlation is close to one in every case. It is also the one number here that would look the same if the correction were systematically half the right size, so read it next to the slope, not instead of it.

Figure 7: The learned residual against the known one. Each point evaluates the trained correction at a state the reference actually visited, at a condition that was never fitted. Black curves are the exact missing terms \(-\kappa\pi r^2\) and \(-\kappa\pi Vr\). Two features are worth reading off. The extrapolation condition follows the shape of the exact term closely while sitting systematically inside it, so its correlation stays near one while its slope does not. And the points form a loop rather than a curve: the exact \(\delta_r\) depends on \(r\) alone, so the width of that loop is dependence on \(V\) that the network absorbed and the true term does not have.

WarningWhat this does and does not establish

The correction recovers the shape of the known missing term along the states the reference visited, including at conditions that were never fitted. That is a real check, and it is more than a trajectory fit. It is not proof that the network encodes \(-\kappa\pi r^2\) as a function: the comparison lives entirely on the observed state manifold, and a flexible correction can absorb errors in fixed parameters, in the drive, or in the initial condition as readily as missing physics.

Here those are held at their known values precisely so that the residual has an interpretable target. On real data, compare alternative parameterizations, include a low-order polynomial residual as a baseline, and check whether a refitted mechanistic model explains the same discrepancy.

Adapting This to Your Own Assumption

The equations matter less than the pattern.

What changes: the reference data, the base model, the conditioning variable, and where the correction enters. If your data comes from a spiking simulator or a recording rather than a second mean field, freeze the measurement protocol first. Rate binning, voltage estimation and transient removal all leave residues that a flexible network cannot distinguish from missing physics.

What does not: put every dependence you can derive into the functional form rather than into the weights, and keep the free network for what is genuinely unknown; condition on the parameter that measures the violation; anchor by subtraction at the limit where the base model is exact; split by condition, never by time point; label interpolation and extrapolation separately, and look at a trajectory before believing an extrapolation error number; check the magnitude of a recovered term and not only its shape; and compare against a refitted mechanistic model and a low-order polynomial residual before crediting the network.

One caveat specific to this example. The violated assumption here does not change the state dimension: the conductance-based model is still exactly two-dimensional in \((r, V)\), so a memoryless correction can be exact, and nearly is. Sparse connectivity, non-Lorentzian heterogeneity and finite population size are not like this. Each introduces fluctuations or higher moments that no function of \((r, V)\) can represent, and there a good trajectory fit can hide a structurally insufficient model. Before adding capacity, test whether nearby states reached through different histories have different conditional drift. If they do, the answer is another state variable, memory or a stochastic term.

Summary

1. The derived mean field was kept in full, with a learned residual added only where its synaptic assumption breaks.
2. The residual was conditioned on \(\kappa\), so one module represents a family of corrections rather than a single one.
3. It was anchored by subtraction, making it exactly zero at \(\kappa = 0\) for every value of the weights, and this was checked as an architectural unit test with untrained weights before any fitting.
4. Conditions were carried on the node axis, so one shared correction received gradients from all of them per update.
5. Generalization was measured across held-out conditions. Interpolation inside the fitted range of \(\kappa\) succeeded; extrapolation beyond it produced trajectories of the wrong qualitative shape.
6. The learned residual was compared against the exact missing term along the observed states, by slope and relative error rather than correlation alone.
7. Rewriting the correction as \(\kappa\,h_\theta(y)\) moved the \(\kappa\) dependence out of the weights and into the form. The anchor stayed exact and interpolation improved by a large factor, but extrapolation only improved: the state dependence is still learned, and a bounded network still saturates where an unbounded term is needed.

References

• Montbrió, E., Pazó, D., and Roxin, A. (2015). Macroscopic description for networks of spiking neurons. Physical Review X, 5(2), 021028.
• Coombes, S., and Byrne, Á. (2019). Next generation neural mass models. In Nonlinear Dynamics in Computational Neuroscience, 1-16.
• 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.