Inference in Spreading Processes with Neural-Network Priors

Stochastic processes on graphs are a powerful tool for modelling complex dynamical systems such as epidemics. A recent line of work focused on the inference problem where one aims to estimate the state of every node at every time, starting from partial observation of a subset of nodes at a subset of times. In these works, the initial state of the process was assumed to be random i.i.d. over nodes. Such an assumption may not be realistic in practice, where one may have access to a set of covariate variables for every node that influence the initial state of the system. In this work, we will assume that the initial state of a node is an unknown function of such covariate variables. Given that functions can be represented by neural networks, we will study a model where the initial state is given by a simple neural network – notably the single-layer perceptron acting on the known node-wise covariate variables. Within a Bayesian framework, we study how such neural-network prior information enhances the recovery of initial states and spreading trajectories. We derive a hybrid belief propagation and approximate message passing (BP-AMP) algorithm that handles both the spreading dynamics and the information included in the node covariates, and we assess its performance against the estimators that either use only the spreading information or use only the information from the covariate variables. We show that in some regimes, the model can exhibit first-order phase transitions when using a Rademacher distribution for the neural-network weights. These transitions create a statistical-to-computational gap where even the BP-AMP algorithm, despite the theoretical possibility of perfect recovery, fails to achieve it.

💡 Research Summary

This paper addresses the problem of inferring the full state trajectory of a stochastic spreading process on a graph when only partial observations are available. Unlike most prior work, which assumes that the initial sources of the spread are independent and identically distributed (i.i.d.) across nodes, the authors propose to incorporate node‑wise covariate information through a neural‑network prior. Specifically, they model the initial binary state x₀ᵢ (source or susceptible) as the output of a single‑layer perceptron acting on an M‑dimensional feature vector Fᵢ associated with each node:

 x₀ᵢ = sign(∑ₐ uₐ F_{ia} − κ),

where the weight vector u is unknown but drawn from a known distribution (Gaussian or Rademacher) and κ controls the overall source density. The ratio α = N/M is kept finite as both the number of nodes N and the feature dimension M tend to infinity, thereby providing a tunable signal‑to‑noise parameter that governs the correlation among the initial states.

The spreading dynamics themselves are defined on a sparse graph G(V,E) and can be either the classic Susceptible‑Infected (SI) model or a deterministic Susceptible‑Infected‑Recovered (dSIR) variant. Transition times tᵢ are used to compactly encode each node’s trajectory, reducing the total number of variables from O(N·T) to O(N). The likelihood of the observed data O (either a set of sensor trajectories or a single snapshot at a fixed time) factorises over nodes given the transition times.

Within a Bayesian framework the posterior over the unknowns {tᵢ}, x₀, and u is proportional to a product of three types of factors: (i) sparse local factors Ψᵢ encoding the spreading kernel, (ii) dense factors Ψ_outᵢ encoding the perceptron prior, and (iii) the prior over the weights ψₐ(uₐ). This hybrid factor graph naturally suggests a combination of two message‑passing schemes: belief propagation (BP) for the sparse part and approximate message passing (AMP) for the dense, high‑dimensional linear part. By applying the cavity method, the authors derive a set of coupled BP‑AMP update equations that iteratively approximate the marginal posterior distributions of tᵢ and x₀ᵢ.

Performance is measured by two overlap metrics: the raw overlap O between the estimated source vector and the ground truth, and the mean overlap MO obtained by the Bayes‑optimal maximum‑mean‑overlap (MMO) estimator, which selects for each node the state with the highest posterior marginal. The authors verify that, under Nishimori conditions, the empirical O and MO coincide, confirming that the algorithm respects Bayes‑optimal self‑consistency.

Extensive simulations on Erdős‑Rényi and random‑regular graphs explore a range of α, weight distributions, and observation settings. When α is large (i.e., many features per node) and the weights are Gaussian, the hybrid BP‑AMP algorithm substantially outperforms either pure BP (which ignores covariates) or pure AMP (which ignores the graph dynamics). In contrast, with Rademacher weights and small α, the system exhibits a first‑order phase transition: the information‑theoretic limit predicts perfect recovery, yet the BP‑AMP dynamics become trapped in a suboptimal fixed point, revealing a statistical‑to‑computational gap. This phenomenon mirrors similar gaps observed in other high‑dimensional inference problems and underscores the importance of algorithmic design when dense priors are involved.

The paper’s contributions are threefold: (1) introducing a tractable neural‑network prior for the initial condition of spreading processes, (2) developing a principled BP‑AMP algorithm that merges sparse and dense graphical structures, and (3) uncovering regimes where the algorithm achieves Bayes‑optimal performance and regimes where a computational barrier prevents it from doing so. The work opens several avenues for future research, including extensions to deeper neural networks, non‑linear activations, and low‑source‑density regimes where replica symmetry breaking may occur. Overall, the study demonstrates that leveraging node‑level covariates through neural‑network priors can dramatically improve inference in epidemic‑type dynamics, with potential applications across epidemiology, information diffusion, and biological network analysis.