A hitchhiker's guide to Poisson gradient estimation
Abstract
Poisson-distributed latent variable models are widely used in computational neuroscience, but differentiating through discrete stochastic samples remains challenging. Two approaches address this: Exponential Arrival Time (EAT) simulation and Gumbel-SoftMax (GSM) relaxation. We provide the first systematic comparison of these methods, along with practical guidance for practitioners. Our main technical contribution is a modification to the EAT method that theoretically guarantees an unbiased first moment (exactly matching the firing rate), and reduces second-moment bias. We evaluate these methods on their distributional fidelity, gradient quality, and performance on two tasks: (1) variational autoencoders with Poisson latents, and (2) partially observable generalized linear models, where latent neural connectivity must be inferred from observed spike trains. Across all metrics, our modified EAT method exhibits better overall performance (often comparable to exact gradients), and substantially higher robustness to hyperparameter choices. These results extend to over-dispersed Negative Binomial latents, where modified EAT again performs best. However, only GSM generalizes to arbitrary non-Poisson distributions, including the under-dispersed regime. Together, our results clarify the trade-offs between these methods and offer concrete recommendations for practitioners working with Poisson latent variable models.
Lay Summary
Neurons communicate by firing brief electrical pulses called spikes, and computational models of the brain need a way to generate random spike counts that match real neural activity. The standard recipe for this produces integer numbers, which clashes with how modern machine-learning methods learn: through small, smooth adjustments. Two workarounds exist, but no one had carefully compared them before. We provide the first head-to-head comparison, introduce an improved version of one of the methods, and show that it gives more accurate, more reliable results without the careful manual tuning the alternatives require.