We appreciate the reviewer’s detailed and constructive comments. We reply by numbering the weaknesses by W1, W2, … and the questions by Q1, Q2, …, grouping weaknesses and questions where relevant. For each reply, we start with a one-sentence TL;DR and then develop our argument.

W1, Q1: Clarity

TL;DR: We clarify our terminology and explain why the predator-prey and Schelling models are representative of ABMs.

We realize that the benefits of modeling individual behaviors were not adequately explained.

ABMs are based on defining rules at the level of individual agents – a perspective known as the micro-level. These rules deal only with the characteristics and behaviors of single agents, such as their type (e.g., predator or prey) or positions. Micro-level dynamics are governed by these rules, that act only on these individual features, such as updating an agent’s position over time. The purpose of ABMs is often to observe how these simple micro-level rules give rise to complex, emergent properties at the system-wide level, referred to as the macro-level. Examples of macro-level outcomes include summary statistics, such as the total number of surviving agents or the total number of agents that are happy of their neighbors, associated to segregation patterns.

With this background, we can now explain the value of modeling individual behaviors directly. In fact, we constrain our surrogate to model only micro-level rules. This ensures that the surrogate remains structurally faithful to the original ABM and does not rely on shortcuts such as directly predicting macro-level outcomes. To illustrate this, consider the predator-prey model. A macro-level surrogate might emulate the system using equations like the Lotka–Volterra model, which predicts population counts over time without modeling individual agents. In contrast, our surrogate expresses rules in terms of agent-level attributes, such as the type and position of each agent. This approach leads to two major benefits. First, by modeling only micro-level rules, the surrogate stays aligned with the logic of the original ABM. Second, the surrogate is expressed directly in terms of micro-level variables, allowing for real-data alignment with respect to individual agent attributes – something standard surrogates cannot do.

Traditional surrogates, in fact, can only perform macro-level calibration (e.g., tuning a global threshold like $\epsilon$ in the Schelling model). Our model, however, is differentiable with respect to micro-level inputs, such as the type or position of each individual agent. Thanks to this property, our approach offers practical advantages in integration with real-world data. For instance, when micro-level data is available for a subset of agents (e.g., demographic or spatial attributes), our surrogate can use this information to simulate possible future states.

Therefore, in response to Q1, to faithfully "align ABMs with real-world data," we mean to align the micro-level properties of individual agents to match observed data. Calibrating macro-level parameters only allows to calibrate some key macro-level quantities of the ABM, but its agents remains totally disconnected from any entity observed in the real-world. Consider again the predator-prey model: if we have a real-world data source enumerating some real-world animal positions, aligning this data with the ABM at the micro-level would mean associating one of such real positions to a simulated agent, instead of simply tuning a global parameter for the whole system. In this sense, aligning micro-level states with real-world observations remains an open challenge in the ABM literature, yet it is essential for ensuring that ABMs can be reliably applied in real-world contexts. Our model is the first differentiable surrogate capable of expressing learned stochastic rules in terms of such micro-level attributes. As such, it represents a crucial step toward building tools that can calibrate the micro-level components of ABMs, an often overlooked but essential aspect of systematically aligning ABMs with real-world data.

Regarding the representativeness of our choice of ABMs, both the Schelling model and predator-prey models are widely recognized as canonical examples of agent-based models. They are commonly used to introduce the field and are frequently employed in studies that explore calibration techniques and surrogate modeling.

The Schelling model is often cited as one of the earliest ABMs in the social sciences [1], and has served as the sole testbed for the development of data-driven model extraction methods [2]. Similarly, predator-prey models form the foundation for a broad range of more complex models in ecological research [3, 6], and they have played a key role as testbeds for the development of novel calibration and surrogate techniques. For example, in "Data assimilation with agent-based models using Markov chain sampling" [4], the authors test their method for an MCMC-based sampler of possible trajectories on a predator-prey ABM, very similar to ours. It is their only case study. Likewise, in "The use of surrogate models to analyse agent-based models" [5], one of the three case studies used to compare macro-level surrogates is a predator-prey model. It is worth noting that in all of these works, the predator-prey model is used as a representative of a broader class of ABMs, likely due to its simplicity, intuitive appeal, and ease of explanation to readers across different disciplines.

We will add these considerations and further clarifications to the final version of the manuscript.

[1] Squazzoni, F. (2010). The impact of agent-based models in the social sciences after 15 years of incursion. History of economic ideas: XVIII, 2, 2010, 1000-1037.

[2] Jamali, Ruhollah, Wannes Vermeiren, and Sanja Lazarova-Molnar. "Data-driven agent-based modeling: Experimenting with the Schelling’s model." Procedia Computer Science 238 (2024): 298-305.

[3] McLane, Adam J., et al. "The role of agent-based models in wildlife ecology and management." Ecological modelling 222.8 (2011): 1544-1556.

[4] Tang, Daniel, and Nick Malleson. "Data assimilation with agent-based models using Markov chain sampling." Open Research Europe 2 (2022): 70.

[5] Ten Broeke, Guus, et al. "The use of surrogate models to analyse agent-based models." Journal of Artificial Societies and Social Simulation 24.2 (2021).

[6] Murphy, Kilian J., Simone Ciuti, and Adam Kane. "An introduction to agent‐based models as an accessible surrogate to field‐based research and teaching." Ecology and evolution 10.22 (2020): 12482-12498.

W2, Q2: Baselines

TL;DR: Comparisons to baselines is difficult given the novel focus of our approach.

All reviewers raised concerns about the lack of baseline comparisons. We acknowledge this and have expanded our evaluation accordingly. However, our work introduces the first individual-level stochastic surrogate for ABMs—a direction that, to our knowledge, has not been explored before. As such, there are no directly comparable baselines. This is similar to the situation faced by Grattarola et al. (NeurIPS 2021), who introduced the first GNN-based surrogate and were also unable to compare against prior methods.

Below, we summarize our responses to the baseline-related concerns raised by the other reviewers.

Reviewer BiVF suggested we focus on macro-level surrogates. However, most of these surrogates do not model system dynamics directly, but rather map parameters to summary statistics. While a few exceptions exist (e.g., Fadikar et al., Kieu et al.), they still model aggregate variables and struggle to generalize beyond in-sample dynamics. For instance, an AR(1) time-series model fails to reproduce even simple non-monotonic patterns like predator-prey cycles, significantly underperforming compared to our individual-level approach.

In response to reviewer McsJ suggestions, we conducted an additional evaluation, training a GNN-only ablation. While this version performs adequately/decently in deterministic settings, it struggles with stochastic dynamics, showing much higher errors than the full GDN, particularly in the predator-prey model. In the Schelling model, where behavior is more deterministic under certain parameterizations, the GNN-only performance was closer to GDN. Reviewer McsJ also suggested that we considered alternative GNNs (e.g., Grattarola et al., Gravina et al.), but these alternatives suffer from the same structural limitation of our ablation study: they predict deterministic outcomes and cannot model stochastic transitions. This is supported by the underperformance of the GNN-only ablation. Another proposed baseline is the work from Casert et al., a promising Transformer-based stochastic emulator. However, their method is not applicable to our setting. It assumes fixed-size state spaces, continuous time, and explicit enumeration of all allowed transitions — constraints incompatible with ABMs like Schelling or Predator-Prey, where state spaces are large and dynamic.

We will include the new experiments, along with the AR(1) baseline and an extended discussion, in the revised paper and supplement.

Finally, the reviewer raised a concern regarding the lack of related work, requesting that we “provide the related works.” While we will revise the manuscript to include the expanded discussion presented above, we believe that Section 2 already covers the key contributions of the most relevant prior work in this area. That said, we are fully open to including additional references and discussion if the reviewer can point to specific missing works or directions that merit inclusion.

We welcome any follow-up questions and would be happy to provide further details or clarifications.

Thanks for the response. It addresses my questions.

We appreciate the reviewer’s detailed and constructive comments. We reply by numbering the weaknesses by W1, W2, … and the questions by Q1, Q2, …, grouping weaknesses and questions where relevant. For each reply, we start with a one-sentence TL;DR and then develop our argument.

W1: Need access to the ABM

TL;DR: ABMs cover lots of different systems, but in principle, the method can be used on data with no existing ABM too, granted that sufficient agent/context pairs and transitions are observed in the data.

Agent-based models (ABMs) are widely used to represent individual behavior and interactions across disciplines, and after more than three decades of development, it is likely that for many real-world datasets, an ABM exists that closely approximates the underlying data-generating process. The central challenge lies in aligning such ABMs with observational data.

That said, our method is not limited to synthetic data. In principle, it could be applied directly to real-world data, provided that the data includes sufficient repeated instances of similar individuals facing similar interaction structures. Taking the Schelling model as an example: in our synthetic setup, we generated 500 alternative futures for an agent moving to a new location given the characteristics of its neighbors. If we had real-world data from, say, 500 households with matching characteristics and neighborhood compositions, their observed behaviors could serve as a real-world ramifications dataset.

While this paper focuses on synthetic data as a necessary first step, we are actively exploring applications to real-world datasets as part of ongoing and future research.

W2, W3, Q2: Baselines

TL;DR: GNN models completely fail to learn stochastic behavior, and thus lead to much lower micro- and macro performance, while alternative baselines are not general enough for ABMs.

The reviewer proposes three further possible choices for ablations/baselines:

1. Perform an ablation that removes the diffusion component while keeping the GNN.
2. Consider alternative GNN models (e.g., Grattarola et al., temporal GNNs).
3. Compare to Casert et al. (2024).

We tackle each of these one by one.

1. Ablation without diffusion

We conducted additional experiments using the ablated model suggested by the reviewer, which excludes the diffusion component while retaining the GNN. This modification is non-trivial, as the original surrogate relies on gradients from the diffusion model to train the GNN. In the ablated version, we instead minimize the MSE between the true and predicted next agent states. The only architectural change is in the MLP following the aggregation pass: its output dimension is set to the agent state dimension rather than the embedding dimension. While the original GDN uses hidden layers of size [32, 64, 128], the ablated model adopts a symmetrical structure: [32, 64, 128, 128, 64, 32]. We trained for 100 epochs with the Adam optimizer (learning rate 2×10⁻⁵) and mini-batch size of 16.

Our results clearly show that the GNN successfully learns to predict agents’ future states when the underlying dynamics are deterministic. However, it performs poorly in the presence of stochastic rules, resulting in significantly lower overall performance compared to the GDN.

The results for the predator-prey model are summarized in the tables below (mean±standard deviation over 8 runs). For each parameter, the second and third columns reproduce the results from Figure 4 (GDN and diffusion-only ablation), while the last column reports results from the GNN-only ablation as suggested by the reviewer. We observe that both macro-level and micro-level errors are substantially higher in the GNN-only ablation, and often worse than the diffusion-only model. This is expected, as stochastic rules are dominant in most parameterizations, making the GNN's deterministic forecasting approach inadequate.

Predator-Prey macro-level:

Predator-Prey micro-level:

A slightly different picture emerges in the Schelling model for parameterization $\xi_1$. Here, nearly all agents remain stationary during the test period, leading to effectively deterministic dynamics. Consequently, the GNN-only surrogate achieves a micro-level EMD of 0.73±0.07, which is comparable to the GDN’s performance (0.71±0.06). However, from preliminary results (mean and std over 6 experiments, instead of 8 due to time) on parameterization $\xi_3$ , the most stochastic configuration, the GNN-only surrogate performs poorly with EMD (9.73±0.06) against GDN’s performance (2.81±0.11).

In the final version of the paper we will include the full results from this additional ablation study.

2. Alternative GNN models

We argue that all GNN-based models share the same structural limitations observed in the GNN-only ablation examined above, and therefore exhibit similar performance. This class includes the approach by Grattarola et al., which in fact was applied to deterministic cellular automata, as well as the temporal GNNs discussed in the reference suggested by the reviewer (Gravina et al. 2024)

3. Casert et al.

The reviewer suggested using the Transformer-based model by Casert et al. as a baseline. While it is the only stochastic emulator we are aware of, it is not applicable to the ABMs we consider. Their method is tailored to physical systems in continuous time and relies on explicitly enumerating allowed transitions between configurations. This is feasible when transitions involve small local changes, as in particle systems, but not in ABMs where agents can make large, unconstrained moves (e.g., in Schelling) or where the number of agents varies over time (e.g., predator-prey). In contrast, our ramifications approach samples from the space of possible configurations without requiring predefined transitions, making it suitable for ABMs but incompatible with the Casert et al. methodology.

W4, Q1: Scaling

TL;DR: Our surrogate model training cost scales linearly w.r.t. time-horizon and number of ramifications and sublinearly w.r.t. number of agents. Increasing time-horizon and number of agents increases stochastic coverage, so the number of ramifications can be reduced.

We acknowledge the need to discuss computational cost and scalability. Training our GDN on 9 timesteps, 500 ramifications, and ~2 000 agents (1 950 Schelling, 2 048 Predator–Prey) takes ≈ 37 s per epoch. Because each step processes one timestep and one ramification, runtime grows linearly with the number of ramifications or timesteps and only sub-linearly with the agent count (see reply to Reviewer BiVF). We use multiple stochastic branches during training to increase the statistical coverage of the ABM micro-dynamics. A new experiment on the predator–prey model (parameter set $\Psi_1$​) confirms that increasing R improves EMD and sMAPE until they match the main-paper results, while training time scales linearly.

*: Mean values from the paper micro and macro evaluations.

We also agree that generating hundreds of futures for each state may be prohibitive in systems with many agents or longer time horizons. However, larger systems naturally provide more independent samples of similar conditions. In such cases, fewer ramifications may suffice to achieve similar statistical coverage. To explore the trade-off between data generation cost and predictive quality, we trained GDN on three Predator-Prey datasets (parameter $\Psi_1$) with increasing agent counts and decreasing ramifications, keeping the number of agent transitions roughly constant. All other hyperparameters remain as in the main experiments. Results are shown below.

*: Mean values from the paper micro and macro evaluations.

We find that sMAPE remains close to the original results (and is notably lower for the 48-grid case), while EMD increases slightly but stays near the reported average. We will include these scaling results and the discussion above in the revised manuscript and supplement to better communicate the computational footprint and scalability of our framework.

Q3: GNN details

We use a simple MessagePassing GNN, where each agent’s state is passed as the message, aggregated, and fed into an MLP. As the reviewer has correctly pointed out, the choice of the MessagePassing aggregation scheme can impact expressiveness and inductive bias. The most commonly used aggregations for this setting are sum, mean and max. In our experiments we used “sum” for Predator-Prey and “mean” for Schelling since there is dependency on the degree in its dynamics. Minimal experimental evaluations allow to direct the practitioner towards the one which seems promising for the task at hand. In the future we would like to explore ‘aggregation-agnostic’ approaches, in order to reduce inductive bias. Promising options are relying on Attention weighting or leveraging PowerMeanAggregation or MultiAggregation operators from Pytorch.

We welcome any follow-up questions and would be happy to provide further details or clarifications.

We thank the reviewer for the thoughtful and constructive feedback. We reply by numbering the weaknesses by W1, W2, … and the questions by Q1, Q2, …, grouping weaknesses and questions where relevant. For each reply, we start with a one-sentence TL;DR and then develop our argument.

W1: Scaling

TL;DR: Our surrogate model has been tested on mid-sized ABMs, the training time scales linearly with number of ramifications and timesteps, and sublinearly with number of agents.

We agree that the main manuscript should state the concrete dimensions of the testbeds and clarify how the method scales. We will move the numbers that are currently in the supplement to the main text (§4) and add the following discussion.

These settings are in the range of mid-sized ABMs: each training epoch processes about $\sim 9\times10^{6}$ agent transitions (9 time steps x 2000 agents x 500 ramified futures). On a single RTX-A6000, one epoch takes ≈ 37 seconds.

The reviewer is correct that an ABM containing millions of agents can stress any learning algorithm. As we have mentioned in the discussion with Reviewer McsJ, the training time of the model grows linearly with the number of ramifications provided to the network during training, as well as the number of timesteps. On the other hand, increasing the number of agents in the system (for example by increasing the grid size and keeping the density of agents constant) increases the training time sublinearly. In fact, with mid-sized systems, we can train the network with batches that correspond to the whole set of agents. Thus, the number of training iterations per epoch does not change, but rather the size of the batch in input to the network. Considering the Predator-Prey model, if we increase the grid size from 32 to 64, going from 2048 agents to 8192 (4 times as many agents), we have an increase from 37.09 s/epoch to 108.37 s/epoch, which is 2.92 times larger. Similarly, by increasing from grid size 32 to 48, we increase by a factor of 2.25 the number of agents, but the training time increases by a factor of 1.72. Regarding inference, the computational cost to produce the next agent state is always the same. However, it will take longer for larger systems to simulate the evolution of the entire system. Below we show a table with the training time and inference time for the 3 settings of grid sizes we have discussed, where inference time refers to the time necessary to generate 1 timestep of a simulation of N. agents with the trained GDN. All other hyper-parameters are the same as in the paper experiments. We include the evaluation metrics at the macro and micro level used in the paper, to assess that even with a larger system the dynamics are learned by our methodology.

*: Mean values from the paper micro and macro evaluations.

Training the model with a ramified dataset of R ramifications is a choice made to allow the model to be exposed to many agent transitions ($\sim 9\times10^{6}$ in our experiments). Increasing the number of agents in the system increases the number of agent transitions available in the dataset, expanding the statistical coverage provided to the model to learn. In our previous experiment, the models with higher numbers of agents performed better both at the macro and micro level than in the original experiments. The number of ramifications R could be decreased to potentially get the same result quality of the paper experiments. We ran an experiment by adjusting the number of ramifications R according to the number of agents so that we have a similar number of agent transitions in the training set as in the original experiments. Given the time constraints of the rebuttal period, we were able to train one model for each grid size configuration on one dataset of Predator-Prey with parameter $\Psi_1$ and evaluate the model as done in the paper. The results are shown in the table below. We find that the sMAPE is around the value found in the original experiments (actually lower for grid size 48), and the mean EMD is slightly higher, but close to the mean value from the original experiments. We will add the clarifying discussion above and additional scaling experiments with more hyperparameter values to the revised manuscript and supplement material.

*: Mean values from the paper micro and macro evaluations.

Inference time refers to the time needed to produce 1 timestep of a simulation of N agents with the trained GDN.

Regarding scaling to even larger systems (e.g. in the range of $10^4$ to $10^6$), it would be more appropriate to train the model on batches made of subgraphs of the system configuration, as it would become troublesome to optimize the model on batches of tens of thousands or millions of samples. However, apart from batch optimization, the methodology would remain unchanged.

W2, Q1: Macro-level baselines

TL;DR: While surrogate models for aggregate dynamics do exist, they struggle to reproduce trends that are absent from the training data. In contrast, our approach (based on learning individual behavior) can generalize to these unseen dynamics. Most existing surrogate approaches do not actually learn system dynamics. Instead, quoting from our manuscript, surrogate models typically learn directly the mapping from parameters $\Theta$ to static summary statistics, disregarding individual behavior and model dynamics.

There are a few exceptions. For example, Fadikar et al. [1] use functional Gaussian Processes to capture the shape of epidemic curves, while Kieu et al. [2] employ LSTMs to forecast time series of pedestrian counts. However, these approaches share two key limitations: (i) they model aggregate variables rather than micro-level behavior, and (ii) they rely on statistical models fitted to in-sample data.

As a result, these models perform well when out-of-sample dynamics resemble those seen during training, but their performance deteriorates when the dynamics change. For instance, Fadikar et al. focus on cumulative case counts in epidemics, which follow a monotonic trajectory. In such cases, using early time steps to predict future ones is a reasonable strategy. But this assumption breaks down in more complex settings.

Consider the predator-prey model shown in Figure 3 (top panels), where both predator and prey populations grow and then decline. These dynamics are far from monotonic. A surrogate trained on aggregate counts during the early phase (when populations are rising) is unlikely to capture the subsequent decline.

To illustrate this point, we implemented a simple AR(1) process as a representative of aggregate time-series-based models. We trained it on just 10 time steps (matching the surrogate’s training window). Given this limited horizon, using more complex models like LSTMs would be inappropriate.

The AR(1) surrogate substantially underperforms our individual-level surrogate. For example, in scenario $\Psi_1$—characterized by oscillatory predator-prey dynamics—the AR(1) either predicts indefinite growth or convergence to a steady state. It fails to capture the true cyclical behavior, resulting in a significantly higher sMAPE. Even in the simpler scenario $\Psi_2$, where dynamics are monotonic, the AR(1) model incorrectly predicts convergence to a steady state (e.g., a predator and prey population of 60), again leading to large errors.

The table below reports the sMAPE (mean and standard deviation) across four scenarios. As shown, the AR(1) surrogate performs poorly for all scenarios, with errors far exceeding those of our micro-founded surrogate reported in Figure 4.

In summary, we deliberately chose not to include comparisons with macro-level dynamic surrogates in the main paper because their limitations are structural: they rely entirely on aggregate time series observed during training. As such, they are fundamentally incapable of extrapolating to qualitatively different out-of-sample dynamics—precisely the type of generalization our setting requires. The AR(1) results shown here support this point, illustrating that even the simplest non-monotonic patterns are missed. For this reason, we maintain that comparing our surrogate to these macro-level approaches would not provide a meaningful benchmark. We will include this discussion and the results on the AR(1) in the revised paper.

[1] Fadikar, A., Higdon, D., Chen, J., Lewis, B., Venkatramanan, S., & Marathe, M. (2018). Calibrating a stochastic, agent-based model using quantile-based emulation. SIAM/ASA Journal on Uncertainty Quantification, 6(4), 1685-1706.

[2] Kieu, M., Nguyen, H., Ward, J. A., & Malleson, N. (2024). Towards real-time predictions using emulators of agent-based models. Journal of Simulation, 18(1), 29-46.