NegBio-VAE: Negative Binomial VAE

• The paper introduces NegBio-VAE, a novel VAE framework that employs a negative binomial distribution to decouple mean and variance in overdispersed count data.
• It features innovative reparameterization methods, including Gumbel-Softmax and continuous-time simulation, enabling differentiable sampling for discrete latent variables.
• Empirical analyses demonstrate that NegBio-VAE outperforms Poisson-VAE and Gaussian baselines in reconstruction fidelity and latent space expressiveness for neural spike trains.

NegBio-VAE is a variational autoencoder (VAE) framework designed to model overdispersed count data, with specific emphasis on biological spike trains. It extends the conventional VAE architecture by introducing the negative binomial (NB) distribution for the latent variables, thus overcoming the limitations of Poisson-based modeling in capturing the variance inherent in neural activity. NegBio-VAE introduces explicit dispersion control, new reparameterization schemes, and tractable evidence lower bound (ELBO) computation strategies, leading to improved fidelity in both reconstruction and latent representation of spike-like data (Zhang et al., 7 Aug 2025).

1. Theoretical Motivation and Background

Biological neural spike trains typically exhibit overdispersion, where the variance in spike counts substantially exceeds the mean. Traditional VAEs, when adapted for count data, often default to the Poisson distribution for latent modeling; this approach, as used in Poisson-VAE, rigidly constrains the variance to equal the mean—failing to capture the observed stochasticity and irregularity in empirical data. NegBio-VAE addresses this deficiency by adopting the negative binomial distribution, parameterized by a dispersion parameter $r$ and a success probability $p$, leading to the following moments:

$$\text{mean} = \frac{r(1-p)}{p}, \quad \text{variance} = \frac{r(1-p)}{p^2}$$

This formulation introduces a degree of flexibility that is absent in the Poisson model, allowing explicit and independent control over dispersion and the mean.

2. Variational Formulation and ELBO Optimization

NegBio-VAE optimizes the ELBO using a latent negative binomial model for both the posterior $q(z|x)$ and the prior $p(z)$. Specifically,

$$\text{ELBO} = \mathbb{E}_{z \sim q}[\log p_\theta(x | z)] - D_{KL}\left[q(z | x) \| p(z)\right]$$

with

$$q(z|x) = NB(z; r \cdot \delta_r(x), p \cdot \delta_p(x)), \quad p(z) = NB(z; r, p)$$

Direct calculation of the KL divergence between general negative binomial distributions lacks an analytical solution. To address this, two strategies are developed:

• Monte Carlo (MC) Estimation: The KL divergence is approximated by sampling from the negative binomial variational posterior. This allows generality but introduces gradient noise during training.
• Dispersion Sharing (DS): The encoder is constrained so that the dispersion parameter is not modified ($\delta_r(x) = 1$), yielding a closed-form KL divergence:

$D_{KL} = r \cdot g(p, \delta_p)$

where

$$g(a, b) = \log b + \frac{1 - ab}{ab} \cdot \log\left(\frac{1 - ab}{1 - a}\right)$$

This expression ensures that the KL divergence vanishes at the prior ($\delta_p = 1$), and possesses local quadratic growth in the vicinity of the prior ($\delta_p \simeq 1 + \epsilon$), stabilizing the optimization landscape.

3. Differentiable Reparameterization of Discrete Latent Variables

The NB latent variable is sampled via its compound construction as a Poisson distribution with a Gamma-distributed rate $\lambda$. Since direct reparameterization through discrete distributions is not differentiable, two relaxation schemes allow for end-to-end training:

• Gumbel-Softmax Relaxation: The Poisson is embedded within a truncated categorical, and Gumbel-^^^^0^^^^max sampling is used with a temperature parameter $\tau$. As $\tau \rightarrow 0$, the ^^^^0^^^^ sample converges to a discrete count, making gradients through sampling possible.
• Continuous-Time Simulation: Exploits the equivalence between Poisson processes and exponential inter-arrival times. A "^^^^0^^^^" count is constructed by simulating inter-arrival times and applying a sigmoid activation, again controlled by $\tau$, producing a relaxation that converges to the Poisson in the low-temperature limit.

Both mechanisms enable gradient-based optimization through otherwise non-differentiable sampling steps.

4. Empirical Performance and Analysis

NegBio-VAE demonstrates substantial empirical gains across a range of datasets:

• Image Reconstruction: On datasets including MNIST, Omniglot, downsampled CIFAR (16×16), and SVHN, NegBio-VAE yields higher-fidelity reconstructions than Poisson-VAE and Gaussian-based baselines. Fine-grained details (e.g., gaps in digit "0", stroke endpoints in "4", or textures in natural images) are better preserved.
• Latent Space Expressiveness: The overdispersion index (ODI) remains above unity for NegBio-VAE, indicating persistent variability beyond the Poisson regime. Downstream linear classification and shattering tests on latent codes indicate stronger separability and discriminative power under NB modeling, even in compressed latent spaces.
• Optimization Strategy Effects: Monte Carlo KL estimation introduces higher oscillations and noisier convergence; dispersion sharing with continuous relaxation (DS+C) achieves smoothest and most stable training, yielding the best empirical performance.

A comparison of ELBO strategies and reparameterization approaches is summarized below:

KL Estimation	Relaxation	Training Stability	Reconstruction
Monte Carlo	Gumbel-Softmax	Moderate	Good
Monte Carlo	Continuous-Time	Less Stable	Moderate
Dispersion Share	Gumbel-Softmax	High	Good
Dispersion Share	Continuous-Time	Very High	Best

5. Methodological Innovations and Implications

NegBio-VAE addresses two critical methodological issues in VAE research for count-based or spike data:

1. Decoupled Mean–Variance Modeling: Explicitly controlling overdispersion via NB latent variables and tractable KL divergence estimation creates a more flexible and biologically relevant generative process.
2. Differentiability in Discrete Latents: Compound relaxations for NB sampling are integrated within standard deep learning libraries (using, for instance, Marsaglia–Tsang sampling for Gamma distributions), facilitating practical model deployment.

A plausible implication is that the NegBio-VAE approach can serve as a foundation for simulating and interpreting spike train variability in computational neuroscience, as well as for building robust generative models where count data with high variance is routine.

6. Broader Impact and Domain-Specific Relevance

Beyond spike train modeling, NegBio-VAE's latent structure and regularization enable applications in any domain where discrete, overdispersed counts are intrinsic. This includes genomics (overdispersed gene counts), text modeling (word counts with heavy-tailed variance), and neuromorphic engineering. By decoupling mean and variance, the approach may facilitate the development of hierarchical generative models with adaptable stochastic behavior, potentially expanding the expressiveness and controllability of VAEs in generative modeling and probabilistic reasoning.

7. Comparison with Related Modeling Paradigms

Conventional VAEs with Poisson or Gaussian latents are restricted in their ability to model high-variance or zero-inflated data. NegBio-VAE, through the introduction of an additional dispersion parameter, generalizes Poisson-latent VAEs. The enhanced expressivity is reflected both in empirical performance (reconstruction error and latent structure) and in the theoretical regularization embodied by closed-form DS KL divergence.

Key distinctions are as follows:

Model	Latent Distribution	Mean-Variance Relationship	Overdispersion Control	Biological Plausibility
Poisson-VAE	Poisson	mean = variance	No	Limited
Gaussian-VAE	Gaussian	Free parameterization	Yes, but non-count	Poor
NegBio-VAE	Negative Binomial	variance > mean	Yes	High

In summary, NegBio-VAE represents a principled and tractable advance in discrete latent variable modeling, combining biological relevance with methodologically robust VAE training and inference (Zhang et al., 7 Aug 2025).