A differentiable model of supply-chain shocks

Abstract

Modelling how shocks propagate in supply chains is an increasingly important challenge in economics. Its relevance has been highlighted in recent years by events such as Covid-19 and the Russian invasion of Ukraine. Agent-based models (ABMs) are a promising approach for this problem. However, calibrating them is hard. We show empirically that it is possible to achieve speed ups of over 3 orders of magnitude when calibrating ABMs of supply networks by running them on GPUs and using automatic differentiation, compared to non-differentiable baselines. This opens the door to scaling ABMs to model the whole global supply network.

1 Introduction

Supply-chain disruptions cause macroeconomic damages IMFSC2022, and yet, despite an established research tradition leontief1936; acemoglu2012network; baqaee_2018_cascading, they remain difficult to model realistically CNN2025Powell. This challenge arises from fundamental limitations in existing approaches. Traditional economic models rely on comparative statics and cannot capture the explicitly dynamical phenomena that characterize disruption propagation through supply networks. Many modern machine-learning methods are similarly constrained by data scarcity: the historical record contains only a few disruptions (e.g., pandemics or natural disasters) on which to train, making it difficult to generalize across different types of shocks or network structures.

Agent-based models (ABMs) (farmer_review) offer a promising alternative to these limitations. ABMs are a computational paradigm where a collection of agents and their interactions are modelled explicitly in a simulator. Unlike statistical methods, ABMs provide mechanistic models that recover macroscopic quantities from the bottom-up through emergent processes. ABMs are therefore a natural choice for modelling supply chains disruptions, by capturing the micro-dynamics of each firm’s production and inventory management (Hallegatte2008-lq; Hallegatte2014-km; Inoue2019-ka).

However, ABMs are hard to calibrate. These models contain latent parameters that must be calibrated to match real-world observations according to some preferred metric. For example, firm productivity may not be directly observable and must be inferred from macroscopic quantities. Calibration requires evaluating the simulator at multiple points in parameter space — potentially thousands or millions of times — incurring very high computational costs when the simulator is slow to run. Furthermore, the intractable likelihood function of ABMs prevents straightforward application of Bayesian calibration methods, hindering proper uncertainty quantification.

Several approaches have emerged to mitigate these calibration challenges. These range from developing surrogate methods that mimic ABM outputs at greater speed (cozzi2025learningindividualbehavioragentbased; june_hist_match), to neural-based approaches that directly learn likelihood or posterior distributions (DYER2024104827). There has also been significant effort in tensorizing ABMs to leverage GPU acceleration, achieving substantial speedups (chopra_differentiable_2023). Furthermore, applying automatic differentiation to ABMs enables fast sensitivity analyses that can greatly accelerate calibration methods (chopra_differentiable_2023; https://doi.org/10.48550/arxiv.2509.03303).

In this work, we present a supply chain ABM implemented in JAX (jax2018github), a differentiable programming, high-performance numerical computing framework for Python. We demonstrate how supply chain models can be made amenable to tensorization across thousands of firm agents, achieving substantial speed-ups compared to traditional implementations. Furthermore, by leveraging JAX’s automatic differentiation (AD) capabilities and its probabilistic programming library Numpyro (phan2019composable), we apply generalized variational inference to perform efficient ABM calibration with proper uncertainty quantification over thousands of parameters.

2 Background and Related Work

Attempts to model production networks date back to Leontief’s Input-Output framework leontief1936, later extended to dynamic and firm-level contexts by Long and Plosser long_plosser_1983_rbc, Acemoğlu et al. acemoglu2012network, and Baqaee and Farhi baqaee_2018_cascading; baqaee_farhi_2019_micro. These models, widely used in the mainstream economic literature, study shock propagation through comparative statics: they compute equilibrium allocations before and after an exogenous productivity shocks, and compare outcomes such as aggregate output. While useful to understand the effects of network topology on shock propagation, this approach misses dynamical features like inventory adjustment or time-dependent recovery.

To overcome this, researchers have turned to agent-based models (ABMs) which model supply-chain dynamics from the bottom-up. Building on the ARIO model of Hallegatte Hallegatte2008-lq; Hallegatte2014-km, recent work has replicated real-world shock propagation such as the 2011 Tohoku earthquake Inoue2019-ka, the effect of the Covid-19 lockdowns in the UK’s economy Pichler2022 or the effect on the Austrian economy of an embargo on Russian gas Pichler2024. However, none of these works attempted a serious calibration of their models beyond parameter exploration.

Parameter estimation in ABMs has traditionally relied on Approximate Bayesian computation (ABC) methods, which sidestep intractable likelihood computation by matching summary statistics between simulated and observed data. While this approach provides a general framework for ABM calibration, it suffers from sensitivity to summary statistic selection and poor scaling in high-dimensional parameter spaces (PLATT2020103859). These limitations have motivated the development of more sophisticated methodological approaches.

One prominent strategy develops surrogate models that approximate ABM behavior while maintaining computational tractability and differentiability. By replacing the original ABM with these more amenable representations, researchers can apply established parameter estimation techniques, including Markov Chain Monte Carlo and gradient-based variational inference june_hist_match; cozzi2025learningindividualbehavioragentbased; monti_learning_2023; lamperti2017agentbasedmodelcalibrationusing.

Neural methods represent another paradigm, directly learning the mapping from simulator parameters to posterior distributions (DYER2024104827; BoeltsDeistler_sbi_2025). These approaches offer compelling advantages through amortized inference—once trained, they provide immediate posterior estimates for new observations without retraining. However, their inability to exploit gradient information limits efficiency in high-dimensional parameter spaces.

Most recently, advances in differentiable programming have enabled a third approach that exploits the inherent computational structure of ABMs. Network-based ABMs in epidemiology have demonstrated high degrees of tensorization, achieving substantial GPU acceleration (chopra_differentiable_2023). By applying relaxation techniques to discrete randomness and control flow structures, these methods integrate ABMs directly into AD frameworks, enabling end-to-end optimization without the approximation errors of surrogate models or the gradient limitations of neural approaches (https://doi.org/10.48550/arxiv.2509.03303).

3 Methods

3.1 Model

Figure 1: A: We model a directed production network where firms produce to serve demand from other firms, households, government, and foreign customers. B: In each period, an exogenous shock may limit each firm’s capacity to produce, then firms receive and place orders, produce subject to input inventories and a capacity limit, and update inventories.

We model a directed production network of $M$ firms, where each firm $i$ produces one good to serve demand from other firms (intermediate demand) as well as households, government, and foreign customers (final demand, $d_{i}$). Time is discrete: in each period, firms receive and place orders, produce subject to input inventories and a capacity limit, and update inventories; at the beginning of each period, an exogenous shock can decrease the maximum output that they can produce. Input requirements are encoded by a technical-coefficient matrix $\mathbf{A}$ (entries $A_{ij}$ are units of input $j$ required to produce one unit of output $i$), and inputs and outputs are tied by the formula $f\left(S_{ij}\right)=z_{i}\min\frac{S_{ij}}{A_{ij}}$, where $z_{i}$ is the productivity and set to $z_{i}=1$ in absence of a shock.

Firms try to maintain a target inventory $S^{\text{target}}_{ij}=n_{i}S_{ij}(0)$, enough to produce for $n_{i}$ iterations if they don’t receive supplies. Let total demand for good $i$ be $D_{i}(t)=d_{i}+\sum_{k}O_{ki}(t)$. Firms place orders to cover final demand and close inventory gaps, with an adjustment speed $1/\tau_{i}$, so $O_{ji}(t)=A_{ji}D_{i}(t)+\big[S^{\text{target}}_{ij}-S_{ji}(t-1)\big]/\tau_{i}.$

The actual production is the smallest among the demand received, and the firms’ production capacity, given productivity and input constraints: $x_{i}(t)=\min\left\{D_{i}(t-1),\;z_{i}(t)f\left(S_{ji}(t)\right)\right\},$where the productivity follows a shock-and-recovery process $z_{i}(t)=\left(1-\delta_{i}\exp\left(-\lambda_{i}(t-t^{*})^{+}\right)\right)$ for shocked firms, with $(t-t^{*})^{+}=\max\{0,t-t^{*}\}$ and $t^{*}$ the time of the shock. Inventories evolve by inflow minus use, $S_{ji}(t)=S_{ji}(t-1)+Z_{ji}(t)-A_{ji}x_{i}(t)$, where $Z_{ji}(t)$ denotes the realized deliveries from $j$ to $i$. We initialize at a steady state (orders match $D_{i}$ and inventories sit at target) and study how shocks propagate through orders, inventories, and capacity constraints.

The model produces individual time series for each firm $x_{i}(t)$, but in practice, we aggregate them to obtain macroeconomic aggregates that would match regional production indices. This is obtained, e.g., by defining $y_{r}(t)=\sum_{i\in r}x_{i}(t)$ where $1\leq r\leq N$ denotes the set of firms in a given sector and a given region.

3.2 Calibration

Let $\mathbf{y}$ denote the aggregate observed time series of firms’ outputs, $\mathbf{n}=\left[n_{1},\ldots,n_{N}\right]$ the model parameters — the array of firms’ inventory levels — with prior $p(\mathbf{n})$, and $\ell\!\left(\mathbf{y};\mathbf{n}\right)$ a loss function measuring the discrepancy between the observed data and the simulated results. We estimate the (generalized) posterior over parameters with Generalized Variational Inference (GVI) (knoblauch_optimization-centric_2022), which approximates $p(\mathbf{n}\mid\mathbf{y})$ by a distribution $q^{*}(\mathbf{n})\in\mathcal{Q}$ solving

$$q^{*}(\mathbf{n})\;=\;\arg\min_{q(\mathbf{n})\in\mathcal{Q}}\;\mathbb{E}_{\mathbf{n}\sim q}\!\left[\ell\!\left(y;\mathbf{n}\right)\right]\;+\;D\!\left(q(\mathbf{n})\,\|\,p(\mathbf{n})\right),$$ (1)

where $D\left(\cdot\|p\left(\mathbf{n}\right)\right)$ penalizes deviations from the prior. In practice, we parameterize $q\left(\mathbf{n}\right)$ using a set of variational parameters $\phi$ (e.g., $\mathcal{Q}$ can be the Gaussian family $q_{\phi}(\mathbf{n})=\mathcal{N}(\mathbf{n};\mu,\Sigma)$ with $\phi=(\mu,\Sigma)$) and minimize the objective in (1) with stochastic gradient descent. Note that when $\ell$ is the negative log-likelihood and $D$ is the Kullback–Leibler divergence, GVI becomes the usual variational inference. For our purposes, we take $l$ to simply be the $L_{2}$ loss, and $D$ to be the Kullback–Leibler divergence. As a gradient-free baseline, we utilise a variant of Approximate Bayesian Computation (ABC), where we repeatedly sample from the prior, and keep the top $k$ samples (for us, 100) with the lowest loss under $\ell$ and $\mathbf{n}$.

4 Results

Figure 2: Two plots showing the advantages of our methodology. A: we compare the speed of SVI calibration on GPU (RTX 5090) vs CPU (Ryzen 9 9950X). B: We show how the accuracy per model evaluation decreases between SVI and ABC.

To demonstrate the advantages of our methodology, we run two separate experiments, shown in Figure 2. Firstly, in Fig. 2 A, we show how our problem is naturally parallelizable by comparing the same code running on a GPU (RTX 5090) vs CPU (Ryzen 9 9950X). The benchmark used for comparison simply runs SVI for a fixed number of epochs, with different numbers of synthetic firms. The CPU-based implementation takes far longer, especially as the number of firms increase. On the other hand, the GPU-based implementation is much quicker, and doesn’t increase dramatically as the number of firms increase, this is because the GPU is not using all it’s parallel capacity, even at 3000 firms. As the number of firms increases, the difference between the two methods becomes only more dramatic.

In Fig. 2 B, we compare to a gradient-free baseline, showing that using a gradient based method yields huge speed-ups. We do this by taking our model, which consists of 1,000 firms, and sampling some true parameters $\mathbf{n}$ from the prior, 2 for each firm. We then generate observed data from $\mathbf{n}$, and then have the two methods train on this data, utilising the true prior. To measure potential over-fitting, we compare both the “in-sample” loss, that is the loss to the observed data, and the “out-of-sample” loss, which is the average loss to observed data re-sampled from the true prior, conditioning on the true parameters $\mathbf{n}$.

As we see from the experiment, SVI makes it possible to efficiently calibrate ABMs with 1000s of parameters, in contrast to gradient free methods like ABC. After only 300 model evaluations, SVI has a lower loss than ABC after 30,000 samples. Note that this comparison is slightly unfair, as in SVI, a model evaluation also includes a gradient pass, which adds computation that isn’t measured here. However, the gradient pass takes comparable time to the forward pass, and as a result, measuring wall clock time does not meaningfully change the results.

5 Conclusion

Identifying and mitigating vulnerabilities in supply chains is crucial to supporting the well-being and prosperity of the general public. Doing this effectively requires accurate models of production networks. Agent-based models (ABMs) are a promising approach, but are hard to calibrate. By implementing a 1,000-firm production-network ABM in JAX, we exploit parallelisation and automatic differentiation to perform Bayesian inference orders of magnitude faster than non-differentiable alternatives. This work not only scales ABMs but also enhances their realism—for example, by enabling simulations of price dynamics, logistics, and network restructuring in response to disruptions.