Forecasting and predicting stochastic agent-based model data with biologically-informed neural networks

This paper demonstrates that biologically-informed neural networks (BINNs) can be trained to generate interpretable partial differential equation models that accurately forecast and predict the behavior of computationally intensive stochastic agent-based models of collective migration, even in parameter regions where traditional mean-field equations fail or require more complex compartmentalization.

The Gist

Imagine you are trying to predict how a massive crowd of people will move through a city square.

The Problem: The "Individual" vs. The "Crowd"
You have two ways to study this:

1. The Individual Approach (Agent-Based Model): You track every single person. You know exactly where Person A is, that they are tired, that they are pulling Person B along, and that Person C is sticking to Person D. This is incredibly accurate, but it's a nightmare to run on a computer. If you have 10,000 people, the computer has to do billions of calculations just to see where they are in the next second. It's like trying to predict the weather by tracking every single water molecule in the atmosphere.
2. The Crowd Approach (Differential Equations): Instead of tracking individuals, you treat the crowd like a fluid (like water or smoke). You say, "The crowd is moving this fast because it's dense here." This is super fast to calculate. But, it often fails. If the crowd suddenly decides to hold hands and stop moving, or if they start pulling each other in weird directions, the "fluid" math breaks down. It gives you the wrong answer, or sometimes, no answer at all.

The Solution: The "Biologically-Informed Neural Network" (BINN)
The author of this paper, John Nardini, came up with a clever middle ground. He created a "smart translator" called a BINN.

Think of the BINN as a super-smart apprentice chef.

• The Old Way (Mean-Field Models): A recipe book that says, "If you have 500 grams of flour, add 2 eggs." This works for simple cakes, but if you try to make a complex soufflé with weird ingredients, the recipe fails.
• The Individual Way (ABM): Watching a master chef cook the dish 1,000 times, recording every tiny movement of their hand, every pinch of salt, and every stir. It's perfect, but you can't watch it 1,000 times because it takes too long.
• The BINN Way: You show the apprentice chef (the Neural Network) the results of the master chef's cooking (the ABM data). You tell the apprentice: "You must learn the rules of cooking (the physics/math), but you can't just guess the recipe. You have to watch the master and learn the secret sauce."

The apprentice learns the hidden "secret sauce" (the complex rules of how cells pull and stick to each other) directly from the data, without needing to know the exact rules beforehand.

What Did They Do?
The paper tests this on three types of "cell crowds" (simulating how cells move in wound healing or tumor growth):

1. The Pullers: Cells that grab their neighbors and drag them along.
2. The Stickers: Cells that glue themselves to neighbors, stopping movement.
3. The Mixed Crowd: A chaotic mix of pullers and stickers.

The Results: Why It Matters

• Forecasting (Predicting the Future): The BINN learned the rules from a short video of the cells moving. Then, it successfully predicted what the cells would do later, even parts of the video it hadn't seen yet. It was as accurate as the slow, heavy computer simulation but much faster.
• Predicting New Scenarios (The "What If" Game): This is the coolest part. The researchers asked the BINN: "What if we change the stickiness of the cells?"
  • The old "fluid" math (Mean-Field) would crash and say, "I can't calculate that!" because the math breaks down when things get too sticky.
  • The BINN, however, used a technique called interpolation. Imagine you know how a car drives at 10 mph and 20 mph. You can guess how it drives at 15 mph. The BINN learned how cells move at different stickiness levels, and then it "guessed" the behavior for a brand new stickiness level it had never seen before. It worked perfectly.

The Trade-off
There is one catch. Training this "apprentice chef" takes a long time (about 11 hours on a supercomputer). However, once trained, the BINN can run simulations in seconds.

• Old Way: Run the simulation 10,000 times to test different scenarios = 6,600 hours of computer time.
• BINN Way: Train once (11 hours), then run 10,000 simulations = 349 hours.
  That's a 19x speedup.

The Bottom Line
This paper shows that we don't have to choose between "slow but accurate" (tracking every cell) and "fast but inaccurate" (treating cells like water). By using AI to learn the hidden rules from the slow simulations, we can create fast, accurate, and understandable models. This helps scientists figure out how to stop cancer cells from spreading or how to heal wounds faster, without waiting years for computer results.

Technical Summary

Here is a detailed technical summary of the paper "Forecasting and predicting stochastic agent-based model data with biologically-informed neural networks" by John T. Nardini.

1. Problem Statement

Collective migration is a fundamental biological process (e.g., wound healing, tumor growth) often modeled using Stochastic Agent-Based Models (ABMs). While ABMs capture the discrete and stochastic nature of individual cell interactions, they suffer from two major limitations:

1. Computational Intensity: Simulating large populations over time is slow, making it difficult to explore the full parameter space or perform data-driven tasks like parameter estimation.
2. Predictive Failure of Mean-Field Approximations: To speed up simulations, modelers often "coarse-grain" ABM rules into continuous Mean-Field Differential Equation (DE) models (ODEs or PDEs). However, these mean-field models often fail to predict ABM behavior accurately, particularly when:
   • The mean-field assumption (agents respond to average neighbor behavior) is violated.
   • The resulting PDE becomes ill-posed (e.g., negative diffusion coefficients) at certain parameter values.
   • The system requires complex multi-compartment models that are difficult to interpret or derive.

The core challenge is to develop a surrogate modeling approach that is as computationally efficient as DE models but as accurate as the underlying stochastic ABMs, even in parameter regimes where traditional mean-field approximations fail.

2. Methodology

The authors propose a framework combining Biologically-Informed Neural Networks (BINNs) with Partial Differential Equation (PDE) simulations. The workflow involves three main stages:

A. Data Generation

Three case study ABMs were simulated to mimic cell migration experiments (barrier and scratch assays):

1. Pulling ABM: Agents pull neighbors.
2. Adhesion ABM: Agents adhere to neighbors (potentially aborting migration).
3. Pulling & Adhesion ABM: A hybrid of both mechanisms.
   Data was generated by averaging 25 stochastic simulations to obtain spatiotemporal agent densities.

B. Model Training (BINN)

Instead of deriving PDEs analytically, the authors use BINNs to learn the governing equations from data.

• Architecture: The BINN consists of two Multi-Layer Perceptrons (MLPs):
  1. $T_{MLP}(x, t)$: Predicts the total agent density.
  2. $D_{MLP}(T)$: Predicts the density-dependent diffusion rate.
• Loss Function: The network is trained to minimize a composite loss function:
  $$L_{BINN} = L_{WLS} + \epsilon L_{PDE} + L_{constr}$$
  • $L_{WLS}$: Weighted Mean Squared Error between the network output and the ABM training data.
  • $L_{PDE}$: Ensures the network satisfies the diffusion PDE framework: $\frac{\partial T}{\partial t} = \frac{\partial}{\partial x} (D(T) \frac{\partial T}{\partial x})$.
  • $L_{constr}$: Penalizes physically impossible values (e.g., negative diffusion).
• Training Strategy: A two-step process is used: first fitting the data, then enforcing the PDE constraint.

C. Forecasting and Prediction

Once trained, the BINN is used in two ways:

1. Forecasting (Fixed Parameters): The learned diffusion term $D_{MLP}(T)$ is extracted and used to simulate the PDE forward in time. This forecasts future ABM states at the same parameter values used for training.
2. Prediction (New Parameters): To predict behavior at unseen parameter values ($p_{new}$), the authors employ multivariate interpolation. They train BINNs on a set of prior parameters, extract the $D_{MLP}(T; p)$ functions, and interpolate them to create a continuous function $D_{interp}(T; p)$. This interpolated diffusion term is then used to simulate the PDE for new parameter sets.

3. Key Contributions

• BINN-Guided PDEs: Demonstrated that BINNs can learn interpretable, single-compartment PDEs that accurately approximate complex, multi-agent stochastic systems.
• Solving Ill-Posed Problems: Showed that BINN-guided models can accurately predict ABM data in parameter regimes where traditional mean-field PDEs are ill-posed (e.g., high adhesion probabilities where mean-field diffusion becomes negative).
• Parameter Space Exploration: Developed a workflow using multivariate interpolation of learned diffusion terms to efficiently predict ABM outcomes at new, unexplored parameter values without re-simulating the expensive ABM.
• Interpretability: Unlike "black-box" neural networks, the BINN approach yields an explicit functional form for the diffusion rate $D(T)$, allowing for biological interpretation of how parameters influence migration.

4. Results

The study evaluated the approach on three ABMs:

• Pulling ABM: Both mean-field PDEs and BINN-guided PDEs performed well. The BINN-guided model learned a diffusion rate that varied with density, matching the ABM data with low Mean Squared Error (MSE).
• Adhesion ABM:
  • Failure of Mean-Field: For high adhesion probabilities ($p_{adh} > 0.75$), the mean-field PDE became ill-posed (negative diffusion) and failed to predict data.
  • BINN Success: The BINN-guided PDE remained well-posed and accurately forecasted data even at $p_{adh} = 0.9$, where the mean-field model failed completely.
• Pulling & Adhesion ABM:
  • The mean-field approach required a complex two-compartment PDE system to describe the interaction.
  • The BINN approach successfully learned a single-compartment PDE that accurately forecasted the total agent density, simplifying the model while maintaining high accuracy.
• Prediction of New Parameters: Using multivariate interpolation, the BINN-guided PDEs successfully predicted ABM data for new parameter combinations (e.g., varying adhesion and migration rates simultaneously) with testing MSEs generally below $2 \times 10^{-4}$.

5. Significance and Implications

• Efficiency: While training a BINN is computationally expensive (11 hours), the resulting PDE simulation is **28 times faster** than running a single ABM simulation. This makes BINN-guided PDEs ideal for tasks requiring thousands of simulations, such as Approximate Bayesian Computation (ABC) for parameter estimation.
• Robustness: The method overcomes the limitations of analytical coarse-graining, providing accurate surrogates even when the underlying physics leads to ill-posed equations.
• Data-Driven Biology: This framework enables modelers to explore complex parameter spaces and estimate parameters from experimental data more efficiently, bridging the gap between stochastic microscopic models and deterministic macroscopic descriptions.
• Future Work: The authors suggest optimizing BINN architectures (e.g., via genetic algorithms) to reduce training time and validating the method on real experimental biological data.

In summary, the paper presents a powerful hybrid methodology that leverages the flexibility of neural networks to learn governing equations from stochastic data, resulting in fast, interpretable, and robust predictive models for collective cell migration.