Branching Diffusion for Point Processes in Time and Space
Abstract
We propose a non-autoregressive branching diffusion model for generating spatio-temporal point processes. Starting from a geometric principle---the Wasserstein-Fisher-Rao (WFR) gradient flow of a generalized KL divergence toward a simple reference intensity---we obtain a tractable forward noising mechanism with two interpretable components: (i) a Langevin-type \emph{drift-diffusion} step that perturbs event locations and times, and (ii) a \emph{birth-death branching} step that changes the event count via location-dependent thinning (deaths) and Poisson offspring replication (births). We learn the reverse-time dynamics using a permutation-equivariant denoiser that predicts a drift field and a net-growth field, and we train it using an entropic-regularized unbalanced optimal transport (UOT), which naturally handles count mismatch between noisy and clean samples. The resulting generator produces complete spatio-temporal event sets without autoregressive simulation or explicit intensity normalization.
Lay Summary
Many real-world phenomena, such as earthquakes, crimes, policy changes, and social-media activity, occur as events distributed across time and space. Modeling these data requires generating not only when and where events happen, but also how many events occur. Existing point-process models often generate events one by one, which can be slow and may miss the global structure of an entire event pattern. This paper proposes a diffusion-based model that generates a complete spatio-temporal event set in a non-autoregressive way. The key idea is to treat events as points that can move, disappear, or give birth to new events during the noising process, and then learn a neural denoising process to reverse these changes. The model is grounded in Wasserstein-Fisher-Rao geometry, which naturally handles both spatial-temporal movement and changing event counts. Experiments show that it can generate realistic event patterns on synthetic and real-world datasets.