Extracted main text — title through conclusion, appendix excluded. This is what our citation measures are computed over, published so the extraction can be checked by eye.
44,111 characters · 8 sections · 47 citation commands
Learning Individual Behavior in Agent-Based Models with Graph Diffusion Networks
Agent-Based Models (ABMs) are computational frameworks in which autonomous “agents” interact with each other and their environment, leading to emergent collective behavior wilensky2015introduction. ABMs are typically characterized by: (i) a well-defined network of interactions, where the state of each agent is influenced by the states of a specific set of other agents, usually from the previous time step; (ii) stochasticity, meaning that agents' decisions incorporate a degree of randomness, producing probability distributions over multiple runs that capture real-world uncertainty and variation. ABMs have proven to be a powerful tool for developing and refining theoretical understanding, particularly in identifying minimal sets of micro-level rules that generate realistic macro-level outcomes railsback2019agent. In this sense, they have been instrumental in modeling a diverse range of phenomena castellano2009statistical, including structure formation in biological systems, pedestrian traffic, urban aggregation, and opinion dynamics. More recently, ABMs have demonstrated their value as forecasting tools poledna2020economic, such as in predicting the economic impacts of the COVID-19 pandemic pichler2022forecasting.
However, this progress is occurring despite the absence of principled methods to systematically align ABMs with real-world data. While various approaches have been proposed for calibrating macro-level parameters of ABMs platt2020comparison, there are still no established methods for tuning the micro-level behaviors of individual agents to match observed data. One potential approach is to manually construct a probabilistic model that replicates the ABM and then use its likelihood function to estimate individual state variables monti2023learning. However, this method requires the manual development of an ad-hoc probabilistic framework that reproduces the original ABM. Thus, what is currently missing is a fully automated method for deriving a learnable, differentiable model directly from an ABM.
In this work, we propose a novel approach to address this challenge: combining a graph neural network with a diffusion model to learn a differentiable surrogate of an ABM, from ABM-generated data. We refer to this method as a Graph Diffusion Network (GDN). In this framework, a graph neural network captures the interactions that govern the evolution of each agent's state in response to other agents, while the diffusion model learns the distribution of possible state transitions, conditioned on these interactions. A central aspect of our approach is its explicit modeling of individual agent behavior. Rather than treating the system as a whole, we focus on how each agent acts as an independent entity, while also incorporating the influence of other agents on its decisions. This approach ensures that the emergent dynamics remain faithful to the decentralized nature of ABMs. By constraining the surrogate to model only micro-level rules, it cannot rely on shortcuts that predict only macro-level outcomes, preserving the distributed logic of the original ABM.
Our approach draws inspiration from previous work on using neural network models to emulate deterministic cellular automata grattarola2021learning. However, we extend this idea to the broader domain of ABMs by introducing a crucial component: stochasticity. By incorporating stochasticity, our architecture can learn directly from ABM-generated data traces, making it adaptable to a wide variety of agent-based models across diverse real-world applications. Furthermore, since our method is trained on data traces, it can seamlessly integrate empirical observations alongside simulated data, thus being potentially applicable to real-world scenarios. In this sense, our work represents a first step toward developing a comprehensive methodology for creating easy-to-use, learnable ABMs.
From a general perspective, an ABM can be represented as a stochastic process $ \mathbf{Z}_t \sim \mathbb{P}_\Theta ( \mathbf{Z}_t \mid \mathbf{Z}_{\tau < t} ) , $ where $\mathbf{Z}_t$ denotes the state variables at time $t$, $\Theta$ is a set of parameters, and $\mathbb{P}$ is a probability distribution implicitly defined by the model structure and parameters. The index $t$ represents discrete time. Typically, $\Theta$ consists of a small number of parameters, remains fixed in dimension, is interpretable by domain experts, and serves as the model’s primary control mechanism. Conversely, each element in $\mathbf{Z}_t$ captures an agent’s state, leading to a high-dimensional state space.
To illustrate this structure, we consider two ABMs used throughout the paper. The first is the well-known model by Schelling schelling1971dynamic, where $\mathbf{Z}_t$ captures agents' positions and colors, and $\Theta$ indicates preference for same-color neighbors. Even with some tolerance for neighbors of different colors, agents often form segregated clusters ubarevivciene2024fifty. This clear mismatch between individual preferences and aggregate outcomes is a classic example of emergence. The second model is a predator-prey model tang2022data,wilensky2015introduction, describing the ecological dynamics between two interacting species, with one acting as predator and the other as prey, similarly to the Lotka-Volterra equations. In this ABM, $\mathbf{Z}_t$ includes agent position and type (prey-predator), while $\Theta$ governs the probability to move, reproduce, or die. This model replicates the cyclical predator-prey population dynamics, typical of Lotka-Volterra systems, while also capturing complex spatial patterns reminiscent of spatial evolutionary games nowak1992evolutionary. Both the Schelling and predator-prey models are widely recognized as canonical ABMs and are frequently used as testbeds for the development of novel calibration and surrogate techniques Squazzoni2010, Jamali2024, McLane2011, tang2022data, TenBroeke2021, Murphy2020.
ABMs have traditionally been powerful for theory generation, but in recent years, they have become increasingly data-driven pangallo2024datadriven. To align ABM output with empirical data, most efforts focus on calibrating parameters $\Theta$ so that model-generated summary statistics match observed ones platt2020comparison, quera2023blackbirds. Less attention, however, has been paid to estimating agent states $\mathbf{Z}_t$, which is key for matching time series further to summary statistics. Some researchers use data assimilation methods like particle filters lux2018estimation or ensemble Kalman filters oswald2025agent to infer $\mathbf{Z}_t$. A more principled alternative is to make ABMs differentiable, enabling the maximization of a likelihood function via gradient descent and automatic differentiation monti2020learning,monti2023learning. While differentiability is straightforward for simple stochastic behaviors, such as those governed by Bernoulli trials arya2022automatic, it becomes far more challenging for complex behaviors like those observed in the Schelling and predator-prey models.
To address this and other challenges in ABMs, researchers have increasingly turned to more tractable surrogates, also known as meta-models or emulators kennedy2001bayesian, farah2014bayesian, glielmo2023reinforcement. Surrogate models typically learn directly the mapping from parameters $\Theta$ to static summary statistics, disregarding individual behavior and model dynamics. For instance, a surrogate in Lamperti et al. lamperti2018agent maps $\Theta$ to the mean growth rate of the economy. More recent research has also explored the emulation of model dynamics. Grattarola et al. grattarola2021learning use Graph Neural Networks to approximate cellular automata, which can be seen as a special case of ABMs with deterministic interaction rules. Dyer et al. dyer2024interventionally propose Ordinary Differential Equation emulators to construct interventionally consistent surrogates, ensuring that micro-state interventions produce results aligned with macro-state interventions. Casert et al. casert2024learning employ Transformers to model the transition of physical systems from one configuration to another, in terms of their transition rates rather than reproducing individual agent behavior. Their method is tailored to physical systems in continuous time, making it inapplicable for interacting agents in the general case, since it requires explicitly enumerating allowed transitions between configurations.
In contrast to these approaches, our work is the first to jointly emulate individual and stochastic interacting agents. This is particularly important, since ABMs are inherently stochastic and rely on individual-level interactions to produce emergent aggregate outcomes. Moreover, since our surrogate is differentiable by design, it paves the way for methods that estimate both individual-level parameters and state variables.
To achieve this goal, our framework relies on a novel combination of graph neural networks and diffusion models. Diffusion models ddpm_original_paper were first introduced in the context of image generation, where they demonstrated impressive generation capabilities diffusion_vision_review, and were then applied to other domains tabddpm. A number of works addressed graph data graph_diffusion_review, for example in molecule modeling molecule_modeling and protein structure generation protein_diffusion. However, these works focus on the generation of graphs, while our architecture learns to generate random samples that are conditioned on information found on a graph. To the best of our knowledge, our work is the first application of this generative framework to individual behavior modeling in simulation systems, such as ABMs.
Denoting the set of agents by $A$, let each agent $i\in A$ at discrete time $t$ be described by a state vector $\mathbf{Z}_t^{(i)}$, which may include both continuous and categorical features. Given the ABM parameters $\Theta$, the update rule of $\mathbf{Z}_t^{(i)}$ follows a stochastic transition process $P_{\Theta}$ given by
where $N_t^{(i)}$ is the set of agents interacting with agent $i$ at time $t$, inducing a (time-varying) interaction graph $G_t = (A, E_t)$ that we assume to be known. This formulation focuses on individual agents, capturing not the dynamics of the entire system, but the evolution of each agent over time. In this way, it makes the two core ingredients of ABMs explicit: (i) relational structure via local neighborhoods $N_t^{(i)}$; (ii) stochasticity in the choice of next states.
To effectively model these components in the same individual-level view, we leverage respectively (i) message‐passing GNNs message_passing_gnn, which model the relationship between the evolution of an agent’s state and the state of its neighbors on the graph; (ii) conditional diffusion models conditionalimagesynthesisdiffusionreview, generative architectures well-suited to learning complex, multimodal distributions, allowing us to capture the intrinsic stochasticity of agent behavior.
Our proposed method, dubbed Graph Diffusion Network (GDN), combines these two components into a single architecture. Together, these components let us learn both how any agent state is affected by its neighbors on the graph, and the inherent randomness driving agent dynamics, yielding a surrogate that can both emulate the original ABM and be differentiated.
\noindentOverview. To learn the distribution $P_{\Theta}$, the training phase requires observations of different outcomes given the same starting conditions. To do so, in our framework, we use the original ABM to generate a data set as a ramification of possible states, namely $ \bigl(\mathbf{Z}_t^{(i)},\{\mathbf{Z}_t^{(j)}\}_{j\in N_t^{(i)}}\bigr)\;\longrightarrow\;\mathbf{Z}_{t+1}^{(i)}$ (see Figure (ref)). Our Graph Diffusion Network then approximates the stochastic kernel $P_{\Theta}$ by integrating a Message‐Passing GNN with a Conditional Diffusion Model, of learnable parameters $\omega$ and $\phi$ respectively. The GNN aggregates each agent’s state $\mathbf{Z}_t^{(i)}$ and its neighbors’ states $\{\mathbf{Z}_t^{(j)}\}_{j \in N_t^{(i)}}$ via permutation‐invariant message and readout functions to produce an interaction embedding $\mathbf{g}_t^{(i)}$. This embedding acts as a compact representation of the information coming from $i$'s neighbors at $t$, affecting the distribution of possible states of agent $i$ at time $t+1$. As such, it is passed to the diffusion model: conditioned on $\mathbf{Z}_t^{(i)}$ and $\mathbf{g}_t^{(i)}$, the diffusion model learns to transform a sample of Gaussian noise into a possible instance of the next state $\mathbf{Z}_{t+1}^{(i)}~$. By minimizing the standard denoising loss over all observed transitions, this hybrid architecture captures both the graph‐structured interactions and the inherent stochasticity of agent dynamics. The trained model $GDN_{\phi,\omega}$ is therefore able to generate, given a graph $G_t$ of interacting agents and the state of each one $\mathbf{Z}^{(i)}_{t}$, a sequence of possible next states $\mathbf{Z}^{(i)}_{t+1}$. The consecutive application of $GDN_{\phi,\omega}$ allows for reproducing the behavior of the original model. We now describe in detail each of these components.
\noindentMessage‐passing GNN. The GNN operates on the provided interaction graph $G_t=(A,E_t)$, that we assume to be known or to be computable from $\mathbf{Z}_t$ (e.g., in the Schelling model, the position of the agents determines who interacts with whom). For each agent $i$, the GNN aggregates its state $\mathbf{Z}_t^{(i)}$ together with each neighbor’s state $\mathbf{Z}_t^{(j)}$ via a permutation‐invariant operator $\bigoplus$, and then feeds the concatenated result through an MLP $f_\omega$ message_passing_gnn, that is $ \label{eq:message_passing_simple} \mathbf{g}_t^{(i)} = f_\omega\!\Bigl(\mathbf{Z}_t^{(i)},\,\bigoplus_{j\in N_t^{(i)}}\bigl(\mathbf{Z}_t^{(i)} , \mathbf{Z}_t^{(j)}\bigr)\Bigr) $. The resulting vector $\mathbf{g}_t^{(i)}$ captures how $i$’s local neighborhood influences its next‐state distribution. In practice, for the operator $\bigoplus$, we used sum aggregation for Predator-Prey and mean aggregation for Schelling, since the latter’s dynamics depend on the degree. Minimal experimental evaluation can guide practitioners toward the most suitable operator.
\noindentConditional diffusion model. The diffusion model then learns the distribution over future states given this output from the graph and the state of a given agent. Diffusion models do so by reversing a fixed Gaussian noising process ddpm_original_paper. The obtained denoising process, indexed by $\tau \in \{\tau_{\text{max}}, \dots, 0 \}$, starts from a sample of Gaussian noise $\mathbf{x}_{\tau_{max}}\sim\mathcal{N}(\mathbf{0,I})$ and in a sequence of denoising diffusion steps transforms it into a possible next state $\mathbf{x}_0 \approx \mathbf{Z}_{t+1}^{(i)}$. In this setting, we denote the general latent $\mathbf{x}_{\tau}$ as $\tilde{\mathbf{Z}}_{t+1}^{(i)}(\tau)$. Each step of this process receives as input (i) the agent’s current state $\mathbf{Z}_t^{(i)}$, (ii) its interaction embedding $\mathbf{g}_t^{(i)}$, and (iii) a sinusoidal positional embedding of $\tau$. These inputs are first transformed by MLPs to form the condition vector $\mathbf{c}_t^{(i)}$. Then, a feed‐forward network $\phi$ is trained to predict the noise residual $\epsilon_\phi$, i.e., the change to apply to the input $\tilde{\mathbf{Z}}_{t+1}^{(i)}(\tau)$ to continue the denoising process.
\noindentRamification data set. Given these two components, our framework uses the original ABM to produce a ramifications data set (see Figure (ref)). Such data set follows one main branch that specifies the evolution of the ABM, and multiple alternative stochastic evolutions of each time step from time $t$ to time $t+1$. This method makes it possible to expose the model to multiple stochastic successors from identical conditioning context, while avoiding exponential growth in the number of histories. Starting from an initial configuration $ \mathbf{Z}_0 = \bigl(\mathbf{Z}_0^{(1)},\dots,\mathbf{Z}_0^{(n)}\bigr)\,, $ we recursively simulate $R+1$ child configurations at each time step $t$, yielding $\{\mathbf{Z}_{t+1}[r]\}_{r=0,\dots,R}$. We designate the branch $r=0$ as the main branch $\{\mathbf{Z}_t[0]\}_{t=0,\dots,T-1}$, from which we extract the conditioning tuples $ \bigl(\mathbf{Z}_t^{(i)},\,\{\mathbf{Z}_t^{(j)}\}_{j\in N_t^{(i)}}\bigr) $ for all agents $i$. The remaining $R$ sibling branches at each $t$ supply the target next states $\mathbf{Z}_{t+1}^{(i)}$, ensuring that each context yields multiple outcomes.
\noindentLearning procedure. Our framework uses these data sets to train the Graph Diffusion Network. It minimizes the expected denoising loss over the outcomes observed in the ramification data (see Algorithm 1). At each training iteration, it uniformly samples a time index $t$ and extracts the conditioning pair $(\mathbf{Z}_t^{(i)},\{\mathbf{Z}_t^{(j)}\})$ from the main branch $\mathbf{Z}_t[0]$. We compute the interaction embedding $\mathbf{g}_t^{(i)}$ via Equation (ref), then draw a diffusion step $\tau$ to form the condition vector $\mathbf{c}_t^{(i)}$, and uniformly select one of the $R$ next‐state realizations to obtain the target $\mathbf{Z}_{t+1}^{(i)}$. Finally, we minimize the denoising loss in Equation (ref) by backpropagating through both the diffusion model and the GNN. More details about the architecture, and a discussion on computational costs and scalability can be found in Supplementary Section (ref).
In this section, we present different experiments to assess and demonstrate our framework's ability to learn micro‐level agent behaviors and faithfully reproduce emergent system‐level dynamics. We evaluate our Graph Diffusion Network on two canonical agent‐based models: the Schelling’s segregation model and the Predator–Prey ecosystem, presented in Section (ref) and detailed in Supplementary Section (ref). We test both its micro‐level and macro‐level fidelity. At the micro level, we measure how well the surrogate reproduces the conditional next‐state distribution of each agent under identical context on an out-of-training ramification data set. At the macro level, we assess whether the surrogate, once trained on the first $T_{\mathrm{train}}=10$ timesteps, can accurately reproduce the subsequent $T_{\mathrm{test}}=25$ timesteps of aggregate summary statistics.
Because no existing method directly accepts graph‐structured agent states and outputs per‐agent state distributions, there are no directly applicable baselines for our approach. Existing surrogates typically operate at the macro level, learning mappings from parameters to aggregate outcomes rather than reproducing full system dynamics. However, such approaches, including standard time-series models such as AR(1), fail to capture non-monotonic or cyclic patterns (e.g., predator–prey oscillations) and do not generalize beyond in-sample dynamics, as we show in Supplementary Section (ref). Therefore, we evaluate against two ablated variants. The first replaces the GNN embedding with a flat concatenation of all agent states, removing relational structure. The second keeps the GNN but removes the diffusion component, predicting deterministic next states instead. Both ablations are trained on the same ramified datasets and under identical protocols.
In the remainder of this section, we first describe the experimental setup, including dataset generation, model variants, and evaluation metrics. We then present a qualitative analysis of emergent patterns, followed by a comprehensive quantitative comparison. All implementation and reproducibility details are provided in the Supplementary Materials. Full code to reproduce our experiments is available at \href{https://github.com/fracozzi/ABM-Graph-Diffusion-Network}{\textcolor{blue}{http://github.com/fracozzi/ABM-Graph-Diffusion-Network}}.
\paragraph{Ablation.} Our core hypothesis is that both relational structure and stochastic modeling are crucial for accurate ABM surrogates. We consider therefore two possible ablations. In the first, we remove the message-passing GNN and replace the interaction graph with a flat concatenation of all agents’ state vectors---this isolates the impact of neglecting agent interactions. The second drops the diffusion component entirely, yielding a purely deterministic model similar in spirit to prior GNN-based approaches for deterministic automata grattarola2021learning. In this ablated version, we predict the next agent state by minimizing the mean squared error (MSE) with respect to the true next agent state. Details about the architecture can be found in Supplementary Section (ref). These ablations allow us to measure the improvement achieved by combining relational and stochastic modeling.
\paragraph{Agent-based models.} We evaluate our approach on the two ABMs described in Section (ref) as case studies. The first is the Schelling segregation model, in which $n$ agents occupy cells on a two-dimensional grid. Each agent has a fixed binary “color” and a position on the grid. At each timestep, an agent is considered happy if the proportion of its (up to eight) immediate neighbors sharing its color exceeds a tolerance threshold $\xi$; otherwise, it is unhappy and relocates to a randomly selected empty cell; thus, the interaction graph $G_t$ links each agent to its neighbors at time $t$. We adopt the standard NetLogo implementation of this model wilensky1999netlogo. The second is a predator–prey ecosystem model, where agents belong to one of two species (predator or prey), inhabit grid cells, and cycle through life phases---Unborn, Alive, Pregnant, and Dead. At each timestep, an Alive agent may move to a neighboring cell, reproduce (becoming Pregnant), or die, with probabilities specified by a parameter matrix $\Psi$ and conditioned on the local presence of predators or prey tang2022data, wilensky2015introduction. Pregnant agents revert to Alive after giving birth; Unborn agents become Alive if their parent is Pregnant; and Dead agents remain inactive. Here, $G_t$ links Alive neighboring agents, with specific rules for Pregnant and Unborn agents. See Supplementary Section (ref) for more details. In both ABMs, each agent’s full state at time $t$ comprises its position, type (color or species), and, for the predator-prey ABM, its life phase. Together, these two models span both simple relocation dynamics and richer birth–death interactions, providing diverse testbeds for our surrogate.
\paragraph{Micro evaluation metrics.} To quantify how faithfully our surrogate captures individual agent behavior, we compare its predicted conditional next‐state distributions against the ABM’s true stochastic transitions using the Earth Mover’s Distance (EMD) EMD. We extend the ramification dataset beyond the training horizon and generate corresponding datasets for both the surrogate and the ablation models. The EMD is then computed as the mean value across timesteps and individual agents. In the Schelling ABM, the EMD compares the distribution of agent positions. This directly measures the surrogate’s ability to relocate unhappy agents correctly, while keeping happy agents stationary. In the predator-prey model, we treat the agent’s categorical life phase as the random variable and compute the EMD over its four‐state distribution. This metric captures both deterministic transitions (e.g., Unborn $\to$ Alive, Dead $\to$ Dead) and stochastic, interaction‐driven transitions (e.g., Alive $\to$ Dead, Alive $\to$ Pregnant).
\paragraph{Macro evaluation metrics.} Next, we test whether agent‐level predictions translate into faithful reproduction of emergent, system‐level behavior. For each model, we track a summary statistic over time: the number of happy agents in Schelling, and the number of active (i.e. Alive and Pregnant) agents in the predator–prey ecosystem. Reusing the same ramification branches as in training would offer little new information, since different stochastic branches from the same state tend to produce very similar macroscopic trajectories. Instead, we generate a fresh ensemble of main‐branch simulations ($100$ independent runs) beyond the training horizon. We then compute the symmetric mean absolute percentage error (sMAPE) between the mean ground‐truth trajectory and the mean surrogate‐predicted trajectory across this ensemble, providing a quantitative measure of the surrogate’s ability to capture oscillations and steady-state behavior truly out-of-sample.
\paragraph{Experimental set-up.} We consider three parameter combinations $\xi$ for the Schelling ABM, each producing distinct segregation outcomes, and four $\Psi$ combinations for the predator-prey ABM, reflecting different oscillatory patterns in the population dynamics. For each ABM and parameter setting, we simulate $T_{\mathrm{train}}=10$ main‐branch steps with $R=500$ stochastic branches per step, yielding the training ramification as in Figure (ref). For macro‐evaluation, we run $100$ independent main‐branch simulations to calculate sMAPE. For micro-evaluation, we generate an out-of-sample ramification dataset of $T=25$ timesteps. We train both surrogate and ablations for $100$ epochs using Adam with learning rate $10^{-5}$ for the diffusion model and Adam with learning rate $2 \cdot 10^{-5}$ for the GNN, batch size equal to number of agents, and diffusion hyper-parameters $\tau_{\max}=100$ (more information in Supplementary Section (ref)).
To build intuition, we first qualitatively compare the surrogate and its ablated variant on their ability to reproduce key emergent patterns of agent-based dynamics. We then consolidate these insights with a comprehensive quantitative evaluation using the macro- and micro-level metrics introduced in the previous section. We report a selection of results in this Section; more in Supplementary Section (ref).
\paragraph{Reproducing emergent segregation.} Let us first consider the Schelling ABM, under the configurations $\xi_1=0.625$, $\xi_2=0.75$, $\xi_3=0.875$. Figure (ref) illustrates how the ground-truth ABM (top row) progresses from a randomized initialization to structured, segregated communities for the first two configurations $\xi_1, \xi_2$, while it remains unsegregated for $\xi_3$. At the first two tolerance levels, in fact, the agents gradually self‐organize into distinct clusters, with segregated communities clearly emerging around $t=20$ (see Supplementary Section (ref)). The middle row represents the evolution of the system according to our surrogate model: we initialize the system with the same starting condition, and then we iterate giving the current state $\textbf{Z}_{t}$ to our model, and using one sample of the generated output as the next state $\textbf{Z}_{t+1}$. We observe that our surrogate exhibits a qualitatively similar pattern of cluster formation over time, distinct for each of the three configurations. Instead, the ablation models fail to meaningfully relocate agents. The best-performing one, according to sMAPE, the diffusion-only model, largely maintains a random configuration.
\paragraph{Reproducing emergent oscillations in predator-prey dynamics.} Next, we consider the Predator-Prey ecological model. Figure (ref) overlays 100 trajectories of prey and predator populations starting from the same state at the end of training, comparing the stochastic trajectories from the ground-truth model with those obtained by our surrogate and by the ablation. For both configurations, the surrogate and the ablation models are trained only with the initial time steps (up to the dashed line in the plots). Under the parameter set $\Psi_1$, the ground-truth ABM (top-left plot) exhibits classical Lotka–Volterra oscillations: a rise in prey growth drives a delayed increase in predators, which then triggers prey decline and a subsequent predator decline. Under $\Psi_4$, instead, only predators show a rise and decay, while preys only decline (bottom-left plot). We observe that the surrogate (second column) accurately captures both the phase lag and amplitude of these oscillations, while the diffusion-only ablation (third column) collapses to near-monotonic trends. The GNN-only ablation (fourth column), besides following a completely deterministic dynamic, completely diverges from the ground truth. The number of preys under $\Psi_1$, for instance, quickly reaches a plateau at a value of 750 (not shown), almost twice the real one. We perform the same analysis for alternative parameterizations $\Psi_2$, $\Psi_3$ (included in Supplementary Section (ref)) that show different types of dynamics, as the population of predators and/or preys may exhibit monotonic extinctions. In all cases, the surrogate faithfully reproduces monotonic declines or single‐peak dynamics and both ablations fail. We also observe (figures in Supplementary Section (ref)) that the surrogate recreates the rich spatial patterns of predator–prey clusters, also seen in similar settings in evolutionary game theory nowak1992evolutionary.
\paragraph{Quantitative results.}
Now we present the results of a quantitative analysis, systematizing the previous comparisons. Here, each comparison with the ground truth is quantified using one of the metrics presented in the previous subsection, i.e. Earth Mover's Distance (EMD) for the micro-level comparisons, and sMAPE for the macro-level ones. Figure (ref) summarizes the results of our experiments: the left panel shows the microscopic evaluation of both our surrogate model and the ablated variant, while the right panel presents the macroscopic evaluation results.
For the Schelling model, we observe that, on the micro level, the surrogate’s mean EMD is lower than the diffusion-only ablation’s mean EMD in all cases. The differences between the surrogate and the diffusion-only ablation are less pronounced at the thresholds $\xi_1$ and $\xi_3$. At $\xi_1$ (few unhappy agents) behavior is almost entirely deterministic and agents rarely move, while at $\xi_3$ (almost all agents unhappy) behavior is uniformly random, so even a flat, “always‐move” or “never‐move” rule yields near‐optimal predictions in these two cases. In contrast, at the intermediate threshold $\xi_2$, where roughly half the agents are unhappy, the difference between the surrogate’s and the diffusion-only ablation’s EMD is more pronounced. A similar pattern is observed in the macroscopic evaluation. The surrogate’s sMAPE remains below 0.2, whereas the diffusion-only ablation fails to distinguish happy from unhappy cases, resulting in large macro‐level errors. In the GNN-only ablated model, at the micro-level EMD increases proportionally with the level of stochasticity introduced by the parameters, with EMD values similar to the surrogate for $\xi_1$ (the most deterministic) and much higher for $\xi_3$ (the most stochastic). At the macro-level, we observe large errors, with sMAPE always higher than the surrogate. These gaps confirm that only the full model, with explicit graph‐based interaction modeling and stochasticity, can learn the conditional relocation rule critical in balanced regimes.
For the Predator–Prey model, regarding micro-level behavior, our surrogate achieves a low EMD from the ground truth on average, and it consistently outperforms both ablation models. These results confirm that our model is able to faithfully reproduce the complex dynamics of this ABM even at the individual agent-level (thus explaining Figure (ref)). The most successful case is $\Psi_2$, where our surrogate exhibits a near-zero difference from the ground truth. In fact, in this configuration, most agents follow deterministic update rules (e.g., dead $\rightarrow$ dead), which are perfectly recovered by our model, but not by the diffusion-only ablation --- which also obtains worse results on stochastic rules, as shown in Supplementary Section (ref). Instead, the GNN-only ablation performs well only in those deterministic cases, but fails in all the others. At the macro level as well, the surrogate consistently outperforms the ablation, generally achieving low error. The best result is obtained with $\Psi_1$, the most complex dynamics, where the surrogate achieves an average sMAPE of approximately 0.08. This configuration produces two distinct population peaks, and the surrogate faithfully reproduces both their timing and amplitude (Figure (ref)). The worst result is obtained with $\Psi_2$, as this configuration is almost monotonic and dominated by long, near-zero tails that are noisy at very small scales, making them difficult for any model to reproduce.
We introduced Graph Diffusion Networks, a differentiable surrogate for agent-based models that combines graph neural networks to model agent-level interactions with diffusion models to capture stochasticity. Our experiments on the Schelling segregation model and a Predator–Prey ecosystem show that this approach not only accurately reproduces individual-level transition distributions, but also faithfully captures emergent, system-level dynamics beyond the training horizon.
\noindentLimitations. Our approach is limited by our assumptions about the characteristics of the ABM to emulate. First, the interaction graph is assumed to be fully known. Future work might remove this limitation by estimating such a graph directly from available data. However, the estimation of a latent interaction graph is a follow-up challenge, for which our GNN-based approach represents a necessary first step. Second, highly sophisticated ABMs may include features not addressed in our framework - such as all-to-all interactions, multiple rounds of decision-making, or sequential stochastic events within a single time step. Capturing these dynamics may require extending our architecture to incorporate sequential or hierarchical components. While our method may not yet fully generalize to such settings, our findings demonstrate that building surrogates capable of replicating individual-level behavior is both feasible and effective, laying the groundwork for broader applications.
\noindentFuture work. Building on this foundation, the differentiability of our surrogate opens up a range of powerful applications. It enables the use of gradient-based methods for any optimization task, such as policy optimization agrawal2025robust. It allows for efficient calibration of macro parameters by treating key parameters as additional inputs to the neural network. Most importantly, our approach naturally allows for the estimation of micro (i.e., agent) level variables — a challenge for ABMs, that often requires the ad hoc development of handcrafted probabilistic models monti2020learning, monti2023learning. In fact, our model already contains such parameters expressed as agents’ individual states ($\textbf{Z}^{(i)}_t$), something typically not available in ABM surrogates glielmo2023reinforcement. Moreover, our method can in principle be applied directly to real-world datasets whenever sufficient observations of comparable agent–context pairs and transitions are available. In doing so, our framework helps make ABMs more data-driven and empirically grounded, with promising applications in several scientific domains, such as economics pangallo2024datadriven, epidemiology gozzi2025epydemix, sustainability lamperti2025complex, urban science birkin2025digital, and ecology stock2024plant.
The authors wish to thank Daniele Grattarola and Federico Cinus for insightful early discussions that supported the initial development of this work. We also thank Alberto Novati for his contribution to the early draft of the code for the original ABM of the predator-prey system.