A hitchhiker's guide to Poisson gradient estimation
Michael Ibrahim, Hanqi Zhao, Eli Sennesh, Zhi Li, Anqi Wu, Jacob Yates, Chengrui Li, Hadi Vafaii
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 simulation (EAT; Vafaii et al., 2024) and Gumbel-SoftMax relaxation (GSM; Li et al., 2024a). 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. Our code, data, and model checkpoints are available here: github.com/hadivafaii/PoissonGradientEstimation
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 561aefd8-e550-4025-b410-649ec18f17f5Builds on5
- Poisson Variational AutoencoderHadi Vafaii, Dekel Galor, Jacob L. YatesNeurIPS 2024 · 18 citations
- Markov Chain Score Ascent: A Unifying Framework of Variational Inference with Markovian GradientsKyurae Kim, Jisu Oh, Jacob R. Gardner, Adji Bousso Dieng et al.NeurIPS 2022 · 12 citations
- Brain-like Variational InferenceHadi Vafaii, Dekel Galor, Jacob L. YatesNeurIPS 2025 · 7 citations
- Negative Binomial Variational Autoencoders for Overdispersed Latent ModelingYixuan Zhang, Jinhao Sheng, Wenxin Zhang, Quyu Kong et al.CVPR 2026 · 3 citations
- A Differentiable Partially Observable Generalized Linear Model with Forward-Backward Message PassingChengrui Li, Weihan Li, Yule Wang, Anqi WuICML 2024 · 3 citations
Related papers
- Cold Analysis of Rao-Blackwellized Straight-Through Gumbel-Softmax Gradient EstimatorAlexander ShekhovtsovICML 2023 · 2 citations
- Scalable inference of functional neural connectivity at submillisecond timescalesArina Medvedeva, Edoardo Balzani, Alex H. Williams, Stephen KeeleyNeurIPS 2025 · 2 citations
- Rao-Blackwellizing the Straight-Through Gumbel-Softmax Gradient EstimatorMax B. Paulus, Chris J. Maddison, Andreas KrauseICLR 2021 · 48 citations
- Training Discrete Deep Generative Models via Gapped Straight-Through EstimatorTing-Han Fan, Ta-Chung Chi, Alexander I. Rudnicky, Peter J. RamadgeICML 2022 · 9 citations
- Efficient Non-conjugate Gaussian Process Factor Models for Spike Count Data using Polynomial ApproximationsStephen L. Keeley, David M. Zoltowski, Yiyi Yu, Spencer L. Smith et al.ICML 2020 · 24 citations
