Poisson Variational Autoencoder (P-VAE)

• P-VAE is a deep generative model that uses factorized Poisson latent distributions to encode discrete spike counts, reflecting biological realism and sparse coding.
• It employs soft Poisson reparameterization and surrogate gradient methods to enable backpropagation through non-differentiable spike sampling in both standard and spiking neural networks.
• The model integrates a metabolic cost via KL divergence, explicitly trading off reconstruction accuracy and energy consumption to mimic biophysical constraints.

A Poisson Variational Autoencoder (P-VAE) is a deep generative framework in which the latent variables are distributed according to a (factorized) Poisson distribution. Unlike conventional VAEs, which typically assume continuous latent variables with Gaussian priors and posteriors, the P-VAE encodes inputs as discrete spike counts, reflecting both biological plausibility and constraints derived from energy efficiency and sparse coding. The P-VAE objective introduces an emergent metabolic cost, coupling the information-theoretic and energetic aspects of neural representation. The architecture has been implemented in both standard neural and spiking neural network (SNN) domains, providing a bridge between Bayesian inference, predictive coding, and biophysical computation (Vafaii et al., 13 Feb 2026, Vafaii et al., 2024, Zhan et al., 2023).

1. Mathematical Formulation of the Poisson Latent Space

P-VAE models parameterize the generative prior and inference (posterior) distributions over integer-valued, non-negative latent variables as independent Poisson processes. Let $\mathbf{x}\in\mathbb{R}^M$ be the observed data and $\mathbf{z}\in\mathbb{N}^K$ a vector of spike counts. The generative model factorizes as follows:

• Prior:

$$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$

where $\lambda_i > 0$ are learnable baseline firing rates.

• Posterior:

$$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$

where $\delta r_i(\mathbf{x}) > 0$ is a data-driven multiplicative "error" ratio output by the encoder.

• Likelihood:

$$p(\mathbf{x}|\mathbf{z}) = \mathcal{N}\bigl(\mathbf{x}; f_{\theta}(\mathbf{z}), \sigma^2\mathbf{I}\bigr),$$

where $f_{\theta}$ is usually a neural or linear decoder.

The same structure is adopted in SNN implementations, where spike trains over a time window of $T$ steps yield counts $z^i = \sum_{t=1}^T z^{i,t}$, and firing rates $\mathbf{z}\in\mathbb{N}^K$0, so Poisson rates are $\mathbf{z}\in\mathbb{N}^K$1 (Zhan et al., 2023).

2. Poisson Sampling and Reparameterization

Backpropagation through Poisson draws is nontrivial due to their non-differentiable, discrete nature. Two major strategies are deployed:

• Soft Poisson reparameterization (for standard networks): The number of spikes in time $\mathbf{z}\in\mathbb{N}^K$2 for rate $\mathbf{z}\in\mathbb{N}^K$3 is the number of independent exponential($\mathbf{z}\in\mathbb{N}^K$4) arrivals before their sum exceeds $\mathbf{z}\in\mathbb{N}^K$5. This yields a differentiable surrogate: for each $\mathbf{z}\in\mathbb{N}^K$6, sample exponential arrivals $\mathbf{z}\in\mathbb{N}^K$7, accumulate partial sums $\mathbf{z}\in\mathbb{N}^K$8, and approximate $\mathbf{z}\in\mathbb{N}^K$9 as $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$0. This is the Rsample procedure (Vafaii et al., 2024).
• Bernoulli approximation and surrogate gradients (for SNNs): Each spike in $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$1 steps is sampled as $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$2 for $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$3, so the spike count $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$4 for small $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$5. The non-differentiable threshold is replaced during backpropagation by a flat or smooth surrogate, e.g., a small window $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$6, leading to

$$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$7

This allows gradients to flow through spike generation (Zhan et al., 2023).

3. Variational Objective, KL Divergence, and Metabolic Cost

The negative evidence lower bound (ELBO) for P-VAE reads:

$$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$8

where $$p(\mathbf{z}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i),$$9 controls the fidelity–sparsity/metabolic trade-off.

The KL between scalar Poisson distributions is:

$\lambda_i > 0$0

For $\lambda_i > 0$1, setting $\lambda_i > 0$2:

$\lambda_i > 0$3

where $\lambda_i > 0$4.

This KL structure generates a metabolic cost: the penalty for deviating from baseline rates is directly proportional to the total firing rate in the prior. In the regime $\lambda_i > 0$5, the leading order is a simple rate-penalty, i.e., $\lambda_i > 0$6 (Vafaii et al., 13 Feb 2026, Vafaii et al., 2024).

A direct implication is that tuning $\lambda_i > 0$7 in the objective enables a smooth tradeoff between reconstruction fidelity (task performance) and expected energy consumption (spiking rate), which is unique to count-based latent variable models like the P-VAE (Vafaii et al., 13 Feb 2026).

4. Architectures and Algorithmic Implementations

P-VAE admits both conventional deep neural network and SNN implementations:

• SNN-based (e.g., ESVAE (Zhan et al., 2023)):
  • Encoder: Stack of spiking convolutional and leaky integrate-and-fire layers over $\lambda_i > 0$8 steps, outputting binary spike trains $\lambda_i > 0$9.
  • Posterior rates: $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$0; latent Poisson rates $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$1.
  • Prior (generation) branch: $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$2; $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$3; latent rates $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$4.
  • Spike-based sampling: $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$5.
  • Decoder: Stack of spiking deconv layers, collapsing time via rate summation or membrane readout.
  • Loss: Mean squared error for reconstruction plus an MMD penalty between $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$6 and $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$7.
• Feedforward encoder/decoder (e.g., (Vafaii et al., 2024)):
  • Encoder: Feedforward deep net with output nonlinearity ensuring positivity (e.g., $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$8 or $$q(\mathbf{z}|\mathbf{x}) = \prod_{i=1}^K \mathrm{Pois}(z_i; \lambda_i \delta r_i(\mathbf{x})),$$9).
  • Posterior: Factorized Poisson, parameterized by prior baseline times learned function of input.
  • Decoder: Neural network ($\delta r_i(\mathbf{x}) > 0$0) or, for sparse coding, $\delta r_i(\mathbf{x}) > 0$1.
  • Loss: Expected mean squared error (analytic or sampled), metabolic cost term from Poisson KL.
  • Poisson reparameterization as described in Section 2.

Algorithmic details, including forward and backward passes with surrogate gradients for SNN-based variants, are provided in (Zhan et al., 2023).

5. Relationship to Sparse Coding and Biological Relevance

P-VAE implements predictive coding with a discrete, spike-based latent representation. For linear decoders $\delta r_i(\mathbf{x}) > 0$2 with $\delta r_i(\mathbf{x}) > 0$3, the reconstruction loss plus Poisson KL can be written in closed form:

$\delta r_i(\mathbf{x}) > 0$4

where $\delta r_i(\mathbf{x}) > 0$5 (Vafaii et al., 2024). The rate-penalty $\delta r_i(\mathbf{x}) > 0$6 serves as a neural analog of the nonnegative $\delta r_i(\mathbf{x}) > 0$7 penalty in classical sparse coding. In the limit, this recovers spiking sparse coding (Olshausen & Field 1996).

Empirically, basis functions learned by P-VAE on image patches are Gabor-like and tile orientation/phase space, paralleling results from ISTA/LCA sparse coding. Lifetime sparsity achieved is $\delta r_i(\mathbf{x}) > 0$8, similar to classical sparse methods, with marked superiority over Gaussian VAEs, which yield much denser codes (Vafaii et al., 2024).

From a biological perspective, the explicit metabolic cost enforced via the KL reflects energy constraints in biophysical systems. P-VAE makes the tradeoff between representational fidelity and energy expenditure transparent and quantifiable (Vafaii et al., 13 Feb 2026).

6. Empirical Evaluation and Comparison to Alternative VAEs

P-VAE has been benchmarked against Gaussian, Laplace, and Categorical VAEs:

Model	Acc (N=200)	Acc (N=1000)	Acc (N=5000)
P-VAE	0.815±0.002	0.919±0.001	0.946±0.017
Gaussian VAE	0.705±0.002	0.800±0.002	0.853±0.040
Laplace VAE	0.757±0.003	0.869±0.002	0.924±0.028
Categorical	0.673±0.003	0.813±0.002	0.891±0.033

Results on MNIST with latent $\delta r_i(\mathbf{x}) > 0$9 show that P-VAE achieves $$p(\mathbf{x}|\mathbf{z}) = \mathcal{N}\bigl(\mathbf{x}; f_{\theta}(\mathbf{z}), \sigma^2\mathbf{I}\bigr),$$082% accuracy for a downstream KNN classifier with only $$p(\mathbf{x}|\mathbf{z}) = \mathcal{N}\bigl(\mathbf{x}; f_{\theta}(\mathbf{z}), \sigma^2\mathbf{I}\bigr),$$1 labels, compared to $$p(\mathbf{x}|\mathbf{z}) = \mathcal{N}\bigl(\mathbf{x}; f_{\theta}(\mathbf{z}), \sigma^2\mathbf{I}\bigr),$$21000 required for the Gaussian VAE to match this level. P-VAE's latent representations have higher "shattering dimension" and maintain $$p(\mathbf{x}|\mathbf{z}) = \mathcal{N}\bigl(\mathbf{x}; f_{\theta}(\mathbf{z}), \sigma^2\mathbf{I}\bigr),$$390% active latent neurons, avoiding posterior collapse common in continuous VAEs (Vafaii et al., 2024).

Metabolic cost (mean spike count) and sparsity (proportion of zero-activity latents) increase monotonically with rising $$p(\mathbf{x}|\mathbf{z}) = \mathcal{N}\bigl(\mathbf{x}; f_{\theta}(\mathbf{z}), \sigma^2\mathbf{I}\bigr),$$4 in P-VAE, but plateau quickly in rectified Gaussian controls (G-reLU-VAE). This demonstrates that metabolic cost is a specific property of Poisson geometry and not merely a side effect of nonnegativity constraints (Vafaii et al., 13 Feb 2026).

7. Implications for Biophysical and Neuromorphic Computation

P-VAE provides a mathematically principled connection between abstract information-theoretic coding and energy use in neural systems. The KL term proportional to baseline firing rate realizes a resource-aware computation framework where bits and biophysical spikes are directly coupled (Vafaii et al., 13 Feb 2026). The model supports predictive coding and efficient coding principles, enforcing the notion that "silence is free," and enables explicit, interpretable trade-offs between accuracy and energy demands.

Extensions to other count-based priors, such as the negative binomial, yield similar metabolic penalization structures, pointing to a broader universality of such emergent costs in count-based variational inference.

A plausible implication is that Poisson variational inference mechanisms, as instantiated in P-VAE, may serve as a computational blueprint for both biological learning and energy-efficient artificial intelligence (Vafaii et al., 13 Feb 2026, Vafaii et al., 2024).