Negative Binomial Variational Autoencoders for Overdispersed Latent Modeling

2025-08-07

AI Summary

Traditional VAEs and Poisson-VAEs fail to capture neural spike count overdispersion—i.e., variance substantially exceeding the mean. To address this, we propose NegBio-VAE, the first VAE that incorporates the negative binomial distribution in its decoder to explicitly model a learnable dispersion parameter, thereby relaxing the rigid mean-equals-variance constraint inherent to the Poisson distribution. We develop a differentiable reparameterization scheme for the negative binomial distribution and introduce two ELBO optimization strategies to enable end-to-end variational inference. Experiments on multi-source neural electrophysiological datasets demonstrate that introducing only a single learnable dispersion parameter significantly improves spike count reconstruction accuracy. NegBio-VAE consistently outperforms Poisson-VAE and other baseline models on overdispersed neural data modeling tasks, establishing a new state-of-the-art for probabilistic spike train representation learning.

Technology Category

• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Machine Learning: Deep Generative Models & Autoencoders
• Reasoning under Uncertainty: Relational Probabilistic Models

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• User Modeling, Personalization and Recommendation: On-Device user modeling, personalization, and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Biological neurons communicate through spike trains, discrete, irregular bursts of activity that exhibit variability far beyond the modeling capacity of conventional variational autoencoders (VAEs). Recent work, such as the Poisson-VAE, makes a biologically inspired move by modeling spike counts using the Poisson distribution. However, they impose a rigid constraint: equal mean and variance, which fails to reflect the true stochastic nature of neural activity. In this work, we challenge this constraint and introduce NegBio-VAE, a principled extension of the VAE framework that models spike counts using the negative binomial distribution. This shift grants explicit control over dispersion, unlocking a broader and more accurate family of neural representations. We further develop two ELBO optimization schemes and two differentiable reparameterization strategies tailored to the negative binomial setting. By introducing one additional dispersion parameter, NegBio-VAE generalizes the Poisson latent model to a negative binomial formulation. Empirical results demonstrate this minor yet impactful change leads to significant gains in reconstruction fidelity, highlighting the importance of explicitly modeling overdispersion in spike-like activations.

Problem

• Model neural spike trains with overdispersed variability
• Overcome Poisson-VAE's rigid mean-variance constraint
• Introduce negative binomial distribution for accurate neural representations

Innovation

• Negative Binomial VAE for overdispersed modeling
• Two ELBO schemes for NB optimization
• Differentiable reparameterization for NB setting