Stochastic Message Passing for Epidemics Modeling on Social Networks

The behavior of stochastic social phenomena such as epidemics or crowd movements is often unknown and rarely easy to model. This work presents a framework for stochastic graph neural networks (SGNN) training designed to learn complex stochastic behavior from graph data. Unlike standard graph neural networks, the proposed architecture introduces stochastic message propagation through a continuous relaxation of Bernoulli transmission events. Repeated rollouts from the same graph initial condition generate multiple possible trajectories, from which node-level marginal state distributions are estimated. The model evolves recurrently over time where at each time step, stochastic messages are propagated along eligible edges, aggregated at the receiving nodes, and combined with the current node states through a multilayer perceptron. We evaluate the model on synthetic graphs and real social networks using the susceptible-infected (SI) epidemic process as a case study. The evaluation examines both the model's ability to reproduce the temporal evolution of node states and its ability to recover the underlying transmission probability. We further analyze the sensitivity of parameter recovery and predictive performance of two factors: the number of model simulations and the initial value of the learnable transmission probability. This work represents the first step towards graph neural networks that can act as stochastic simulations of graph-based phenomena.