Technical Summary: A Hitchhiker's Guide to Poisson Gradient Estimation

1. Problem Statement

Discrete stochastic latent variable models, particularly those utilizing Poisson distributions, are fundamental to computational neuroscience for modeling neural spike counts. However, training these models via gradient-based optimization is challenging because standard backpropagation cannot differentiate through discrete, integer-valued samples.

While Gaussian Variational Autoencoders (VAEs) utilize the reparameterization trick for low-variance pathwise gradients, Poisson-distributed latents lack an analogous closed-form solution. Two primary approaches have emerged to address this:

1. Exponential Arrival Time (EAT): A relaxation based on point processes (Vafaii et al., 2024).
2. Gumbel-SoftMax (GSM): A relaxation based on the Concrete distribution applied to the Poisson probability mass function (Li et al., 2024a).

Despite both methods addressing the same fundamental challenge, there has been no systematic comparison between them. Consequently, practitioners lack guidance on method selection. Furthermore, existing implementations (specifically EAT using a sigmoid approximation) exhibit high sensitivity to the temperature hyperparameter ($\tau$), leading to unstable training and biased gradients.

2. Methodology and Technical Contributions

The paper introduces a systematic comparison of EAT and GSM, alongside a novel modification to the EAT method to improve its theoretical and empirical performance.

2.1. Theoretical Analysis of EAT

The authors analyze the standard EAT method, which approximates a Poisson sample $z \sim \text{Pois}(\lambda)$ by summing soft indicators of event arrival times:
$$z = \sum_{m=1}^M f_{\text{approx}}\left(\frac{1 - \sum_{i=1}^m \Delta t_i}{\tau}\right)$$
where $\Delta t_i \sim \text{Exp}(\lambda)$.

Using Campbell's theorem from point process theory, the authors derive closed-form expressions for the first and second moments of the relaxed distribution as a function of the temperature $\tau$ and the choice of soft indicator $f_{\text{approx}}$.

• Finding: The standard choice, the sigmoid function ($f_{\text{sigmoid}}$), possesses infinite support. This causes "ghost spikes" (fractional contributions from events far outside the unit interval), resulting in significant bias in the mean and a severe reduction in variance as $\tau$ increases.
• Correction: The authors propose replacing the sigmoid with a cubic Hermite interpolant (smoothstep), denoted as $f_{\text{cubic}}$. This function has compact support on $[-1, 1]$, eliminating contributions from events outside the relevant interval.

2.2. The Proposed Method: EATcubic

The modified method, EATcubic, replaces the sigmoid with the cubic smoothstep:
$$f_{\text{cubic}}(u) = \begin{cases} 0 & u < -1 \\ 3w^2 - 2w^3 & -1 \le u \le 1, \text{ where } w = \frac{u+1}{2} \\ 1 & u > 1 \end{cases}$$
Theoretical Guarantees:

• Unbiased Mean: For $\tau \le 1$, the mean factor $c(\tau)$ is exactly 1, guaranteeing an unbiased first moment that matches the true firing rate.
• Reduced Variance Bias: The variance factor $v(\tau)$ remains significantly closer to 1 compared to the sigmoid approximation.
• Fano Factor: The cubic approximation maintains a Fano factor (variance-to-mean ratio) closer to the Poisson value of 1, staying within biologically plausible ranges (0.5–1.5) across a wider temperature spectrum.

2.3. Evaluation Framework

The authors evaluate EATcubic against EATsigmoid, GSM, score-function estimators, and exact gradients (where analytically tractable) across three axes:

1. Distributional Fidelity: Measured via moment ratios (mean/variance) and Wasserstein-1 distance.
2. Gradient Quality: Measured via cosine similarity, bias energy, and noise energy (curvature-weighted metrics derived from a second-order Taylor expansion of the loss).
3. Downstream Performance: Evaluated on two tasks:
   • Poisson VAE (P-VAE): Generative modeling of natural image patches.
   • Partially Observable Generalized Linear Model (POGLM): Inferring neural connectivity from partially observed spike trains.

3. Key Results

3.1. Distributional Fidelity

• EATcubic achieves near-perfect mean fidelity ($c(\tau) \approx 1$) and substantially better variance fidelity across all tested temperatures and firing rates ($\lambda$).
• EATsigmoid exhibits severe mean inflation and variance collapse (e.g., retaining <20% of true variance at high $\lambda$ and $\tau$).
• GSM maintains low mean bias but suffers from higher variance bias compared to EATcubic.
• In terms of Wasserstein-1 distance, EATcubic consistently provides the closest approximation to the true Poisson distribution, outperforming alternatives by up to 7$\times$ at higher rates.

3.2. Gradient Quality

• Directional Alignment: Both EAT methods maintain near-perfect cosine similarity with exact gradients across temperatures. GSM performs well at higher temperatures but degrades at low $\tau$.
• Bias and Noise:
  • EATcubic exhibits extremely low bias energy ($\sim 10^{-2}$) across all temperatures.
  • EATsigmoid shows low noise energy but increasing bias energy as temperature rises.
  • GSM suffers from high noise energy at low temperatures.
• Dissociation: The paper notes a key dissociation: EATsigmoid, despite having the worst distributional fidelity, often produces gradient estimates with excellent noise characteristics, suggesting that distributional fidelity and gradient quality capture complementary aspects of relaxation quality.

3.3. Empirical Performance (ELBO and Task Metrics)

• Robustness to Temperature: EATcubic is the only method that maintains near-optimal ELBO performance across the entire tested temperature range ($\tau \in [0.02, 0.5]$).
• Sensitivity of Alternatives: EATsigmoid and GSM degrade significantly at higher temperatures (e.g., $\tau = 0.5$), often dropping to the performance level of high-variance score-function baselines.
• Real Data: On a retinal ganglion cell (RGC) dataset, EATcubic consistently recovers model parameters and achieves stable ELBO, whereas other methods fail to converge reliably without precise temperature tuning.
• Scaling: In larger-scale experiments (40$\times$40 ImageNet patches), EATcubic closes the gap to exact gradients, while EATsigmoid remains structurally inferior regardless of temperature tuning.

4. Significance and Recommendations

The paper concludes that EATcubic is the superior choice for Poisson latent variable models due to its:

1. Theoretical Unbiasedness: Exact matching of the first moment for $\tau \le 1$.
2. Temperature Robustness: Stable performance across a wide range of hyperparameters, reducing the engineering burden of grid searches.
3. Distributional Fidelity: Superior approximation of the true Poisson statistics, which is critical for tasks where the latent distribution's shape matters (e.g., neuroscience applications).

Limitations and Trade-offs:

• Generalizability: While EATcubic is optimal for Poisson distributions, the GSM method retains the unique ability to generalize to arbitrary discrete distributions (including over-dispersed Negative Binomial and under-dispersed regimes) without modification.
• Smoothness: The cubic smoothstep is $C^1$ continuous, whereas the sigmoid is $C^\infty$. In specific downstream tasks involving decoder weight visualization (feature learning), EATsigmoid produced slightly cleaner, more Gabor-like features, likely due to its higher-order smoothness.

Practical Recommendation:
For practitioners working with Poisson latent variable models (e.g., P-VAE, POGLM), EATcubic is recommended as the default choice. It offers the best balance of distributional accuracy, gradient quality, and robustness to hyperparameter selection. GSM should be reserved for scenarios requiring non-Poisson discrete distributions.

The authors provide code, data, and model checkpoints to facilitate the adoption of these methods.