Branching Diffusion for Point Processes in Time and Space

Chao Yang, Wenjie Shen, Shuang Li

Keywords: Temporal/Spatio-Temporal Point Process, Diffusion Model, Unbalanced Optimal Transport

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 drift-diffusion step that perturbs event locations and times, and (ii) a 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.

Email Sharing: We authorize the sharing of all author emails with Program Chairs.

Data Release: We authorize the release of our submission and author names to the public in the event of acceptance.

Submission Number: 23