Poisson Variational Autoencoder
Authors: Hadi Vafaii
Dekel Galor
Jacob L. Yates
Venue: NeurIPS 2024 spotlight (2024)
Keywords: NeuroAI
Bayesian Inference
Predictive Coding
Sparse Coding
Variational Autoencoder
Primary area: generative_models
TL;DR: We introduce the Poisson VAE, a generative model that encodes inputs into discrete spike counts and unifies established theoretical concepts in neuroscience with modern machine learning.

Abstract: Variational autoencoders (VAE) employ Bayesian inference to interpret sensory inputs, mirroring processes that occur in primate vision across both ventral (Higgins et al., 2021) and dorsal (Vafaii et al., 2023) pathways. Despite their success, traditional VAEs rely on continuous latent variables, which significantly deviates from the discrete nature of biological neurons. Here, we developed the Poisson VAE (P-VAE), a novel architecture that combines principles of predictive coding with a VAE that encodes inputs into discrete spike counts. Combining Poisson-distributed latent variables with predictive coding introduces a metabolic cost term in the model loss function, suggesting a relationship with sparse coding which we verify empirically. Additionally, we analyze the geometry of learned representations, contrasting the P-VAE to alternative VAE models. We find that the P-VAE encodes its inputs in relatively higher dimensions, facilitating linear separability of categories in a downstream classification task with a much better (5x) sample efficiency. Our work provides an interpretable computational framework to study brain-like sensory processing and paves the way for a deeper understanding of perception as an inferential process.

Poisson Variational Autoencoder
Hadi Vafaii vafaii@berkeley.co edu
Dekel Galor! galorrberkeley. edu
Jacob L. Yates! yates@berkeley.. edu
UC Berkeley
Abstract
Variational autoencoders (VAEs) employ Bayesian inference to interpret sensory inputs, mirroring processes that occur in primate vision across both ventral [1] and dorsal 2 pathways. Despite their success, traditional VAEs rely on continuous la- tent varaales, which deviates sharply from the discrete nature of biological neurons. Here, we developed the Poisson VAE (P- VAE), a novel architecture that combines principles of predictive coding with VAE that encodes inputs into discrete spike counts. Combining Poisson-distribute distributed latent variables with predictive coding in- troduces metabolic cost term in the model loss function, suggesting relationship with sparse coding which we verify empirically. Additionally, we analyze the geometry of learned representations, contrasting the P-VAE to alternative VAE models. We find that the P-VAE encodes its imputs in relatively higher dimensions, facilitating linear separability of categories in downstream classification task with much better (5x) sample efficiency. Our work provides an interpretable computational framework study brain-like sensory processing and paves the way for a deeper understanding of perception as an inferential process.
Perception as Inference
Rate Coding
Predictive Coding
Amortized Sparse Coding is special case the Poisson VAE
Alfezon, -1000 AD
Adrtnn Zotterman,
Neisar, 1967
Heimholtz, 1850s
1925
Gregery, 1980
Barlow,
Srinkassan at.
Linear Poisson VAE Amortized Sparse Coding
Lee S Mumford, 2003
Ballard, 1999
Yulle 8 Kerdeen,, 2006
Pougett
Barlow;
Friston, 2005
Zohary et 1994
Clark, 2013
Poisson Variational Autoencoder
Linear Gaussian VAE
Probabilistic PCA
Encodes inputs into discrete spike counts, significantly enhancing the model's bio-realism and interpretabillty
A metabolic cost term emerges in the model objective "for free," suggesting connection Sparse Coding.
Brings major theories in neuroscience closer logether. under the unifying umbrella of Bayesian Inference.
Figure Graphical abstract. Introducing the Poisson Variational Autoencoder (P-VAE) VAE), which draws on key concepts in neuroscience. When trained on natural image patches, P- with linear decoder develops Gabor-like feature selectivity, reminiscent of Sparse Coding (3). In sharp contrast, the standard Gaussian VAE learns the principal components (4). Our code, data, and model checkpoints are available at this repository: https://gituub..oo//aaaivaaa
VAE
38th Conference Neural Information Processing Systems (NeurIPS 2024).

---

Introduction
The study of artificial neural networks (ANN) and neuroscience has always been closely linked, driving advancements in both fields [5) Despite the close proximity of the two fields, most ANN
models deviate substantially from biological brains [1] 12]. A major challenge is designing models that not only perform well computationally but also exhibit "brain-like" structure and function. This
seen both goal for improving ANNS [3,15], and better understanding biological brains 8 9
16-19..
which has recently been referred to as the neuroconnectionis research programme 22
Drawing from neuroscience, a major guiding idea is that perception is process of inference 21 where the brain constructs a representation of the external world by inferring the causes of sensory inputs [23 (3). This concept is mirrored in "generative Al" where models learn the generative process underlying their inputs (27 However, in this vein, there is a tension between small well-understood models that directly inspired by cortex, such sparse coding (3) and predictive coding (30. and deep generative models that perform well 31 34..
22
The variational autoencoder (VAE; 35 36) model family is promising candidate for neurocon- nectionist goals for multiple reasons. First, VAEs learn probabilistic generative models of their inputs and are grounded in Bayesian probability theory, providing a solid theoretical foundation that directly incorporates the concept of perceptual inference [10, 22). Second, the VAE model family, specifically hierarchical VAEs, is broad with other generative models, such diffusion models, understood as special cases of hierarchical VAEs (37.39]. Finally, VAEs learn representations that are similar to cortex [1 2 40), exhibit cortex-like topographic organization [41 42 and make perceptual errors that mimic those of humans (43), indicating a significant degree of neural, organizational, and psychophysical alignment with the brain.
However, standard VAEs diverge from brains the way they encode information. Biological neurons fire all-or-none action potentials (44), and are thought to represent information via firing rate 49). These firing rates must be positive and generate discrete "spike" counts, which exhibit conditionally Poisson-like statisties in small counting windows (4). In contrast, VAEs are typically parameterized with real- valued, continuous, Gaussian distributions (52).
Contributions. In this work, we address this discrepancy by introducing the Poisson Variational Autoencoder (P-VAE) a novel architecture that combines perceptual inference with two other inspirations from neuroscience (Fig. [. First, that information is encoded in the rates of discrete spike counts, which are approximately Poisson-distributed on short time intervals. And second, that feedforward connections encode deviations from expectations contained in feedback connections (Fig. 2h: (30 (3)) We introduce a reparameterization trick for Poisson samples (Algorithm [ and derive the evidence lower bound (EI (BBO) objective for the P-VAE (eq. (3). Overall, we believe P-VAE introduces a promising new model at the intersection of computational neuroscience and machine learning that offers several appealing features over existing VAE architectures:
The P-VAE loss derivation (eq. naturally results in a metabolic cost term that penalizes high firing rates, such that VAE with linear decoder implements amortized sparse coding (Fig. (p).
We validate this prediction empirically.
P-VAE largely avoids the prevalent posterior collapse issue, maintaining many more active latents compared to alternative VAE models (Table I. especially the continuous ones.
P-VAE encodes its inputs in relatively higher dimensions, facilitating linear separability of categories in a downstream classification task with a much better (5x) sample efficiency.
We evaluate these results on two natural image datasets and MNIST. The P-V paves the way for the future development of interpretable hierarchical models that perform "brain- liken inference.
Background & Related work
Perception as inference: connections to neuroscience and machine learning. A centuries-old idea 21 22). "perception as inference" argues that coherent perception the world results from the unconscious inference over the causes of the senses. In other words, the brain learns generative model of the sensory inputs. This has led fruitful theoretical work in neuroscience [33 56) and machine learning
[5 58) including VAEs [2] See Marino
10 for review.

---

Efficient, predictive, and sparse coding. Another longstanding idea in neuroscience is that brains are adapted to the statistics of the environment. Efficient coding states that brains represent as much information about the environment as possible while minimizing neural resource use (59 60)
61 62 postulates that the brain generates statistical prediction of its inputs, with feedforward networks carrying only the prediction errors or unexplained information [63] More recently, ANNs based on predictive coding have been shown capture wide range of phenomena in biological neurons across the visual system 66 65] More broadly, prediction in time has emerged
Predictive coding 30
as an objective that lends itself to brain-like representations 66 67
Sparse coding (SC) is directly inspired by efficient coding, aiming to explain inputs as sparsely as possible 44 68 SC was the first unsupervised model to learn representations closely resembling the receptive fields of V1 neurons (3 and predicts an array of empirical features of neural activity (99. - SC is formalized with generative model where neural activations
sampled from sparsity-inducing prior, 2 ~ p(z), and the input image a: is reconstructed as a linear combination of basis vectors 4, plus additive Gaussian noise, ã 4z + E. The SC loss is as follows:
LsparseCoding (æ; 4, 2) = || 22 + + B|1=1|1-
Commonly used algorithms for sparse coding include the locally competitive algorithm (LCA; [80 which is a biologically plausible algorithm to optimize eq. ( and iterative sarikkage-hhresholding algorithm (ISTA; [8] 82). which has shown robust performance in learning sparse codes given fixed dictionary +.
VAE objective. VAEs define probabilistic generative model p(æ, 2). where denotes the observed data and x are some latent variables. The generative process samples z from a prior distribution p(z) and then generates the observed data from the conditional distribution pe (arl z), also known as the "decoder". The "encoder", (0 (z)r), performs approximate inference on the inputs. Model parameters are learned by maximizing the evidence lower bound (ELBO) objective, which is derived from variational inference (see appendix B for the full set of derivations). The EI ELOO is given by:
log ((0) 2 (= (o)))()))))
Dea. (2)(2))))) p[z)) LLVEE Ø)
The first term captures the reconstruction performance of the decoder, and the second term, the "KL term," captures the divergence of the approximate posterior from the prior.
The specific form of these distributions is up to the practitioner. In standard VAEs, factorized 2² (z)) and p (2; 0, 1). The likelihood, pe po =))z) Gaussians are typically used: q q= [2x; p(z)) o
is also typically modeled as Gaussian conditioned on parameterized neural network deco
Amortized inference in VAEs. A major contribution of VAEs is the idea amortizing inference over the latents with black box ANN 83 84). "Amortized" inference borrows term from finance to capture the idea of spreading out costs costs-here, the cost of performing inference over multiple samples. In amortized inference, neural network learns (during training) how to map data sample to a distribution over latent variables given the sample. The cost is paid during training, but the trained model can then be used to perform inference on future samples efficiently. It has been argued that the brain performs amortized inference for computational efficiency (85).
VAEs connection to biology. VAEs have been shown to contain individual latents that resemble neurons, capturing a wide range of the phenomena observed in visual cortical areas (40) and human
perceptual judgments (43) Like many other ANN models [86] 887, VAEs have been found to learn
representations that are predictive of single-neuron activity in both the ventral [T] and dorsal (2)
streams. However, unlike most ANNs, the mapping from certain VAEs to neural activity is incredibly sparse, even one-to-one in some cases 2).
Discrete VAEs. VAEs with discrete latent spaces, such as VQ-VAE 88 and Categorical VAE (89), are designed to capture complex data structures by mapping inputs to a finite set of latent variables. Unlike traditional VAEs that use continuous latent spaces, these models leverage discrete representations to enhance interpretability and can yield high performance with lower capacity (90).
3

---

Algorithm Reparameterized sampling (rsample) for Poisson distribution.
Input:
À E IR BxX n_exp
rate parameter; B, batch size; K, latent dimensionality
20
# number of exponential samples to generate # controls the sharpness of the thresholding
temperature
1: procedure RSAMPLE(A, mt_eeee exp, temperature)
2: Exp <- Exponential
D create exponential distribution sample inter-event times, At: [n_exp x B K] D compute arrival times, same shape as At D soft indicator for events within unit time > event counts, number of spikes, [B K]
3: A/ Exp.rsample (n_exp.))
4: times <- cumsum At, dim=0)
5: indicator f- sigmoid
1-times
temperature
6: - sum(indicatorr, dim=0)
7: return z
8: end procedure
VAEs connection to sparse coding. Previous work has attempted to connect sparse coding and VAEs directly 93). with each approaching the problem differently. Geadah et al. (9T) introduced sparsity-indccing priors (such as Laplace or Cauchy) and a linear decoder with an overcomplete latent space. Tonolini al. [92] introduced spike and slab prior into modified ELBO, and Xiao et al. (93] added sparse coding layer learned by ISTA to the latent space of a VA-. Notably, none of the three ended up minimizing the sparse coding loss. Two of the three maintain the linear generative model with an overcomplete latent space, but the ELBO in both requires an additional approximation step for the KL term [9 92
3 Introducing the Poisson Variational Autoencoder (P-VAE)
Our main contribution is integrating Poisson-distributer latents into VAEs, where both the approximate posterior a
sample
r
Pois(z; (x))))
and the prior are parameterized as Poisson distributions. Critically, the latents are no longer continuous variables, but rather they discrete spike counts. To perform inference over discrete latents, introduce Poisson reparameterization trick. We then derive the term and obtain the full P- P-VAE objective.
Encoder
Decoder
Sr
-C
2 20
Indepencent
(0.1.2....
b
Legend
linear decoder
Encoder
Spaccura Mostel
BA
Poisson reparameterization trick. For homogeneous Poisson process 96). given window size At = 1, and rate X, we can generate Pois- distributed counts by drawing ran- domly distributed wait-times from an exponential distribution with mean 1/A and counting all events where the cumulative time is less than 1. Because the exponential distribution trivially reparameterized (35), and Py Torch contains an implementation (97). need only to approximate the hard threshold for comparing cumula- tive wait times with the window size. We accomplish this by replacing the indicator function with sigmoid in refs.
D
Fro une
Z
2
avercomplete
> M
Detccey deneer 1 D
latent space:
Figure 2: (a) Model architecture. Colored shapes indicate learnable model parameters, including the prior firing rates, r. We color code the model's inference and generative compo- nents using red and blue, respectively. The P-VAE encodes its inputs in discrete spike counts, significantly enhancing its biological realism. (b) "Amortized Sparse Coding" is a special case within the P-VAE model family: it's P-VAE with a linear decoder and an overcomplete latent space.
89 98

---

Algorithm I demonstrates the steps: Given a matrix of rates A, sample n _exp wait times f2, n for each element of À by sampling from exponential distribu- tion with mean 1/x. We then calculate the cumulative se_e
Temperature
T 1.0
Percent
== 0.1
event times ((n_exp) 2==1 tj. pass them through ttemperature ), and sum over samples to get event counts, 2. The temperature controls the sharpness of the thresholding. We adaptively scale the number of samples, n n_exp, by keeping track of the maximum rate in each batch, Amax. and then the inverse cumulative density function (cdf) for Poisson to find the number of samples, nxexp, such that edf(n_expp Amax) 0.99999.
T = 0.01
a sigmoid ( 1-8
T 0.0
3
Figure 3: Relaxed Poisson distribution. Samples are drawn using Algorithm | for À 1. non-zero temperatures. samples are non-integer, but approach the true Poisson distribution as 0.
At non-zero temperatures, our parameterization algorithm provides continuous relaxation of the Poisson distri- bution. Figure 3 shows histograms of samples drawn using Algorithm for rate À = and temperatures
= 1.0, 0.1, 0.01, and O. The latter case (T' 0, true Poisson) equivalent to sarch.poissone ().
P-VAE architecture and residual parameterization. The architecture of P-VAE captures the interactions between feedforward and feedback connections that are present in all visual cortical arcas [99] [00). Feedforward areas carry sensory information and feedback connections are thought to carry modulatory signals such as attention [33 or prediction (30). which interact multiplicatively with feedforward inputs [53 [01].
P-VAE embodies this idea by having the posterior rates depend on the prior, such that Tprior = 7 and post = r 3 (r()), where @ the Hadamard (element-wise) product. The prior rates, r E K are learnable parameters that capture expectations about the statistics of the input. The encoder outputs, 6r(x) E R capture deviations from the prior. Thus, P-VAE models the interaction between prior expectations, and deviations from them, in a multiplicative and symmetric way. This results in posterior, 9(2 (x) = Pois(2; r @ ör(x)), and prior, p(z) Pois(2;;) where is the spike count variable and Pois(2; ; x //2 is the Poisson distribution. Notably, this multiplicative relationship is maximally general, as any pair of positive variables, Pprorr and Tost can expressed as base variable, == Tproors multiplied by their relative ratio, ôr == Tpos./rr See Fig. 2h.
P-VAE loss function. For a comprehensive derivation of the P-VAE objective, see appendix B Here, we report the final result:
LeVaE = EgwPois(s: || dec(2)112 + 1=1
(rrr))
(3)
where dee(-) (-) is the decoder neural network, and f(y) = 1----- y y log (see supplementary Fig. 6).
P-VAE relationship to sparse coding. The KL term in eq. 3 penalizes firing rates. Both and 8r are positive by definition, and fyy 2 0, strongly resembling the sparsity penalty in Olshausen and Field (3. To make this connection more explicit, we make two additional assumptions (Fig. (b):
1. The decoder is linear generative model: a dz with E IR and d E ig MxK
2. The latent space is overcomplete: K > M.
Because both EegpPois(x:)) [zs] and Eg [æ:2,] have closed-form solutions (eq. (22 ). the reconstruction term in eq. (3) can be computed analytically for linear decoder, resulting in:
K
LSC-PVAE 5r 8r r, b) = | + 1 + Adiag( 4 TT) + 3
((rr))
where À = (o) are the posterior firing rates, (1) is defined as above, and /3 is hyperparameter that scales the contribution of the KL term 102). and changes the sparsity penalty for the P-VAE.
5

---

Table 1: Models considered in this paper.
Discrete
Continuous
Poisson VAE
Categorical VAE (C-VAF; 89 98)
Gaussian VAE
Laplace VAE (C-VAF; [40] 91)
(P-VAE)
(G-VAE; 35 36)
The relationship between the linear P-VAE loss (eq. 4 and the sparse coding loss (eq. (II can now be seen. Both contain a term that minimizes the squared error of the reconstruction and a term (two terms for P-VAE) that penalizes non-zero firing rates. Unlike prior work that directly implemented amortized sparse coding [91] 92]. here the activity penalty naturally emerges from the derivations, overcomplete linear generative model. The inference is
and the only additional assumption accomplished using parameterized feed-forward neural network, ör(æ), thus, it is amortized [83] We call this specific case of P-VAE "Amortized Sparse Coding" (Fig. (p).
Note that a closed-form derivation of the reconstruction term possible for any VAE with linear decoder and generating distribution that has mean and variance (see eq. (2
This closed- form expression of the loss given linear decoder is useful because we can see how different parameters contribute to the loss. Furthermore, we can compute gradients of the whole loss exactly, and use this to evaluate our Poisson reparameterization.
Experiments
To evaluate the P-VAE, we perform three sets of experiments. First, we utilize the theoretical results for a linear decoder (eqs. (4) and (2T) to test the effectiveness of our reparameterization algorithm. We compare to alternative VAE models with established reparameterization tricks (e.g., Gaussian).
Second, to confirm p- P-VAE with linear decoder not only resembles amortized sparse coding but practically performs like sparse coding, we compare to standard and well-established sparse coding algorithms such as the locally competitive algorithm (LCA: (80)) and the widely-used iterative trinkagg-thresholldee algorithm (ISTA; 81 82 to see if P- VAE reproduces their results.
Third, we test the P-VAE in generic representation learning context and evaluate the geometry of learned representations for downstream tasks. For these experiments, both the encoder and decoder's architecture ResNet (see appendix C for full architecture and training details).
Architecture notation. We experimented with both convolutional and linear architectures. We highlight the encoder and decoder networks using red and blue, respectively. We use the (encldec) convention to clearly specify which architecture type was used. For example, (conv(lin) represents model with convolutional encoder and linear decoder. Using this notation, note that [lin|]img and (conv|lin) architectures were used for the first and second sets experiments, while (conv]conv) architectures employed for the third.
Alternative models. We compare P-VAE to both discrete and continuous VAEs (Table I]. Other than the traditional Gaussian, we compare to Laplace-distributed VAEs because previous work found the Laplace distribution supported robust sparse representations [40] 91). Additionally, we
compare to Categorical VAEs, trained using the Gumbel-Softne trick [89 98 implementation which is based on Maddison et al. (98).
We use Py PyTorch's
Finally, we test models where Gaussian latents are passed through an activation function before passing to the decoder. We call these models G- +att where act {relu, exp). capturing other families of distributions (truncated Gaussian and log-normal), We include these to test the hypothesis that positive constraints (and not discrete latents) are the key contribution of Poisson [103]
Datasets. For sparse coding results, we use 101 natural images from the van Hateren dataset (104). We tile the images to extract x 16 patches and apply whitening and contrast normalization, as
typically done in sparse coding literature (3, 105 To test the generalizability of our sparse coding results, we repeat these steps on CIFAR 10 [066,, dataset we call CIFAR16x11 For the general
representation learning results, use MNIST. See appendix C for additional details.

---

Table 2: Reparameterized gradient estimators perform comparably exact ones across datasets and encoder architectures (linear vs. convolutional). Exact gradients are only computable for linear decoders (see eqs. (IT). (23 and (24). Values represent percent drop in validation loss (lower is better), shown as mean1.0% confidence interval calculated from n 5 random initializations. The best-performing case was selected as the single best random seed for models of the same architecture and dataset across gradient methods (1 out of: 15 for P-VAE, 10 for G-VAE), See supplementary Fig. 7 for a visualization of the same data presented in this table. For actual loss EX: exact; MC: Monte Carlo; ST: straight-through [107].
values, see supplementary Table 3
van Hateren
CIFARI6x 16
MNIST
Model
[lin[lin)
(conv|
(1in|lin)
(conv|1in)
[lin[lin)
(conv|lin)
EX
0.62.5
0.1
0.0+.1
0.0
0.1
0.5
P-VAE
MC
0.0
0.7#.1
0.2+.0
0.5+.1
0.7+.4
0.9
ST
7.3#.1
10.5+.1
9.1+.1
12.5+.1
8.1+.3
11.8#.2
EX
0.1
0.0
0.0
0.04.0
0.1
0.1
G-VAE
MC
0.1
0.02.0
0.1
0.0
0.4
0.3
Statistical tests. In the VAE literature, it is known that random seeds can have large effect compared to architecture or regularization [108]. Therefore, we train each configuration using different random initializations. We report 99% confidence intervals throughout, and perform paired (-tests, reporting significance for p < 0.01 (FDR corrected using the Henjamini-Hoohheee method).
Evaluating the Poisson reparameterization algorithm. P-VAE with a linear decoder has a closed form solution (eq- (4). which lets us evaluate how well our reparameterized gradients perform compared to the exact ones. We compare our results to the gold-standard Gaussian (Table 2 as well as Categorical and Laplace VAEs (supplementary Table (5). In Table 2 report the percent performance drop relative to the best fit, enabling meaningful comparisons across architectures and datasets. Monte Carlo sampling with Poisson reparameterization closely matches exact inference just like established methods for Gaussian and Laplace. In contrast, the straight-through (ST; [107) estimator performs poorly (Table
see also supplementary Fig. 7
Annealing the temperature. The temperature parameter (T) is a crucial hyperparameter in our Poisson reparameterization trick (Algorithm [T. To assess its impact, we followed standard practice (89) and annealed '' during the first half of training, starting from large value (Tatar 1) and gradually decreasing it to a small value (Zinall = 0.05 in the main paper). Figure 9 shows the performance on the van Hateren dataset as function of various Tinal, two architectures ([lin]lin [lin) and (conv|lin)) as well as two annealing schedules (linear vs. exponential; see inset). We find that final temperatures Timal S 0.1 and either annealing strategy work well.
During training, we maintain T > O, which results in continuous (floating) latent variables, x. At test time, we set T 0 produce genuine integer Poisson samples. Crucially, all reported results use T = 0 at test time. We also explored a fordard scheme during the latter half of training, where T remains nonzero only in the backward pass. This surrogate gradients approach provides integer latents in the forward pass but, somewhat unexpectedly, underperformed our "relaxed Poisson" method (Fig. (9) These findings suggest that surrogate gradient methods might benefit from relaxing the hard-forward strategy during training. We believe this observation will be of particular interest to the spiking neural network community, which often relies on surrogate gradients for training.
The P-VAE learns basis vectors similar to those from sparse coding. A major result from sparse coding is that it learns basis vectors (dictionaries) that resemble the "Gabor-like" receptive fields of cortical neurons [3] 109 [10]. Inspecting the dictionaries learned by different models demonstrates
this is not trivial (Fig, 4. As expected from theoretical results 4 g- VAE (top left) lear probabilistic PCA, but with many noisy elements. As demonstrated previously 44 91). L- (lower left) learn Gabor-like elements. However, there are a large number of noisy basis vectors. It is of note that previous work did not show complete dictionaries for their results with Laplace latents (40) 9T]. In contrast, P-VAE (top middle) learns Gabor-like filters that cover space, orientation, and spatial
7

---

Gaussian VAE (# dead neurons: 4011
Poisson VAE (A deed neurons: 5)
Locally competitive algorithm (LCA)
Laplace VAE dead 416)
Categorical VAE (6 dladd neurons:
iterative shrinkage thresholding algoritm (STA)
Figure 4: Learned basis elements for various (1in|lin) VAEs (first two columns) and standard sparse coding models (last column). There are total of K = 512 elements, each made of 16 x 16 256 pixels (i.c., 2 E 255 Features are ordered from top-left to bottom- right, in ascending order of their associated KL divergence (P-VAE, G-VAE, L- C-V,,), the magnitude of posterior logits (C-VAE). The sparse coding results (LCA and ISTA) are ordered randomly.
frequency. The quality is comparable to sparse coding dictionaries learned with LCA/ISTA (top/lower right panels). C-VAE also learns Gabors, although there are significantly more noisy basis elements.
The P-VAE avoids posterior collapse. A striking feature of Fig. Dis the sheer number of noisy basis vectors for both continuous VAEs (G-VAE, C-VAE). We suspected this reflected dead neurons collapsed latent dimension that's no longer encoding
with vanishing KL, which is indicative information. To quantify this, we binned the distribution of KL values and thresholded the resulting distribution at discontinuous points (see supplemental Fig. TO Table shows the results of this analysis for all VAFs with valid KL terms. Across all datasets, both continuous VAEs suffered from large numbers of dead neurons, whereas P-VAE largely avoided this problem. On both natural image datasets, P-VAI VAE had 2% dead neurons compared to ~ 80% for G-VAE and C-VAE. Having more expressive encoder slightly increases this percentage, but a dramatic diffcrence between P-VAE and continuous VAES (G-VAE, C-VAE) persists.
The P-VAE learns sparse representations. To quantify whether P-VAE learns sparse representa- tions, we compared our VAE models to sparse coding trained with LCA and ISTA and quantified the lifetime sparsity [69]. The lifetime sparsity of the j-th latent is:
8j (
Σ, 24y Σ
(5)
N
where N is the number of images, and ziy is sampled from the posterior for the i-th image. Intuitively, 8y = whenever neuron j responds to a single stimulus out of the entire set (highly selective). In contrast, 8j O whenever the neuron responds equally well to all stimuli indiseriminately.
Fig. 5p shows the reconstruction performance (MSE) compared to lifetime sparsity (s. eq. (5)) for all VAES. Empty and solid circles represent (conv|lin) and (lin architectures, respectively. The G-VAE finds good reconstructions (MSE = 71.49) but with low sparsity (s = 0.37). Because the explored different 3 values for P-VAE with both (lim|lin) and (conv]lin) architectures (Fig. 54. blue curves). This maps out rate-distortion curves, enabling to compare the sparsity levels at which P-VAE matches G-VAE performance.
P-VAE KL term explicitly penalizes rate (eq. (3)
With a simpler (linear) encoder, (lin|1in) P-VAE matches (conv) [in) G-VAE performance while achieving 1.7x greater sparsity at 3 = 0.6. A (conv|lin) P-VAE further increases this gap to 2.4x greater sparsity. Adding relu activation to G-VAE also increases sparsity (s = 0.69). By comparing [lin|] and (conv) (in) P-VAE models, we observe that enhancing encoder complexity for the same 3 (gray arrows) preserves MSE performance while achieving greater sparsity. This
highlights how amortization quality can significantly influence rate-distortion curves [33 111-1133

---

Table 3: Proportion of active neurons. All models considered in this table had latent dimensionality of K = -112, with either [lin[lin) or (conv| [lin) architectures. See also supplementary Fig. I0
van Hateren
CIFAR16x0
MNIST
Model
linear
conv
linear
conv
linear
conv
P-VAE
0.984 0.819#.011
0.999 002 0.928
0.537
0.426#t.01
C-
0.188=.000
0.22221.00
0.1932.003
0.23011000
0.027=.000
0.034.1.002
G-VAE
0.218
0.2461.000
0.10511.008
0.24611000
0.027
0.0311..00
a 1.0
b 1.00
0.95
Lifelime sparsity
(cout)]-in)
cateponcal
Lifelime sparsity
0.90
0.8-
LCA
goussianvee
0.05
ISTA
0.7-
peupaken-ecs
E- 0.0L
0.80
5o0n
besn W LCA inheance (u 0.033)
J- 0.20
toon w LCA intrennee (ua 5.000
J- 0.40
0.75
po w LGA inference dua 6.01)
E- 0.50
500s y LCA infrrenee (hra 0.055
0.5-
(asallio)
0.50
0.70
soon LCA infornnce (Ac 0.11
butn LCA infennce (Aca
0.4-
0.65
tean LCA infernnce (Ace
LCA infernnce (ua
0.3
0.60
0
50 100
150
200
250
60
100
120
Reconstruction performance (MSE)
Reconstruction performance (MSE)
Figure 5: Reconstruction performance vs. sparsity of representations. (a) Results for the VAE model family. The curves are sigmoid fit to (lin 1in) and (conv/lin) P- P-VAE results across varying B values (3 from eq. (4). Empty circles correspond for P-VAE (blue open circle) compared to sparse coding (LCA/ISTA), Solid points show results from applying the L.C inference algorithm to P-VAE basis vectors at different sparsity levels (9LCA from eq. D.. The purple curve is sigmoid fit, and curves from part (a) are also included for comparison.
(conv|lin) architectures. (b) Amortization gap
Does P-VAE match the performance of traditional sparse coding trained with LCA or ISTA? Figure compares P- P-V sparse coding models that were trained using wide range of hyperparameters, and the best models were selected for each class (appendix C. P- VAV achieves a similar sparsity to LCA and ISTA (s = 0.94, 0.91, and 0.96, respectively), but the best LCA model drastically outperforms P-VAE on MSE for similar levels of sparsity. This suggests our convolutional encoder is struggling to close the amortization gap. To test this hypothesis, we performed LCA inference on basis elements learned by V-V (Fig. 55 curve/solid points). We explored range of hyperparameters to determine whether the MSF improved for similar sparsity levels. Indeed, I.CA inference using P- P-VAE dictionary was able to nearly match the performance of sparse coding LCA for similar levels of sparsity. This confirms our hypothesis that large amortization gap remains for the specific encoder architectures we tested, highlighting the need for improved inference algorithms/architcctree. 112)
The P-VAE more sample efficient in downstream tasks. To assess downstream performance, we trained (conv) conv) VAE models with K = 10 latent dimension on MNIST (see supplementary Fig. 12 for generated samples and reconstructions from these models). We then extracted represen- tations from the trained encoders and evaluated their ability to classify MNIST digits. We define representations as mean vectors ya for continuous VAEs (G- (A-, L- C-VAE) following conventions in the VAE literature [088] and use log dr for P-VAF, and logits for C-VAE.
We split the MNIST validation set into two 5,000 sample sets, used as train/test sets for this task. We train K-nearest neighbors (KNN) classifiers with a varying number of limited supervised samples (N 2000, 1000, 5000) drawn without replacement from the first (train), to measure classification accuracy on the withheld set (test). KNN is nonparametric, and its performance is directly influenced by the geometry of representations by explicitly capturing the distance between encoded samples [14). We find that using only N 200 samples, P-VAE achieves ~ 82% accuracy in held out data; whereas, G- VAE achieves the same level of accuracy at N = 1000 samples (Table By this measure, P-VAE is 5x more sample efficient. But from Alleman et al. [115), we know that the choice of activation function changes the geometry of learned representations. Therefore, we also tested G-VAE models with an activation function (relu and exp) applied to latents after sampling from the

---

Table 4: Geometry of representations (K 10 only; see Table 6 for the full set of results).
Latent
KNN classification (N, labeled samples)
Shattering dim.
Model
dim.
N 200
= 1,000
N = 5,000
P-VAE
0.815#.002
0.919
0.946+.017
0.797 4.09
C-VAE
0.705=.002
0.800#+.002
0.853+.040
0.795#.006
L-VAE
0.7571.003
0.8691.002
0.9241.028
0.751 1:008
= 10
G-VAE
0.6731.003
0.81311002
0.8911.033
0.758.1.007
G-VAE +relu
0.69411.03
0.8171.00
0.8771.045
0.762.1.07
G-VAE +oxp
0.6422.003
0.784+.002
0.863+.085 032
0.737+.008
posterior. This biological constraint improved G-VAE, but it still underperformed P-VAE (Table We also found this result held for higher dimensional latent spaces (supplementary Table
In supplementary analyses (Fig. [1], we evaluated the representations using logistic regression trained on the full dataset. For larger latent dimensionalities (K 50, 100), P-VAE outperformed all other VAEs, but at lower dimensionalities (K 10), underperforms both G-VAE and C-VAF.
The P- P-VAE learns representations with higher dimensional geometry. The preceding results are indicative of substantial differences in the geometry of the representations learned by P--AE compared to other VAE families (Table 4 To test this more explicitly, we calculated the "shattering dimensionality" of the latent space 118). Shattering dim measures the average accuracy over all possible pairwise classification tasks. This called "shattering" because if the model shatters data points around into high dimensional space, they will become more linearly separable. For MNIST with 10 classes, there are (15) 252 possible classifications. We trained logistic regression on the entire training set to classify each of the 252 arbitrary splits and measured the average performance on the entire validation set. The far right column of Table 4howws the measured shattering dims. For K 10, the shattering dim was significantly higher for discrete VAEs (P-VAE, C-VAE). For higher dimensional latent spaces P- V-VAE strongly outperformed alternative models (Ta
5 Conclusions
In this paper, we introduced the P-VAE, generative model that encodes inputs into discrete spike counts and unifies established theoretical concepts in neuroscience with modern machine learning. We introduced Poisson reparameterization algorithm and derived the ELBO for Poisson-distribute latent variables. The P-VAE objective results in term that penalizes firing rates, like sparse coding. We showed that P- P-VAE with a linear decoder reduces to amortized sparse coding. We evaluated the representations on downstream classification tasks and found that P. VAE encodes its inputs in a higher dimensional space, enabling good linear separability between classes
Limitations. P-VAE samples Poisson latents, Although this is inspired by the statistics of spike counts in the brain over short time intervals (50), there are deviations from Poisson throughout the cortex over longer time windows [51). We discuss this point in appendix A A second limitation is the amortization gap between our current implementation of P. VAE and traditional sparse coding. This could likely be closed with more expressive encoders 119 or through iterative inference [113 120). but is open area of research [112]
Neuroscience implications and future directions. Like biological neurons, the P- P-V generates spikes. This non-negative, discrete representational form closely parallels neuronal spiking activity. Therefore, the P- P-VAE can be more directly compared to neuronal circuits than unconstrained, continuous VAEs. This analogy facilitates in silico perturbation experiments (e.g.. "stimulating" or "silencing" PP-VAE neurons) to mirror in vivo causal manipulations. It also allows applying methods like Most Exciting Inputs (MEI; [21]), ), which assume non-negative activations. Future work could explore hierarchical P-VAEs, finding sweet spot between interpretability and performance. Overall, the biologically inspired representational form of VAV brings computational modeling closer to experimental neuroscience and opens new avenues for advancing NeuroAI research
[3 20
10

---

Code and data
Our code, data, and model checkpoints are available here: http://gtthbbbcommmaaaaaa////ooo
7 Acknowledgments
This work was supported by the National Institute of Health under award number NEI EY032179, Additionally, this material is based upon work supported by the National Science Foundation Grad- uate Research Fellowship Program under Grant No. DGE-175281 (DG). Any opinions, findings, conclusions, or recommendations expressed this material are those the author(s) and do not necessarily reflect the views of the National Science Foundation. We thank anonymous reviewers for their helpful comments, and the developers of the software packages used in this project, including PyTorch [97], NumPy [122]. SciPy [123] scikit-learn
[24]. pandas [25], matplotlib [126), and
seaborn 127
References
[1] Irina Higgins et "Unsupervised deep learning identifies semantic disentanglement in single inferotemporal face patch neurons". In: Nature Communications 12.1 (2021), p. 6456. DO1:
10. 1038/ 10.1000//s444 s41467-4444777 02100021-22222 267-511--
[2] Hadi Vafaii et al. "Hierarchical VAEs provide a normative account of motion processing in the primate brain".. In: Thirt--seventh Conference Neural Information Processing Systems.
2023. URL: attpss:/opeeneviieee wttp://oooo/iidddoototttibteeee
[3] Bruno A Olshausen and David Field. "Emergence of simple-cell receptive field properties by learning sparse code for natural images". In: Nature 3.1.6583 (1996), pp. 607-609, DOI:
10. 1038/38160788
[4] Michael E. Tipping and Christopher M. Bishop. "Probabilistic Principal Component Analy- sis". In: Journal of the Royal Statistical Society Series B: Statistical Methodology 61.3 (Jan. 1999). pp. 611-622. ISSN: 1369-7412. DOT: 10.1/10..111////....1.......1111
00196
[5] Warren McCulloch and Walter Pitts. "A logical calculus of the ideas immanent in nervous activity". In: The bulletin of mathematicul biophysics 5 (1943), pp. 115-133. DOI: 10. 1007/ BF02478259
[6] Patricia S Churchland and Terrence J Sejnowski. "Perspectives on cognitive neuroscience". In: Science 242.4879 (1988), pp. 741-745 DOr: 10. 1111/11ciennce 3055294
[7] Michael SC Thomas and James L McClelland. "Connectionist models of cognition", In:
The Cambridge handbook of computational psychology (2008), pp. 23-58. URL: http
//www7 bbk. uk psychology/ / dnl/wps content /pploads 2023 2002/10/// T0Tomaas
McClelland-proom pdf
[8] Nikolaus Kriegeskorte. "Deep neural networks: new framework for modeling biological vision and brain information processing". In: Annual Review of Vision Science (2015),
pp. 417-446. DOL: 10. 1101/029876
Grace W. Lindsay. "Convolutional Neural Networks Model of the Visual System: Past, Present, and Future". In: Journal of Cogmitive Neuroscience 33.10 (Sept. 2021), pp. 2017-
2031. ISSN: 0898-929X. DOI: 10. 1ew./12//ja http:/jon_a__15550
[10] Joseph Marino. "Predictive coding, variational autoencoders, and biological connections". In:
Neural Computation 34.1 (2022), pp. 1-44. DOI: 10. 1 www/6///cc... ________0_00 a__1458
[11] Jeffrey S Bowers et al. "Deep problems with neural network models of human vision". In: Behavioral and Brain Sciences 46 (2023), e385. DOT: 1001101/001445552220222
[12] Felix A. Wichmann and Robert Geirhos. "Are Deep Neural Networks Adequate Behavioral Models of Human Visual Perception In: Annual Review of Vision Science 9. Volume
2023 (2023), pp. 524.524. ISSN: 2374-4650, DOI: 10. 1 mnnrev-vvsiion-1220
031739
[13] Anthony Zador et al. "Catalyzing next-generation Artificial Intelligence through NeuroAI". In: Nature Communications 14.1 (2023), p. 1597. DOI: 10. 1038/s41467 -41467-023-37180
[14] Fabian Sinz et "Engineering less artificial intelligence", In: Neuron 103.6 (2019). pp. 967-979, DOI: 10. 1016/j neuron. 2019. 08.

---

[15] Demis Hassabis et al. "Neuroscience-inspired artificial intelligence". In: Neuron 95.2 (2017),
pp. 245-258. DOI: 10. 1016/j neuron. 2017.06 .6.001
[16] Nancy Kanwisher et al. "Using artificial neural networks to ask "why' questions of minds
and brains". In: Trends in Neurosciences (2023). DO1: 10. 1016/ tins 2022. 12. 008
[17] Blake Richards al. "The application of artificial intelligence to biology and neuroscience".
In: Cell 185.15 15 (2022), pp. 2640-2643. DO1: 10. 1et.10166...1111 2022. 06 06.047
[18] Blake A Richards et al. "A deep learning framework for neuroscience", In: Nature Neuro- science 22.11 (2019), pp. 1761-1770. DOI: 10. 1038 1111///111111000000-0-00000
[19] David GT Barrett et al. "Analyzing biological and artificial neural networks: challenges with opportunities for synergy? In: Current Opinion in Neurobiology 55 (2019), pp. 55-64. DOI:
10. 10.6/. conb 2019.01..00 01.007
[20] Adrien Doerig et al. "The neuroconnectionist research programme". In: Nature Reviews Neuroscience (2023), pp. 1-20. DO1: 10. 0ttp:///44458000200700000
[21]
Ibn al-Haytham. Book of optics (Kitab Al-Manazir), 1011-1021 AD.
[22]
Hermann Von Helmholtz. Handbuch der physiologischer Optik. Vol. 9. Voss, 1867.
[23]
Tai Sing Lee and David Mumford. "Hierarchical Bayesian inference in the visual cortex". In: JOSA 20.7 (2003), pp. 1444.1448. DOI: 10. 1334//JSSAAA 20. 001434
[24]
Bruno A. Olshausen. "Perception as an Inference Problem".. In: The Cognitive Neurosciences (Sth edition) (2014). Ed. by Michael Gazzaniga and George R. Mangun. DOL: 10. 7551/
mitpress/ 9504 003 0037
URL: http::/rctai org/ woruo/ppppet perceptionn
as-inference. pdf
[25]
Edwin Garrigues Boring. "Perception of objects". In: American Journal of Physics (1946). D01: 10. 11w.1119/1..11 1990807
[26]
Karl Friston. "The free-energy principle: unified brain theory?" In: Nature Reviews Neuro- science 11 11.2 (2010), pp. 127-138. DO1: 10. 1008/0003///// nrn2777
[27] Sam Bond-Taylor et al. "Deep generative modelling: A comparative review of vaes, gans, normalizing flows, energy-based and autoregressive models". In: IEEE transactions pattern analysis and machine intelligence 44.11 (2021), pp. 7327-7347. DOI: 10. 1109/
TPAMI. 200.2021..3000 3116661
[28]
Stanley H. Chan. Tutorial on Diffussio Models for Imaging and Vision. 2024. arviv: 2403
18103 [G]]
[29]
Wayne Xin Zhao et al. A Survey of Large Language Models. 2023. arXiv: 2303 18223
[cs. CL]
[30]
Rajesh PN Rao and Dana H Ballard. "Predictive coding in the visual cortex: functional interpretation of some extra-classical receptive-fiele effects". In: Nature Neuroscience 2.1
(1999), pp. 79-87. DOI: 10. 1038/ /4880
[31] Robin Rombach et al.
Image Synthesis With Latent Diffusion Models". In: Proceedings of the JEEE/CVF Conference on Computer Vision and Pattern Recogni- tion (CVPR). June 2022, pp. 10684-10695 URL: /
thecvf com/ content/CVPR200 / html (tatttmm//ommaae htt////omaaaahhhecceeee Htwebahhhhhggg-Resosstiiiioo Resolutionn Iattaanaaaaeeeaiiiiiiaaaaa With Latent_Diffus LatiffuusionMoodellsse Modeloodool..ooo Seceeesssc_P__. httpeels__c____2______222 WWtCCC___222__eeeeeeee html
[32] Tero Karras et al. "A Style- Based Generator Architecture for Generative Adversarial Net- works". In: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition (CVPR). June 2019. URL: https: http:::///pp/acccseeeeee thecvf com/ content_cvvii CVPR_ 2019/html/Kee Karrasrra__a.. Style-yyleee Based Generaterrrrrrre Beeeeeeeatooddeeeetttrrrrrrr ArrceerrattrrArcctttttuureee Wwwhitecture for wow.ffrreeeriieee..mmm
Adversarial_Nettorre CVPR httpiPp__o11_ppppp.__ee html
[33] Arash Vahdat and Jan Kautz. "NVAE: A Deep Hierarchical Variational Autoencoder". In: Advances in Neural Information Processing Systems. Vol. 33. Curran Associates, Inc., 2020. pp. 19667719679, URL: https: htpp::tsss//pppr...ntit/s wip.. wttp:ecccppppppeaeeeeeee http://pppp..leeeeep//////////aa 20000/ash0/ ht1p:666655577772222/6bb111.3333
html
[34]
Rewon Child. "Very Deep VAEJs Generalize Autoregressive Models and Can Outperform
Them on Images". In: International Conference on Learning Representations. 2021. URL:
https:// /ttps:///peerreiex net/ http://wwneet/oomriiaaaaaaa. RtR./rrrmm----RLXXXVEE
[35]
Diederik Kingma and Max Welling. " "Auto-encoding variational bayes". In: (2014). arXiv:
1312. 6114v11 [stat. ML]
12

---

[36] Danilo Jimenez Rezende et al. "Stochastic backpropagation and approximate inference in deep generative models". In: Internatiional Conference on Machine Learning. PMLR. 2014, pp. 1278-1286 1286. URT.: https: wtte://poeeeddigs.ooo mlr. press/ /v2sss///2//e htt////22/eeeeeee rezende14 html
[37] Diederik P Kingma and Ruiqi Gao. "Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation" In: Thirty-seventhh Conference on Neural Information Processing Systems. 2023. URL: https: https: httpp:://ppeereeee.... http:/ewweeeeeeeeeeeeeeeeeeeeeee htte//oommmmiii id N
[38] Karsten Kreis et al. NeurIPS 2023 Tutorial on Latent Diffiusion Models. https
heurips2023-1dmm tetnn-uttorall... wtt.oaa...ggguu.bbbb....... i0/ 2023.
[39] Diederik Kingma et al. "Variational Diffusion Models". In: Advances in Neural Informa- tion Processing Systems. Ed. by M. Ranzato et al. Vol. 34. Curran Associates, Inc., 2021, pp. 21196-217777 URL: https: /ttp::::///ppocceieeeeee orr.aa/222.....1ee222200060ae..2 wbstrrttttaaattttttttt.....a.... html
neurips cc paper /021 hash
[40]
Ferenc Csikor et al. "Top-down perceptual inference shaping the activity of early visual cortex". In: bioRxiv (2023). DOI: 100.0111012222333 29. 2696622
[41] T. Anderson Keller et al. "Modeling Category-Seleeeee Selective Cortical Regions with Topographic Variational Autoencoders". In: SVRHM 202/ Workshop @ NeurlPS. 2021. URL: https
[42] T. Anderson Keller and Max Welling. "Topographic VAFs learn Equivariant Capsules In: Advances in Neural Information Processing Systems. Ed. by M. Ranzato al. Vol. 34. Curran Associates, Inc., 2021, 1, pp.
URL:
neurips cc/
[43]
Katherine R Storts et al. "Unsupervised leamning predicts human perception and misperception of gloss". In: Nature Human Behaviour 5.10 (2021), pp. 1402-1417, DoI: 10. 10../0013//41152
021-01
[44]
Edgar Adrian. The activity of the nerve medicine/1993
1932.
[45]
Edgar Douglas Adrian and Yngve Zotterman. "The impulses produced by sensory nerve- endings: Part II. The response of Single End-Organ". In: The Journal of Physiology (1926), pp. 151-71. DO1: 11.10.11110 1111/ jphyssopppiol..... 19926..p0022..
[46]
Donald Perkel and Theodore H Bullock, "Neural coding". In: Neurosciences Research
Program Bulletin (1968). URL: http::
nasa. wwwaaagg.aaa...... cwttttiion//9990//222200
[47] Horace B Barlow. "Single units and sensation: a neuron doctrine for perceptual psychology?"
In: Perception 1.4 (1972), pp. 371-394. DOr: 10. 1068//0103711
[48]
Ehud Zohary et al. "Comelated neuronal discharge rate and its implications for psychophysical performance". In: Nature 370.6485 (1994), pp. 140-143. DOI: 10. 1038/370140ad
[49] Fred Rieke et al. Spikes: exploring the neural code. MIT press, 1999.
URL: https : /
mi tpress mit. eeait..du/9770 wwt.e///77800211177777111/111/11
[50] Malvin C Teich. "Fractal character of the auditory neural spike train". In: IEEE Transactions on Biomedical Engineering 36.1 (1989), pp. 150-160. DOI: 10. 1109/10 16460
[51] Robbe LT Goris et al. "Partitioning neuronal variability". In: Nature Neuroscience 17.6 (2014), pp. 858-865. DOI: 10. 1038/nn nn.3111
[52] Diederik P Kingma and Max Welling. "An introduction to variational autoencoders". In: Foundations and Trends@ in Machine Learning 12.4 (2019), pp. 307-392. DOI: 10. 1561/
2200000056]
[53] Charles D Gilbert and Wu Li. "Top-down influences on visual processing". In: Nature Reviews Neuroscience 14.5 (2013), pp. 350-363. DO1: 10. 10w./////////666...
[54] David C Knill and Alexandre Pouget. "The Bayesian brain: the role of uncertainty in neural coding and computation" In: Trends in Neurosciences 27.12 (2004), pp. 712-719. DOI: 10. 1016/j. tins...ii..... 2004 10.
[55] Horace Barlow. "Redundancy reduction revisited". In: Network: computation in neural systems 12.3 (2001), p. 241. DO1: 10. 1080/net. 12.3. 241.253
[56] Karl Friston. "The free-energy principle: a rough guide the brain?" In: Trends in Cognitive
Sciences 13.7 (2009), pp. 293-301. DO1: 10. 1016/j.e.... ticss....is....... 2009. 04. .05
[57] Peter Dayan al. "The Helmholtz machine". In: Neural Computation 7.5 (1995), pp. 889- 904. DOI: 10. www1111ennccccc 1999.97.5.0000
13

---

[58] William Lotter et al. "Deep Predictive Coding Networks for Video Prediction and Unsu-
pervised Learning.. In: International Conference on Learning Representations. 2017. URL:
httpss://pperreviex htt.pppereeieeeeeeteorruunnaaaaa
[59]
Fred Attneave. "Some informational aspects of visual perception." In: Psychological review 61.3 (1954), p. 183. DOI: 10. 1037 1037/00054663
[60]
Horace B. Barlow. "Possible principles underlying the transformation of sensory messages" In: Sensory communication 1.01 (1961), 217 URL: httph://ppss cnbc. cmu, edu/ http://ww..pppppp....aaoooaaaaae
1961 pdf
[61]
Mandyam Veerambudi Srinivasan et al. "Predictive coding: fresh view of inhibition in the retina". " In: Proceedings of the Royal Society of London. Series B. Biological Sciences 216.1205 (1982), pp. 427-459, DOL: 10. 1098 1098/rspb. 1982. 0085
[62]
Karl Friston. "A theory of cortical responses", In: Philosophical transactions of the Royal Society B: Biological Sciences 360. (2005), pp. 815-836, DOI: 10. 1099/rstb..... 2005
1622
[63]
Andy Clark. "Whatever next? Predictive brains, situated agents, and the future of cognitive science". In: Behavioral and brain sciences 36.3 (2013), pp. 181-204. DOI: 10. 1017
[64]
William Lotter et al. "A neural network trained for prediction mimics diverse features of biological neurons and perception". In: Nature machine intelligence 2.4 (2020); pp. 210-219.
DOL: 10.
[65]
Beren Millidge et al. "Predictive coding networks for temporal prediction". In: PLOS Compu-
tational Biology 20.4 (2024), c1011183, DO1: 10. 1371/
[66] Yosef Singer et al. "Hierarchical temporal prediction captures motion processing along the visual pathway". In: Elife 12 (2023), c52599. DO1: 10. 75544/5ife...... 52599
[67]
Pierre-fitienne H Fiquet and Eero P Simoncelli. "A polar prediction model for learning to represent visual transformations" In: Whirty Conference Neural Information Processing Systems. 2023. URL.: ht https::/ppenrreieee..eee.
1d=hy
[68] Bruno A Olshausen and David J Field. "Sparse coding of sensory inputs" In: Current opinion in neurobiology 14.4 (2004). pp. 481-487. Do1: 1..11.1.1.111 conb. 2004. 07. 007
[69] William E Vinje and Jack L Gallant. "Sparse coding and decorrelation in primary visual cortex during natural vision". In: Science 287.5456 (2000), pp. 1273-1276. DOI: 10. 1126/
science. 0287.5456... 1273
[70] Alison L Barth and James FA Poulet. "Experimental evidence for sparse firing in the ncocor- tex". In: Trends in neurosciences 35.6 (2012), pp. 345-355. Do1: 10. 1016/ j tins. 2012
03.008
[71] R Quian Quiroga et al. "Sparse but not fraandmotthr-cell'coooiie in the medial temporal lobe". In: Trends cognitive sciences 12.3 (2008), pp. 87-91. Do1: 10. 1016/j tics. 2007
12.003
[72] Tomáš Hromádka al. "Sparse representation of sounds in the unanesthetized auditory cortex". In: PLoS biology 6.1 (2008), e16. DOl: 10. /371//jurnal..ppp 0060006
[73] Cindy Poo and Jeffry S Isaacson. "Odor representations in olfactory cortex: "sparse" coding, global inhibition, and oscillations". In: Neuron 62.6 (2009), pp. 861. DOI: 10. 1016/] neuron. 2009.0000.. 05. 022)
[74]
Jason Wolfe et al. "Sparse and powerful cortical spikes". In: Current opinion in neurobiology 20.3 (2010), pp. 306-312 312. DOI: [00 1016/j conb. 2010. 03. 00.006)
Ben DB Willmore et al. "Sparse coding in striate and extrastriate visual cortex". In: Journal of neurophysiology 105.6 (2011), pp. 2907-2919. DO1: 10. 1152/jn. 00594. 2010 [76] Bilal Haider al. "Synaptic and network mechanisms of sparse and reliable visual cortical activity during nonclassical receptive field stimulation". In: Neuron 65.1 (2010), pp. 107-121.
[75]
DOL: 10.6/.0../. www.neuuunn0000
12. 005
[77] Sylvain Crochet et al. "Synaptic mechanisms underlying sparse coding of active touch". In:
Neuron 69.6 (201 1). pp. 1160-1175 DOL: 10. 1016/j /016//.neuroo........... www.weuuooo.om 20w.urrrnn001..2222
[78] Carl CH Petersen. "Sensorimotor processing the rodent barrel cortex" In: Nature Reviews
Neuroscience 20.9 (2019), pp. 533-546. Do1: 10. 00039///11583-199--222
14

---

[79] Emmanouil Froudarakis et al. "Population code in mouse VI facilitates readout of natural scenes through increased sparseness". In: Nature neuroscience 17.6 (2014), pp. 851-857. DOI: [0. 1038/nn .307
[80] Christopher Rozell et al. "Sparse coding via thresholding and local competition in neural circuits". In: Neural Computation 20.10 (2008), pp. 2526-2563. DOI: 10. www.11662ccccoooo 2008
03-07-486)
[81] Daubechies al. "An iterative thresholding algorithm for linear inverse problems with sparsity constraint". In: Communications on Pure and Applied Mathematics 57.11 (2004),
pp. 1413-1455 1457. DOI: 1002100/ppp 20042
[82] Amir Beck and Mare Teboulle. "A Fast Iterative Shrinkage- Thresholding Algorithm for Linear Inverse Problems". In: SIAM Journal Imaging Sciences 2.1 (2009), pp. 183-202
DOL: 1111110700000165555
[83]
Ankush Ganguly et al. "Amortized Variational Inference: A Systematic Review". In: Journal of Artificial Intelligence Research 78 (2023), pp. 167-215. DOI: 10. www.161/jaii... 14258 Brandon Amos. "Tutorial on Amortized Optimization". In: Foundations and Trends@ in Machine Learning 16.5 (2023), pp. 592-732. ISSN: 1935-8237. p01:[10.
[84]
[85] Samuel Gershman and Noah Goodman. "Amortized inference in probabilistic reasoning". In: Proceedings of the annual meeting of the cognitive science society. Vol. 36. 36. 2014. URT.: https://eschol arship wtwp uc/ /eeeee/33/hh
[86] Colin Conwell et "What can 1.8 billion regressions tell us about the pressures shaping high-level visual representation in brains and machines?" In: bioRxiv (2023). DO1: 10. 1101/
20022.3..2 4858688
[87] Eric Elmoznino and Michael Bonner. "High-pen performing neural network models of visual cortex benefit from high latent dimensionality". In: bioRxiv (2022), pp. 2022-07, DOE:
10. 10.1001/0022.0 07. 07.15 13. 499969
[88]
Aaron van den Oord et al. "Neural Discrete Representation Learning". In: Advances in Neural Information Processing Systems. Ed. by Guyon et al. Vol. 30. Curran Associates. Inc., 2017. URL: https papers/ppppssseea/aaa nips cc paperrpppppppppp files / paper / 2017 / hash /
/a98af17e63a0oe /Blaaa/eeeaaaa0/0000022000000000
[89] Eric Jang et al. "Categorical Reparameterizatioo with Gumbel-SSfftaxx. International Conference Learning Representations. 2017. URL: https://openreviee wwwp:rreewiie.ene/ournnne
id=rkE3y85ee
[90] Hiromichi Kamata et al. "Fully spiking variational autoencoder". In: Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 36. 6. 2022, pp. 7059-7067. DO1: 10. 1609/ aaai
v3616.20665
[91] Victor Geadah et al. "Sparse-Coding Variational Autoencoders"". In: Neural Computation 36.12 (Nov. 2024), pp. 2571-2601. ISSN: 0899-7667. DO1: [0. 1162/necoo_oo.e _ap_a_______ a_01715
[92] Francesco Tonolini et al. "Variational Sparse Coding". In: Proceedings of The 35th Un- certainty in Artificial Intelligence Conference. Fd. by Ryan Adams and Vibhav Gogate. Vol. 115. Proceedings of Machine Learning Research, PMLR, July 2020, pp. 690-700. URI.: https: :ttps:://proceedingg mlr. press/ v//vvv11111 wtwl:/tooooliiiieeeoooooo ttso/inoonnninnn 20a. html
[93] Pan Xiao et al. "SC- VAE: Sparse bodding-bssed Variational Autoencoder with Learned ISTA".
In: (2024). arXiv: 2303. 16666 [cs. CV]
[94]
David Roxbee Cox and Valerie Isham. Point processes. Vol. 12. CRC Press, 1980.
[95]
Oleksandr Shchur al. "Fast and Flexible Temporal Point Processes with Triangular Maps".
In: Advances in Neural Information Processing Systems. Ed, by H. Larochelle et al. Vol. 33. Curran Associates, Inc., 2020. pp. 73.84. URL: http:/ttttss//ceeapiiiiis neurips...........eee..rrrr.r. cc/ paper seger_files/ppeeer 2020/ htt.haaahhhhaa http////w//..ccso/eeeoooooooaaa. html
[96] Oleksandr Shchur. "Modeling Continuous-time Event Data with Neural Temporal Point Processes". PhD thesis. Technische Universität München, 2022. URL: htttps:/meed: dissertation ub..um.aem de/ http://66611//22aaaee33332222210 ott.://ec.....oooooooooo....aaa. pdff
[97] Adam Paszke et al. "PyTorch: An Imperative Style, High-Performanee
Deep Learning Library". In: Advances Neural Information Processing Systems. Vol. 32. Curran Associates, Inc., 2019. URL: https papers nips cc /apperpp_... files /paper/ / 2019 / hash / 2cbba2288feer799922fffaaaa
15

---

[98] Chris J. Maddison et al. "The Conerete Distribution: A Continuous Relaxation of Discrete Random Variables". In: International Conference on Learning Representations, 2017. URL:
httpp://openreviee.. httperrwwwwweeeeeeeeeeeeeeeeeeee
[99] Daniel J Felleman and David C Van Essen. "Distributed hierarchical processing in the primate cerebral cortex". In: Cerebral Cortex 1.1 (1991), pp. 1-47. DOI: 10.
[100] Nikola T Markov et al. "Anatomy of hierarchy: feedforward and feedback pathways in macaque visual cortex". In: Journal of Comparative Neurology 522.1 (2014), pp. 225-259.
DO1: 10. 1002/ cne. 23458
[101] Anita A Disney. "Neuromodulatory control of early visual processing in macaque". In: Annual Review of Vision Science 7 (2021), pp. 181-199. Do1: 10. 1146/annurevy annurev-viste
100119-125747
[102]
Irina Higgins et al. "beta-VAE: Learning Basic Visual Concepts with Constrained Variational Framework". In: International Conference on Learning Representations. 2017. URL: https
//openrvview. ataee/fforummiddSS2222lggg
[103] James C. R. Whittington et al. "Disentanglement with Biological Constraints: Theory of Functional Cell Types". . In: The Eleventh International Conference on Learning Representa- tions. 2023. URL: httph: btttts::/openreviiee
met/forumiideeeeeeeeeee GfnzGGH
[104] Hans Van Hateren and Arjen van der Schaaf. "Independent component filters of natural images compared with simple cells in primary visual cortex". In: Proceedings of the Royal Society of London. Series B: Biological Sciences 265.1394 (1998), pp- 359-366. DOE: 10
1098/rspb. 1998. 0303
[105] Victor Boutin al. "Sparse deep predictive coding captures contour integration capabilities of the early visual system". In: PLoS computational biology 17.1 (2021), e e1008629. DOT: 10.1331/journal.1 www.pp..i..0022999
[106] Alex Krizhevsky, Geoffrey Hinton, et al. "Learning multiple layers of features from tiny im- ages". In: (2009). URL: http::ttt://// http://wwwwcco cs. toronto. wdw/"krri/leerr http://rriii/eeeannnneeeeee features. 2009-TR.
[107] Yoshua Bengio et al. "Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation". In: (2013). arXiv: 1308.3432
[108]
Francesco Locatello et al. "Challenging common assumptions in the unsupervised learning of disentangled representations" In: international conference machine learning. PMLR. URL:
2019, pp. 4114-4124. html
mlr press/ 1oc/lccatlalllll
[109]
David H. Hubel and Torsten N. Wiesel. "Receptive fields of single neurones in the cat's striate cortex". In: The Journal of Physiology 148 (1959). DO1: 10. 10.1113//physiol. 1959
sp006308
[110] David H. Hubel and Torsten N. Wiesel. "Receptive fields and functional architecture of monkey striate cortex". In: The Journal of Physiology 195 (1968). DOI: 10. 1113/jphysiol 1968. sp008455
[111] Alexander Alemi et al. "Fixing a Broken ELBO". In: Proceedings of the 35th International Conference on Machine Learning. Ed. by Jennifer Dy and Andreas Krause. Vol. 80. Pro- ceedings of Machine Learning Research. PMLR, July 2018, pp. 159-168. URL: https \/proceedings. mlr www.ppreeee/veee alemi 18a, html
[112] Chris Cremer et al. "Inference suboptimality in variational autoencoders". In: Interna- tional Conference on Machine Leurning. PMLR. 2018, pp. 1078-1086. URL.: https: proceedings.o mlr prwss/v80 c::::0//eeeee11em 18a.. html
[113] Joe Marino et al. "Iterative Amortized Inference". In: Proceedings of the 35th International Conference Machine Learning. Ed. by Jennifer Dy and Andreas Krause. Vol. 80. Pro- ceedings of Machine Learning Research, PMLR, July 2018, pp. 3403-3412 URL: https //proceedings.noo .mr..mmmm....mm wttp:peess//////riiiiaaaaa
html
[114] Kilian Q. Weinberger and Lawrence K. Saul. "Distance Metric Learning for Large Margin Nearest Neighbor Classification". In: Journal of Machine Learning Research 10.9 (2009),
pp. 207-244 207-244. URL.: http:://jmil
org/ pappers/v10/viinerrgeeeee
html
[115] Mattco Alleman et al. "Task structure and nonlincarity jointly determine learned represen- tational geometry". In: The Twelfth International Conference on Learning Representations.
2024. URL:: hattps://operrevie. ppperrovww..eefffooodr///aaaaaa
16

---

[116] Mattia Rigotti et al. "The importance of mixed selectivity in complex cognitive tasks". In: Nature 497.7451 (2013), pp. 585-590. DOI: 10. 1038/nature121661
[117] Silvia Bernardi et al. "The geometry of abstraction in the hippocampus and prefrontal cortex".
In: Cell 183.4 (2020), pp. 954-967. DO1: 10. 1016/j cell. 2020. .900
[118] Matthew T Kaufman et al. "The implications of categorical and category- -ree mixed selectivity on representational geometries". " In: Current opinion in neurobiology 77 (2022), P. 102644. DOI: 10. 1 conb. 2022. 102644
[119] Karol Gregor and Yann LeCun. "Learning fast approximations of sparse coding". In: Proceed- ings of the 27th international conference on international conference on machine learning, www.o.rg/.......oo.a....a..oo... 2010. pp. 399-406. URL: bittps://di acm. org/ doi /ato//////////a/// 10. 5555/3104322. 3104374
[120] Yoon Kim et al. -eoot--mot Variational Autoencoders".. In: I Proceedings of the 35th International Conference on Machine Learning. Ed. by Jennifer Dy and Andreas Krause. Vol. 80. Proceedings of Machine Learning Research. PMLR, July 2018, pp. 2678-26871 2687. URl. https://proceedan
mlr. prwss/esseeeee vtt5//80///// wwwlme..........m....m....m..m.. html
[121] Edgar Y Walker et al. "Inception loops discover what excites neurons most using deep predictive models". In: Nature neuroscience 22. 12.12(6 (2019), pp. 2060-2065, DO1: 10. 1038/
841593-019-0511
[122] Charles R. Harris et al. "Array programming with NumPy". In: Nature 585.7825 (Sept. 2020),
pp. 357-362. DOI: 10.
020-586-020--2800 2049-2669-22
[123] Pauli Virtanen et al. "SciPy 1.0: Fundamental Algorithms for Scientific Computing in Python".
In: Nature Methods 17 (2020), pp. 261-272. DOI: [0. 1038/ 14119991199-966-------
[124] Fabian Pedregosa et al. "Scikit-lear: Machine learning in Python" In: the Journal of machine
Learning research 12 (2011), 1,, pp.
DOI: 10. 5555/ 1953048 2078195
[125] The pandas development team. dandas-dev/pandas: Pandas. Version latest. Feb. 2020. DOT:
10.528:11 2281//eeodo.3509134
[126] John D Hunter. "Matplotib: A 2D graphics environment". In: Computing in science & engineering 9.03 (2007), pp. 90-95. DOI: 10. 1109/MCSE, 2007 .5
[127] Michael L Waskom. "Seaborn: statistical data visualization". In: Journal of Open Source Software 6.60 (2021), 3021. DOt: 10. 20105/ jts///30ss00001
[128] David J Tolhurst et al. "The statistical reliability of signals in single neurons in cat and monkey visual cortex". In: Vision research 23.8 (1983), pp. 775-785, DOI: 10. 1016/0042- 6989 (83) 90200-6
[129] AF Dean. "The variability of discharge of simple cells in the cat striate cortex". In: Experi- mental Brain Research 44.4 (1981), pp. 437-440. DOI: 10. 1007 /F000/FF0038837
[130] Michael N Shadlen and William T Newsome. "The variable discharge of cortical neurons: im- plications for connectivity, computation. and information coding" In: Journal of neuroscience 18. 10 (1998), pp. 3870-3896. DO1: 10. 1523/ JNEUROSCI. 18 10- 03870 1998
[131] Bruno B Averbeck et al. "Neural correlations, population coding and computation". In: Nature Reviews Neuroscience 7.5 (2006), pp. 358-366. DOL: 10. 103.00.0100/mm188...3
[132] Peter Dayan and Laurence F Abbott. "Theoretical Neuroscience". In: (2001). URL: https
//mitpress wit.mmittmmmmmmmm
[133] Wilson Truccolo et al. "A point process framework for relating neural spiking activity to spik-
ing history, neural ensemble, and extrinsic covariate effects". In: Journal of neurophysiology
93.2 (2005), pp. 1074-1089. DOT: 10. 1152/jn...n.... jo 00697 2004
[134] Zachary F Mainen and Terrence J Sejnowski. "Reliability of spike timing in neocortical
neurons". In: Science 268.5216 (1995), pp. 1503-1506. DOI: [10
7770778
[135] Michael R WeWeese et al. "Binary spiking in auditory cortex". " In: Journal of Neuroscience 23.21 (2003), pp. 7940-7949. DOL: 10. JNER/SNE 23-21-07944 2003
[136] William H Calvin and CHARLES F Stevens. "Synaptic noise and other sources of randomness in motoneuron interspike intervals". In: Journal of neurophysiology 31.4 (1968), pp. 574-587.
DOI: 10. 1112/ jn. 1968. 31. 4.574
[137] Christina Allen and Charles F Stevens. "An evaluation of causes for unreliability of synaptic transmission". In: Proceedings of the National Academy of Sciences 91.22 (1994), pp. 10380- 10383. DO1: 10. 20000///aa.../. [138] Matteo Carandini. " "Amplification of trial-to-trial response variability by neurons in visual cortex". In: PLoS biology 2.9 (2004), e264. Do1: 10. 1171/ http:/joonnal.ppii 2
22. 10380
17

---

[139] Alison I Weber and Jonathan W Pillow. "Capturing the dynamical repertoire of single neurons with generalized linear models". In: Neural Computation 29.12 (2017), pp. 3260-3289. DOI:
10. 1162//eccooooo 01021
[140] Daniel A Butts al. "Nonlinear computations shaping temporal processing of precortical vision". In: Journal of Neurophysiology 116.3 (2016), pp. 1344 1357. DOI: 10. j152/jn
00878.2015
[141] Jagruti J Pattadkal et "Synchrony dynamics underlie irregular neocortical spiking". In: bioRxiv (2024), pp. 2024-10. DOI: 10. 1101/2024 10. 15. 618398
[142] Casper Kaae Sønderby et al. "Ladder Variational Autoencoders". In: Advances in Neural Infor-
mation Processing Systems. Vol. 29. Curran Associates, Inc., 2016. URL: https: httt:hhtt:::/pppeer
http://w......ppp.o////////aaaa2 Abstract html
[143] Cina Aghamohammadi et al. "A doubly stochastic renewal framework for partitioning spiking variability". In: bioRxiv (2024), pp. 2024-02.
[144] David M Blei et al. "Variational inference: review for statisticians". In: Journal of the American statistical Association 1 1125118 (2017), pp. 859-8777 DOT: 10.
2017. 1285773
[145]
Ashish Vaswani et al. "Attention is All you Need". In: Advances in Neural Infor- mation Processing Systems. Ed. by I. Guyon et al. Vol. 30. Curran Associates, Inc., 2017. URL: https / papers nips cc paper files / paper / 2017 / hash / Abstaaattbbttatt..hh.
[146]
Prajit Ramachandran et "Searching for Activation Functions". In: International Confer- ence on Learning Representations. 2018. URL: httpps:/ wwt. forum? id=
SkBYYyZRZ
[147]
Stefan Elfwing et al. "Sigmoid-weighte linear units for neural network function approxima- tion in reinforcement learning". In: Neural Networks 107 (2018), pp. 3-11. DOI: 10 1016/
neunet. Sounet.2017.12. 2017.12.012
[148]
Diederik P Kingma and Jimmy Ba. "Adam: A method for stochastic optimization". In: (2014). Xrxiv: 1412. 6980 [cs.LG]
[149] Ilya Loshchilov and Frank Hutter. "SGDR: Stochastic Gradient Descent with Warm Restarts". In: International Conference Learning Representations. 2017. URL: https //openreiiew. net/ seeet/forummice idtrummmdddSS3388cxr
[150] Samuel R. Bowman et al. "Generating Sentences from a Continuous Space". In: Proceedings of the 20th SIGNLL Conference on Computational Natural Language Learning. Berlin, Germany: Association for Computational Linguistics, Aug. 2016, pp. 10-21. DO1: 10. 18653/ v1/K16-1002
[151] Hao Fu al. "Cyclical Annealing Schedule: A Simple Approach Mitigating KL Vanishing". In: Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Pupers). Minneapolis, Minnesota: Association for Computational Linguistics, June 2019,
pp. 240-250, DOI: 1.0.18865/v1//// 1/1///11119912/11 1021
[152]
Michael Teti. LCA-PyTorch. [Computer Software] https //doi org/ 11578/1111/66 20230728. 4 June 2023. DOI: 10. www11/17//cc.. www.dddcc.2222....
URL: https://doi.org///
[153]
Shakir Mohamed et al. "Monte Carlo Gradient Estimation in Machine Learning". In: Journal of Machine Learning Research 21.132 (2020), pp. 1-62. URL: http://jmlr... org/ papers/
v21/19 2011/99-2 146119999444.... html
18

---

A Are real neurons truly Poisson?
In this section, we discuss empirical and theoretical observations from neuroscience that motivated our Poisson assumption.
"Poisson-like" noise in neuroscience has long history. It begins with observations that neurons do not fire the same sequence of spikes to repeated presentations of the same input, and that the variance proportional the mean [128. 29 and was followed by the observation that for short counting windows, that proportionality is one 49 50 132). Larger windows and higher visual areas are notably super-Poisson, but that can be attributed to modulation of the rate of an inhomogeneous Poisson process [51].
In other words, neurons are conditionally Poisson, not marginally Poisson
133
Spike-generation, is argued, is not noisy [134 [36). but synaptic noise 137
noise on the
membrane potential, can create a Poisson-like distribution of spikes (138). An important caveat is
that the well-known example of precision in spike generation by Mainen and Sejnowski [34] is effectively captured by Poisson-prosse process Generalized Linear Model (GLM: Weber and Pillow 139 ).
However, this precision relies on a Bernoulli approximation to Poisson process, allowing only 0 or 1 spikes. There is widely- misconception that precise timing cannot be produced by spike-rate models, but inhomogeneous rate models can produce precise spiking patterns at high time resolution [10).. In contrast, recent work has shown that correlations in excitatory inputs drive Poisson-like variability, explaining the widespread observation of Poisson-like noise in real neurons [14]..
In summary, neurons are not literally Poisson, but it is a good choice. To set up the ELBO, one has to choose approximate posterior and prior. Because spike counts are integer and cannot be negative, Poisson is a more natural choice than Gaussian without knowing anything about neural firing statistics. Here, we found that the Poisson assumption led to model with interesting theoretical and empirical properties, where sparse coding emerged from the ELBO with Poisson.
Extending the P-VAE hierarchical architectures 2 Poisson, but not marginally Poisson (as they are modul ated by top-down rates). Further extensions could implement doubly-stochastie spike generation [1 143
142 will make the latents conditionally
B Full derivations
In this section, we provide self-contained and pedagogical introduction to VAEs, derive the P-VAE loss function, and highlight how combining to the emergence of a metabolic cost term in the PAVAE loss. For the case of a linear decoder, the reconstruction loss assumes closed-form solution. This means compute the gradients analytically, which we can then use to evaluate the Poisson reparameterizatioo trick.
latents with predictive coding leads
B.1 Deriving the evidence lower bound (ELBO) loss
For completeness, let's first go over the basics. This section will provide a quick refresher on variational inference and how to derive the VAE loss from scratch. Assume the data E R.M and K -idimensional latent variables are jointly distributed as p(x, z), with the data generated through the following process:
p(x) f p(zx dz p(zlz))((z)
/ dz,
(6)
In Bayesian posterior inference, the goal is to identify which latents z are likely given data æ. In other words, we want to approximate P(z|z). the optimal but (typically) intractable posterior distribution.
B.1.1 Variational inference and VAE loss function
To achieve approximate Bayesian inference, common approach is to define family of variational densities @ and find a member g[>|==) e @ such that it sufficiently approximates the optimal posterior [14].. We call <<((=) the approximate posterior. The general aim of variational inference (VI) can be summarized as follows:
VI: find (2((x) (x) E Q such that g(z) (x) is good approximation of p(z(x))
(7)
19

---

The goodness of our approximate posterior, or its closeness to the true posterior, is measured using the Kullback-Leibler (KL) divergence:
qº' = argmin DEL (((2|z)) p(z
(8)
We cannot directly optimize eq. 8 because p(z|ar) ()) is often intractable. Instead, we rearrange some terms and arrive at the following loss function:
(4) -Egaa(())) [ogp(z|z) + DxL
D (1(2)2) ||
NELBO stands for negative ELBO, also known as "variational free energy." Notably, finding a q E Q that minimizes (NELBO(q) in eq. is equivalent to finding the optimal q" in eq.
The first term in eq. ( often called the reconstruction term, captures the likelihood of the observed data given latents 2, under the approximate posterior. For all our VAE models, we assume Gaussian conditional likelihood with a fixed variance, as is typically done in the literature. This approximates the reconstruction term the mean squared error between input data and their recon- structed version. The second term, known as the KL term, is more interesting. This term can assume very different forms depending on the distribution used.
B.2
The KL term
In this section, will derive closed-form expressions for the KL term for different choices of the distributions ((>(<) and p(z). Specifically, we will focus on Gaussian and Poisson parameterizations.
Predictive coding assumption. We will draw inspiration from predictive coding and assume that the bottom-up inference pathway only encodes the residual information relative to the top-down, or predicted information. We will apply this idea to both Gaussian and Poisson cases, and find that only in the Poisson case, the outcome becomes interpretable and resembles sparse coding objective.
B.2.1
KL term derivation: Gaussian
Let g(= a) N(z; (x), ag (a)) and p(z) =
where the mean and variance are either outputs of the encoder network or parameters of the decoder network.
Now, 1et us implement the predictive coding assumption, where the encoder only keeps track of residual information that is not already contained in the prior information. Mathematically, this idea can be formalized as follows:
Hp 14, o - ,
My p 8p
(10)
Tg + a 80
With these modifications, the Gaussians KL term becomes:
1 84 DEL (q (1p) = -
- + 500 plo--ooobo
log d0 -
(11)
In standard Gaussian VAEs, the prior has no learnable parameter. Instead, we have μ - and - 1. Therefore, the final form of the KL term for standard Gaussian VAE is:
Dra (q (O, 1)) = (00² + 000 log do² - 11.
6n²²
(12)
We observe that the KL term vanishes when مز - and d0 1. This happens whenever information propagated through the encoder, phenomenon known as posterior collapse.
Other than this trivial observation, eq. 12p does not really lend itself to interpretation. In contrast, will show below that a Poisson parameterization of VAFs leads much more interpretable outcome for the KL term
20

---

B.2.2 KL term derivation: Poisson
Now suppose ((=[æ) = Pois(2;rór(a))) and p(z) (ois(2; r), where 2 is literally the spike count of single latent dimension -r shall we say, neuron?
In the Poisson case, the KL term becomes more interpretable, as we will show below. Recall that the Poisson distribution for a single variable 2, given rate À 250. is given by:
A
Pois(2; A)
(13)
z!
Plug this expressions into the KL divergence definition to get:
Dau (q = Exng
log
( /2/
E8 1gz~g
log
r6E-1/z!
log (nár
rrr
E Egzz
--8rtt
Egzg
log dr + +log -rertr
(14)
E Ezng
log 8r rōr
E Ezng
log 8r r8r +
rrr log 8r rōr +
(1 + log dr)
f(rr),
where we have define - y log y-
To examine the behavior of the Poisson KL term, we assume 8r = 1 + c, where e < then Taylor expand f. Calculating the first and second derivatives of f(y) y + y log y gives f'()) log
and f"()) = 1/8. Thus:
(0+
('1)+ 0(6))
21
=0+0+
2! +
(15)
2
Plug this back into eq. (14) to get:
DeC (p) = r((rr)
r(1 +
(16)
~
2
For small deviations c, the KL term simplifies to the product of the prior firing rate, r, and c See 6 for a visualization of the full function, f(ōr) ör + 8r log ör, along with its quadratic approximation near ör 1.
In general, there are two ways minimize the term: dead prior neurons (r - 0), or posterior collapse (ör 1).
21

---

4
f(ö
0.5 * (1 5r)
6r
109
or
2
+
1
0
1
2
3
4 5
1
2
4
5
5r
5r
Figure 6: Left, residual term ((ór) from eq. (4).. Right, quadratic approximation of f from eq.
115
Together with the reconstruction loss, the NELBO for a 1-dimensional P- reads:
LPVAR (r, 5r) Crecon. (r, + 5r + 6r log ör).
(17)
Finally, it is easy to show that for K K-dimensional latent space, eq. (14) generalizes to:
DEL Pois((z; @ dr(())) Po r f(ōr),
(18)
where and denote the Hadamard (element-wise) and vector products, respectively.
B.3 Connection to sparse coding
Equation 17 mirrors sparse coding due the presence of the firing rate in the objective function. Furthermore, it follows the principle of predictive coding by design. Thus, our Poisson formulation of VAEs effectively unifies these major themes in theoretical neuroscience. Let's explore this curious connection to sparse coding more closely below.
B.4 Statistically independent neurons
Suppose our P- P-VAE has K statistically independent neurons, and Z the spike count variable, where Zzo {0, 1, ... is the set of non-negative integers. Let us use bold font and ör refer to the firing rate vectors of the representation and error units, respectively. Recall that we allowed these variables to interact in a multiplicative way construct the posterior rates, >i(x) (w)) More explicitly, we have:
K
X
<<=xx) P 2 ôr) = II Pois(za3;00))
zd!
(19)
K
K
r
p(z) Pois(=;
II (ois(2
2,!
i=1
Note that, unlike a standard Gaussian VAF, the prior in P-VAE is parameterized using which is learned from data along with the other parameters. Similar standard Gaussian VAEs, 8r(x)) parameterized as neural network.
22

---

B.5 Linear decoder
Following the sparse coding literature [3, we will now assume our decoder generates the imput image IR linear sum K basis elements, & R RMX Additionally, we choose diagonal Gaussian distribution with fixed variance as our conditional likelihood, resulting in a mean squared error between the input æ, and its reconstruction 42.
Given these assumptions, the reconstruction loss for VAE with approximate posterior y can be expressed as follows: Lrooon. (x; = L2 =(2xx()()x==)
2211
(20)
For a linear decoder, the reconstruction term x 22 contains only the first and second moments
of Consequently, the expectation in eq. (20) can be analytically computed. This results in
close-form expression for the reconstruction loss, and consequently, its gradients as well.
In general, whenever the VAE decoder is linear, the following result holds:
recon.
x q, ④ =
x + E 1 + Var,I2] diag +T
(21)
Note that linear decoder is the only assumption we needed to obtain this closed-form solution. There are no restrictions the form of the encoder: it can be linear, complicated as want. We only have to compute the mean and variance of the posterior.
Specifically, for the Poisson case, only need know the following expectation values:
Bg~Pois(=:>)) [2 = Xis
(22)
E E=~
[22 = Xx2g + 5.z22.
Here are the reconstruction losses for both Poisson and Gaussian VAEs with linear decoders, put side-by-side for comparison:
Poisson:
Creoon. (x;; X, ④ = x +b A1 + X diag( o Φ);
(23)
Gaussian:
Lrecon. (x: pl,
x 1 ? (a²)² diag( T)
Given these assumptions, the NELBO (eq. ( for P-VAE with linear decoder becomes:
K
LScE (æ;; dr, (b) = || +> 2 + d'diag( o (T - B rs (örss)
(24)
Recall that we have fly) = 1 log (see Fig. 60. We introduced the 3 term here to control
the trade-off between the reconstruction and the term 102 Additionally, we dropped the explicit dependence of (r(x) on the input image æ to enhance readability.
B.6
Linear encoder
We can further simplify the P-VAE architecture by making the encoder also linear. Let W E R KxM denote the encoder's weight matrix, and assume exponential link function mapping the input to residual firing rates, i.e., dr exp( (Vx).
Starting from eq. 12 substituting log ör = W æ, and rearranging terms yields the following loss function for the {lin lin} P-VAR:
LLin.VAF = Th
X &diag( (o -33) (3W 20 )) +
Ti + x æ. (25)
23

---

C Architecture, training, and hyperparameter details
C.1 Datasets: additional details
We consider three datasets in this paper. We tile up the van Hateren dataset of natural images [04 and CIFARIO into 16 x 16 patches and apply whitening and contrast normalization using the code made available by Boutin al. [105]. This operation results in the following total number of samples:
van Hateren:
#train = 107,520,
#validation = 28,224, #train = 200,000, #validation = 40,000.
CCIARR6616:
We use the MNIST dataset primarily for the downstream classification task. After the training is done, we use the following train/validation split to evaluate the models:
K-nearest neighbor classification (tables 4 and 6): For this task, we only make use of the validation set for both training and testing of the classifier. We divide up the N 10,000 validation samples into two disjoint sets of N = 5,000 samples each. We then draw random samples (without replacement) from the first half and use them for training the KNN classifier. We then test the performance on the other half.
Shattering dimensionality (tables and 6 last column): We use the entire MNIST training set (N 60,000 samples) to train logistic regression classifiers on extracted representations. We then test the results using the entire validation set (N 10,000 samples).
C.2 Architecture details
For sparse coding results, focused on models with linear decoders. For the fully linear models (Figs. 4 and 10) both the encoder and decoder were linear layers, without bias
For the convolutional components, we use residual layers without batch norm. For van Hateren and CIFAR 16x 16 datasets, the encoders had layers (2 x conv each). The decoders had convolutional layers (1 x conv each). For the MNIST dataset, the encoders had 7 layers (2 conv each). The decoders had 10 convolutional layers (1 x conv each). For all convolutional encoders, the output from ResNet was followed by a learned pooling layer. The pooled output was then fed into feed-forward layer inspired by Transformers [155,, which includes a layer norm as the final operation, the output of which fed into a linear layer that projects features into posterior distribution parameters. For all convolutional decoders, nearest neighbor upsampling was performed to scale up the spatial dimension of reconstructions, suggested by Child (34).
We experimented with both leaky_relu and swish activation functions 146
and found that swish consistently outperformed leaky_relu in all our experiments across datasets and VAE models.
147
Please see our code for the full architecture details.
C.3 Training details
We used a variety of learning rates and batch sizes, depending on the dataset and architecture. For (lin|lin) and (conv]lin) models, we used (r = 0.005, and for (conv) (conv/conv) models we used 1 0.002. All models were trained using the AdaMax optimizer [148) with cosine learning rate schedule 149 Please see our code for the full details of training hyperparameters, Overall, we trained 195 VAE models, n = 5 seeds each, resulting in a total of 195 5 975 VAEs. For sparse coding models, we ran ISTA [81] [2) and LCA [80] with 270 hyperparameter combinations each. Training all models took roughly a week on RTX 6000 Ada GPUs
Temperature annealing for discrete VAEs. We also annealed the temperature from a large value to smaller value during the same first half of training for P-VAE and C-VAE. We found that the specific functional form of temperature annealing (e.g., linear, exponential, etc.) did not matter as much as the final temperature (Fig. 9). For both P-VAE and C-VAE, we start from T 1.0 and
anneal down to T Tapp 0.05 for P-VAE, and T Tump 0.1 for C-VAE, We found that the C-VAE performance was not very sensitive to the choice of T stp.. corroborating previous reports 89 98
24

---

The P-VAE was relatively more sensitive to the value of Topp. and we found marginal improvements when reducing from 0.1 0.05. See for comprehensive experiments exploring the effect of the final temperature, as well as a "hard-forward" training method where we set 7 = 0 in the forward pass (ensuring integer samples) and use non-7ero T only during the backward pass (surrogate gradients). We find that our "relaxed Poisson" approach (Fig. 3 consistently outperforms the hard-forward approach.
KL annealing for VAEs. For all VAE models, we annealed the KL term during the first half of the training, which is known be an effective trick in training VAEs 33
142
150
15T
C.3.1 Training: sparse coding models
To fit LCA and ISTA models, we explored a combination of 3 schedules (same 3 as in eq. ( 3 numbers of iteration (for inference), 3 learning rates, and 5 different seeds (for dictionary mi- tialization). The code for LCA was obtained from the public python library 'Ica-rytorch' (152)) and the code for ISTA obtained from public 'sparsecoding" repository of the Redwood Center for Theoretical Neuroscience (with added clipping of coefficients to be nonnegative, following the thresholding step).
We explored learning rates of x 10-1, x 10-2, and 10-3. We trained all models for 100 epochs. We scheduled the / parameters linearly, starting from Buant and stepped it up every five epochs by Sstpp until it reached Bend. We explored the following 3 schedules (expressed as Snnatt
0.05:0.7:0.1, 0.01:0.1:0.01, 0.1:1.0:0.1,
0.05:0.7:0005,
0.05:0.5:0.05,
0.1:0.1:0
We also explored the inference iteration limits of 100, 500, and 900 iterations. We selected the best fits to include in the main results shown in Figs. and 5
D Supplementary results
In this section, include additional results that further support those reported in the main paper, including:
Table 5 contains the negative ELBO values for all VAE models with linear decoder. This table reveals a comparable performance between using Monte Carlo samples to estimate gradients, versus optimizing the exact loss (see eqs.
and 24 highlighting the effectiveness of our
Poisson reparameterization algorithm.
Figure uses the same data from the main paper Table 2 to visualize the effects.
Figure & the dependence of loss latent dimensionality. We find that increasing the number of latent dimensions consistently improves ELBO for (conv) lin) architectures, but [lin]lin) models either overfit (for van Hateren) or fail improve (for CIFAR16x16) once becomes large.
Figure D demonstrates the robustness of our Poisson reparameterization trick (Algorithm D to also explore "hard-forward" training
variations in the temperature parameter. Importantly, approach, where we fix T = 0 during the forward pass but allow T > 0 in the backward pass. This is also known surrogate gradients. We find that, somewhat surprisingly, this hard-forward method performs significantly worse than our "relaxed Poisson" approach (Fig.
Figure T0 shows how the distribution of KL values (or the norm of decoder weights in the case of linear decoders) can be used to determine dead neurons that don't 't contribute to the encoding of information.
contains the full set of downstream classification results. Related Table 4
Figure shows the performance of a simple linear classifier (logistic regression) trained on unsu- pervised representations learned by various (conv)con) conv) VAEs. We find that increasing the latent dimension (K) generally improves the performance of P-VAE, but at lower dimensions, other methods like CAE and G-VAE can outperform it.
Figure 12 shows MNIST samples generated from the latent space of models, as well their reconstruction performance.
different (conv/conv) VAE
25

---

Table 5: The reparameterized gradient estimators work well exact ones, datasets and encoder architectures (linear vs. conv). Note that exact gradients are only computable for linear decoders (see eqs. (21). (23) and (24). The values are negative ELBO (lower is better), shown as 99an=99% confidence interval calculated from nt = 5 different random initializations. For MNIST, our use of Gaussian conditional likelihoods means the numerical performance values are not directly comparable to studies that use binarized MNIST with cross-entropy decoder. EX, exact, MC, Monte-Carlo, ST, straight-through [107].
See also Table 2 and supplementary Fig.
van Hateren
16x16
MNIST
Model
[1in|lin)
(conv|1
(lin|lin)
(conv
[lin|lin)
(conv[lin)
EX
168.0#.8
162.4+.2
167.1
16..1+.1
41.3
39.7+.2
P-VAE
MC
167. 2
163. .6..4
167.3+.1
162.9
41.72.2
40.1+.2
ST
179.34.1
179.44.1
182.34.1
182.31.2
44.8
44.2
EX 160.3+.1
154.4+.1
165.9
149.2+.0
40.6#.1
40.0=.1
G-VAE
MC
160.3+.1
154.4+.1
165.9
149.2+.1
40.7#.1
40.1+.0
EX 174.9..1
186.31.8
177.11.1
180.61.5
56.11.7
59.11.0
C-VAE
MC
170.5
171.9+.2
174.7=.1
176.5+.1
39.7=.2
59.1#.0
ST
174.2+.2
181.1
180.24.0
185.6+.2
49.3+.1
63.8+3.4
EX 167.3:t.
159.01.2
170.1
154.31.1
42.1
41.04.0
L-VAE
MC
167.3#t0
159.2t.2
170.1
154.5+.1
42.1#.0
41.0#.0
van Hateren
CIFAR16x16
MNIST
Relative performace drop [%]
10
Linear encoder
5
10
Conv encoder
5
0
EX MC
ST
EX MC
ST
EX MC ST
Figure 7: Performance drop relative to the best fit. Blue circles indicate P- VAE results, red circles indicate G-VAE results, and each set of n circles corresponds to five random initializations. Using Monte Carlo samples [153] and our Poisson reparameterization trick (Algorithm T to estimate gradients performs comparably to using exact gradients (see eqs. (21). (23) and (24)). Table provides tabular summary of these results. EX, exact, MC, Monte-Carlo ST. straight- throuhhhhhhghhhh [00
26

---

Negative ELBO (lower is better)
170
Dataset: Hateren
168
{lin [1in)
(conv [in)
166
Dataset: CIFAR16x15
164
(lin[lin)
162
(conv [lin)
256 256
=
384
213
2024
The effect of latent dimensionality model performance across datasets and encoder archi- tectures. For all convolutional encoder cases, BLBB improves as a function latent dimensionality. However, for linear encoders, we see that the van Hateren dataset starts to overfit for K > 512, and it stagnates for the C1FAR16x16 dataset. conclusion, more expressive encoders can find nonlinear features, represented using additional latent dimensions, but simple linear encoders struggle to utilize additional dimensions. The gray triangle indicates the setting used in the main results.
Figure 8:
Anneal schedule: lin
Anneal schedule: exp
176
hard fwd
False
Negative ELBO (lower is better)
174
True
172
170
Teun x
176
174
Thmal
172
lar [portio)
170
168
166
(T = Tenal)
T=000 0.01 T
7 T
10000
1=00 0.4
500
TO
20
100
E3
0.4
1001
1.05
0.2
0.3
Figure 9: Performance function of the final temperature (Tina), annealing schedule (linear vs. exponential; inset), and the "hard-forward" approach. The hard-forward approach uses exact integer samples (T 0) in the forward pass and applies nonzero temperatures only in the backward pass (i.e., "surrogate gradients"). Although all results are evaluated at hard-forward approach still underperforms our "relaxed Poisson" method (Fig. 3). which employs
during testing, the
continuous (floating) samples during training due to non-zero T' (Algorithm l). The gray triangle
indicates the setting used in the main results: Ti 0.05 with linear annealing schedule.
27

---

Gaussian (# dead neurons: 401)
Poisson (#f dead neurons: 8)
dead alive
dead alive
10'
10
10
0.0
0.1
0.2
0.3
0.5
0.7
0.00
0.02
0.04
0.06
0.08
0.10
Laplace (# dead neurons: 416)
Categorical (#) dead neurons: 4)
dead alive
dead
10a
allve
10a
10c
Figure 10: Identifying dead neurons using histogram- method. We bin the values and determine the gap between small values and larger ones. We identify neurons with KL values lower than the identified threshold (black dashed lines) and pronounce them dead. The figure shows the distribution of KL values over all neurons (K 512) for P-VAE, G- G-V,E, and LAVAE. The KL term single number the C-VAF because its latent space consists of single one-hot categorical distribution with K = 512 categories. Therefore, for the C- VAE, we use the distribution of decoder weight norms instead. These are the same models shown Fig. 4 where both encoder and decoder linear. Table 3 uses this method to quantifyy the proportion of active neurons for VAEs across different datasets and the choice of encoder architectures.
Table Geometry of representations. Full set of results. Related to Table
Latent
KNN classification (N, labeled samples)
Shattering dim.
Model
dim.
N 200
= 1,000
N = 5,000
P-VAE
0.815 0.002
0.919
0.946
0.797 +.00
C-VAE
0.7055:.002
0.8004.002
0.853+.040
0.795-1.006
C-VAE
0.757=.003
0.869+.00 002
0.924+.025
0.751+.00 0008
= 10
G-VAE
0.673=.003
0.813+.002
0.891 003
0.758+..00
G-VAE +relu
0.694#.003
0.817+.000 003
0.877
0.762+ 007
G-VAE +exp
0.642
0.784.1.000
0.8631.032
0.737:1.000
P-VAE
0.825 .002
0.927+.00
0.957+.00s
0.935+.00 003
C-VAE
0.770=.002
0.880+.001
0.920+.009
0.899+.004
C-VAE
0.710
0.836:t.003
0.902+.038
0.770:t.007
50
G-VAE
0.604
0.746:+.002
0.837+.022
0.743:+.007
G-VAE +relu
0.710
0,844+.002
0,904+.022
0.786 006
G-VAE +exp
0.694 .033
0.836+
0.906+.02
0.762+ 007
P.VAE
0.807 1.02
0.925
0.9581.013
0.949
C-VAE
0.753=.002
0.876+.001 001
0.925+.005
0.8844.000
C-VAE
0.701#.001
0.830+.003
0.896+.016
0.767+.007
K 100
G-VAE
0.636#.008
0.789++002
0.875+.024
0.763+ 007
VAE +relu
0.757+.002
0.8814.001
0.9334.019
0.818
G- VAE +exp
0.693
0.846+t.02
0.918+.024
0.793+ 006
28

---

1.00
0.95
Accuracy
0.90
type
poisson
0.85
categorical
laplace
gaussian
0.80
gaussian-relu
gaussian-exp
0.75
= 10
50
K 100
Figure 11: Downstream classification performance using a simple linear classifier. After unsupervised
training of (conv)conv) VAFs on MNIST, we extracted latent representations and applied logistic regression. For K 100, P-VAE achieves the highest accuracy, while for 10, both C- C-E and G-VA G-VAE outperform it.
Gaussian VAE
Poisson VAE
Gaussian VAE
l
8
0 4
7 2 0
4
/
4
O
6 6
9 0
0
15 5
9
1
5 2
a
00 6 6600 000000 000000 2 2906903000
0
2
4
6
2 2 0
4 6
9
0
6699
6
8 2
0
0060003606
Laplace VAE
? 7
27350
5
0684667990
9 9869240007
7 2 6
4 9
6 6
9 9
0
15 5
&
4 6690960200
8
7
7 7 4 / 4 S
9
0
0
666640715
Laplace VAE
Categorical VAE
Poisson VAE
9
9 9
7 2 6 1
4 A
9
0 6 6
9 9
0 0
15 5
8
9 2
0
2
511
9
5
8
9
0
5
1
5
9
7 2
4
9
2
8 2
Categorical VAE
6
2
6 9
7
1
9
0
7 2
/ 4
9
0 6 9 5
9
6
8
O
7
6
7
7
/
9
9
6
9
3
6
9
0
3
O
7
Figure 12: Generated samples (left) and reconstruction performance (right). These results shown here are from models with a (conv/cony) architectures and latent dimensionality of K 10.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: provide comprehensive theoretical and empirical evidence support claims of (1) introducing the P-VAE and its reparameterization trick; (2) P- P-VAE containing amortized sparse coding special case; (3) P. VAE largely avoiding posterior collapse; and (4) P-VAE facilitating linear separabilityy of categories at better sample efficiency, in and supplemental appendices to
sections B and 4
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation as long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
The limitations of (1) Poisson possibly not being a perfect description of cortical activity, and (2) amortization gap, are shown explicitly and thoroughly discussed in sections 4 [nd Specifically, we have dedicated paragraph for limitations in section 5 We evaluated our claims using multiple well-known datasets such as the van Hateren natural images (104), CIFARIO, and MNIST, tasks such reconstruction, sparse coding, and classification.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification. asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best
30

---

judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: We provide the full derivation of the P-VAE loss function, which is self- contained in the paper (section and supplement (appendix B
Guidelines:
The answer NA means that the paper does include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We disclose all details relating to the algorithm, including the optimization objective (eq. (, architecture and training details (appendix [ and pseudo-code for
Poisson reparameterized sampling (Algorithm Tp. In addition, we intend release all code and data needed for replicating our work.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data provided or not.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (c.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
31

---

(c) If the contribution is a new model (e.g.. a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: Our code, data, and model checkpoints are available from the following GitHub repository: http://gtuub..commaaavaaaaiiiii Poisson VAE
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/ //nips. wwtpppblcc/guueeeeeeeeoeeeppeeee hawpl///gguddd///oeeubbbiisssss
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: All details about how the data used for training and testing. well as which hyperparameters were used, are available at appendix C In addition, the provided code replicates results and therefore contains all details of implementation.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
32

---

Answer: [Yes]
Justification: As stated in section 4 the paper reports confidence intervals and Ł-test signifi- cance tests, using false discovery rate (FDR) correction for multiple comparisons. The exact implementation details are included in the provided code for reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by enror bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard crror of the mean.
OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide details about our compute resources (GPUs), and duration of training in section
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///wwwwp...coooooccccccaaaa
Answer: [Yes]
Justification: Our paper follows the code of ethics, including preserving anonymity (such as in releasing code anonymously).
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
33

---

The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our paper is considered foundational research, and does not target practical tasks that can be deployed outside of the research field. Thus we do not anticipate negative social impacts from this work.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation. generating fake profiles, surveillance), faireess considerations (e.g.. deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative socictal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA]
Question:
Justification: Our paper utilizes publicly domain datasets (not scraped), and poses no safety risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
34

---

Answer: [Yes]
Justification: We properly cite papers that introduce algorithms (such as LCA), datasets (such as MNIST, CIFARIO, van Hateren), and code (such LCA and ISTA).
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
If assets released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: New assets introduced in the paper consist of codebase which includes notebooks to replicate our experiments and analyses, and contains documentation.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
35

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with
human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
36


Reviews:
[{"review_id": "0dPeK32NZa", "reviewer": "Reviewer_2Wt3", "summary": "In this paper, the authors propose a VAE model, the Poisson-VAE ($\\mathcal{P}$-VAE), with a Poisson distributed prior and approximate posterior such that it works with Poisson distributed latents and demonstrate its sparse coding abilities on the van Hateren dataset and its efficacy on downstream classification tasks using the MNIST dataset.  The authors also present a Poisson reparametrization trick.", "questions": "- The authors list this as a limitation of their work later, but is it possible to discuss potential reasons for a large amortization gap of the P-VAE as compared to the LCA/ISTA?", "limitations": "The authors adequately discuss limitations of their work, and its potential societal impact.\n\nI would like to note another limitation of their work as not very illustrative experiments on downstream classification tasks (MNIST and CIFAR have been considered in the paper). The paper could benefit from experiments showing greater impact of sparse coding.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The authors combine predictive coding and Poisson-distributed latents to obtain a neat P-VAE objective emphasizing sparsity without additional design constraints-- hence, elegantly buying amortized sparse coding.\n- The authors also propose a new reparametrization trick, the Poisson reparametrization trick, which can potentially be more generally applicable.\n\nWriting:\n- The writing and presentation of the paper is very clear, making the derivation and rest of the math easily approachable.\n\nResults:\n- The authors seem to get very good results in terms of the values of the ELBOs obtained on MNIST as compared to pre-existing work. They also showcase linear separability of learnt features on MNIST.", "weaknesses": "These questions/remarks might also have arisen due to my lack of proper understanding, so I am willing to increase my score if these can be clarified:\n\n- How do we know that the approximate posterior is Poisson? \n- The temperature is controlling the sharpness of the thresholding-- what is the temperature being considered generally? How do we know that the cdf = 0.99999 across all these temperatures?\n\nExperiments:\n- How does the ELBO vary with the dimension of the overcomplete latent space considered?\n- Is the VQ-VAE a valid baseline? If so, it would be helpful if the authors can mention why it hasn't been considered as a baseline in the paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and thoughtful comments. We are also glad the reviewer finds our results \"neat\" and \"elegant\"!\n\n> How do we know that the approximate posterior is Poisson?\n\nDuring training, we use non-zero temperatures in our reparameterization algorithm. As a result, during training, the approximate posterior is a relaxed approximation to Poisson. In Fig. R1b we show samples drawn from our \"relaxed Poisson\" distribution at different temperatures. As $T \\rightarrow 0$, the sampled distribution converges toward the true Poisson distribution at $T = 0$.\n\nHowever, we want to emphasize that we used relaxed Poisson only during training. At validation time, we always set $T = 0$. Therefore, the approximate posterior is exactly Poisson at test time.\n\n> …what is the temperature being considered generally?\n\nDuring the first half of training, we anneal temperatures from a large initial value,  such as $T_\\mathrm{start} = 1.0$, to a small final value, $T_\\mathrm{final}$. In the paper, as well as Figs. R1c and d, we report results obtained using $T_\\mathrm{final} = 0.05$. Our new extensive experiments (Fig R1a) suggest this choice was reasonable.\n\n> How do we know that the cdf = 0.99999 across all these temperatures?\n\nThank you for this interesting question, which we overlooked before. We determine ```n_exp``` in Algorithm 1 using the largest posterior rate in a given batch, $r_\\mathrm{max}$. We observed that the distribution of rates in a batch is typically skewed and long-tailed. Consequently, the vast majority of the rates are much lower than $r_\\mathrm{max}$. Therefore, this is a very conservative way of choosing ```n_exp```, and even if the \"cdf = 0.99999\" condition is not met for certain temperatures, it will be so only for a vanishingly small subset of rates.\n\nWith that said, we explored this question empirically for temperatures encountered during training, using a few reasonable rate values. We found that this condition holds regardless of temperature (even more strongly for non-zero temperatures when the rate is small).\n\n> How does the ELBO vary with the dimension of the overcomplete latent space considered?\n\nWe investigated this across datasets (van Hateren and CIFAR), and encoder architectures (linear versus convolutional). We report the results in Fig. R1d. We found that for all convolutional encoder cases, ELBO improves as a function of latent dimensionality. However, for linear encoders, we observed that the van Hateren dataset started to overfit for $K > 512$, and it stagnated for the $\\mathrm{CIFAR}_{16 \\times 16}$ dataset. In conclusion, more expressive encoders can find nonlinear features, represented using additional latent dimensions, but simple linear encoders struggle to utilize additional dimensions.\n\n> Is the VQ-VAE a valid baseline?\n\nWe did not consider VQ-VAE as a baseline, because VQ-VAE does not optimize the ELBO loss, and thus, it is not technically a VAE. We think the naming is unfortunate and we will discuss this point in the final paper.\n\n> …discuss potential reasons for a large amortization gap of the P-VAE as compared to the LCA/ISTA?\n\nWe suspect both LCA and ISTA perform well because they are iterative algorithms. Our P-VAE, in its current format, uses a single forward pass to perform inference. Prior work by [Marino and colleagues](https://proceedings.mlr.press/v80/marino18a.html) has shown that iterative inference can significantly decrease the amortization gap. As a future work, we are interested in developing iterative versions of P-VAE, which we hope will close the amortization gap and beat the best LCA and ISTA fits.\n\n> …not very illustrative experiments on downstream classification tasks…\n\nWe agree with the reviewer here. We primarily used these downstream tasks to assess whether the geometry of representations is quantitatively different between P-VAE and alternative models. Our limited results in this area suggest they are different, which we plan to explore more rigorously later. The point of the current experiments was to establish there is indeed a difference in the geometry of representations, rather than fully exploring the differences and similarities, or highlighting particular applications of P-VAE.\n\n> The paper could benefit from experiments showing greater impact of sparse coding.\n\nWe agree that a limitation of our paper is that we did not highlight applications where a P-VAE would have a greater impact. As described in the general rebuttal, we will discuss several potential future directions in the final paper, and we believe these warrant more attention than we have room for in this paper.\n\nWe developed P-VAE with neuroscience applications in mind and one of the major advantages over other VAEs is that the latents can now be interpreted as neurons. That said, there are applications where sparse coding is applied, such as image processing and computational imaging, where P-VAE might shine. We have included a few of these potential directions in our general rebuttal, but still feel that they would require substantial experimental evaluation and should be left to future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "cUNrcaSaE4", "reviewer": "Reviewer_w8Jn", "summary": "Inspired by biological neurons, a new type of variational autoencoder, the Poisson variational autoencoder ($\\mathcal{P}$-VAE) is proposed. The $\\mathcal{P}$-VAE uses discrete latent states with Poisson priors, and learns sparse discrete representations of the data similar to sparse coding methods. The authors compare their method with classical sparse coding algorithms and obtain learned representations that are less prone to posterior collapse and are better for downstream classification tasks.", "questions": "- I am confused by the statement on lines 56-57: “facilitating linear separability of categories in a downstream classification task with a much better (5x) sample efficiency.” Isn’t it true that the increase in sample efficiency is for KNN classification but not in the linear separability? In that case this sentence is misleading.\n- What’s the significance of the global prior parameter $r$? Have you tried ablation studies where the posteriors are directly parameterized?\n- What’s the takeaway for the “$\\mathcal{P}$-VAE learns sparse representations” section? How should we view the amortization gap you mentioned on line 277?\n- There are multiple bold entries for some columns in tables 3 and 5, what do these mean?\n- Can you talk more about where you believe the $\\mathcal{P}$-VAE model should be applied and how it can be impactful to the scientific community?", "limitations": "The authors have adequately addressed the limitations of their work.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- This paper is very well-written and ideas are presented very clearly. The writing is clear and concise with little to no typos. The contributions are outlined and emphasized throughout the paper. Background to the work are explained thoroughly. The figure and algorithm box serve to communicate the main algorithm clearly to the reader. The color coding of inference and generative components in section 3 makes it easy to parse the equations and grok the main $\\mathcal{P}$-VAE algorithm.\n- The main arguments of the paper are supported well by experiments. The experiment section presents extensive analysis on the representation learning capabilities of the $\\mathcal{P}$-VAE, and compared it to many existing VAE models, including both continuous and discrete ones. In-depth discussions are made with respect to the Gabor-like quality of learned filters, avoiding posterior collapse, sparsity and effectiveness in downstream tasks.", "weaknesses": "My main concern for this work is in its limited impact. While to my knowledge it is true that applying the Poisson prior to discrete VAEs as a sparsity-inducing constraint is novel, I fail to see how this is fundamentally different that other discrete VAEs with regularization. Further to this point, the comparisons to VAE with the concrete distribution ($\\mathcal{C}$-VAE) seems relatively weak, as the latter has less dead neurons and also seems to achieve more sparsity for the same reconstruction performance (Figure 3a). The authors also commented that LCA models drastically outperform the $\\mathcal{P}$-VAE with the convolution encoder, making it questionable if one should use $\\mathcal{P}$-VAE for its sparsity.\n\nTo the authors’ credit, the paper does state that part of the appeal of the $\\mathcal{P}$-VAE is in its biological plausibility. Unfortunately I don’t think this point is expanded upon to a satisfactory degree in the text, leaving much to be desired. Personally I would love to see more technical discussion on the biological plausibility of the $\\mathcal{P}$-VAE, and concrete examples on the types of tasks and inquiries it unlocks for neuroscientific studies that similar models cannot solve adequately.\n\nDespite the above points, I still think that this is a solid paper with potential contributions to the computational neuroscience community.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and thoughtful comments.\n\n> My main concern for this work is in its limited impact.\n\nWe hope that we have addressed this concern in the general rebuttal. We developed P-VAE with neuroscience applications in mind, where we feel it will have a big impact. However, there are several other exciting future directions that we have discussed in the global rebuttal, which we will expand upon and include in the final paper.\n\n> …the comparisons to VAE with the concrete distribution (𝐶-VAE) seems relatively weak, as the latter has less dead neurons and also seems to achieve more sparsity for the same reconstruction performance.\n\nC-VAE is hard to compare to for a number of reasons which we'll discuss here. But first, we want to emphasize the P-VAE with a convolutional encoder has much better reconstruction performance than C-VAE, with a comparable sparsity level. That said, the sparsity comparison is hard to make because C-VAE is a one-hot categorical distribution, which means it has a fixed sparsity of $1-1/K$ (~99% when $K=512$). P-VAE can sweep out a tradeoff between sparsity and reconstruction (something akin to a rate-distortion curve). In contrast, the sparsity of C-VAE is solely determined by the latent dimensionality.\n\nFurther, Because C-VAE has a single KL value, the comparison between dead neurons is also not straightforward, which is why we excluded C-VAE from Table 2. For the C-VAE, we used the norm of basis elements as a measure of the dead neurons. We explain this in the caption of Fig. 5.\n\nWe will add text to discuss these points more explicitly in the final paper.\n\n> …The authors also commented that LCA models drastically outperform the 𝑃-VAE with the convolution encoder, making it questionable if one should use 𝑃-VAE for its sparsity.\n\nWe were surprised to see such a strong performance for LCA, a \"shallow\" model from 2008. We speculate this is because LCA performs iterative inference, whereas P-VAE, in its current form, uses a single forward pass for inference (amortized inference). This observation has inspired us to develop an iterative version of the P-VAE in subsequent work, which we hope will reduce the amortization gap and either match or surpass LCA performance.\n\nThat said, it is important to note that LCA necessarily has to work with a linear generative model ([Rozell et al., 2008](https://direct.mit.edu/neco/article-abstract/20/10/2526/7343/Sparse-Coding-via-Thresholding-and-Local?redirectedFrom=fulltext)), whereas P-VAE (like other VAEs) could have a nonlinear (deep) decoder. We primarily focused on linear decoders in this paper because of the close connection between the P-VAE loss and sparse coding. Two future directions could be to explore P-VAEs with more typical deep decoders or, as we mentioned above, with iterative inference, which we believe is a promising future direction for closing the amortization gap in VAEs more generally.\n\n> Personally I would love to see more technical discussion on the biological plausibility of the 𝑃-VAE, and concrete examples on the types of tasks and inquiries it unlocks for neuroscientific studies that similar models cannot solve adequately.\n\nThank you for this comment. We have included a discussion on neuroscience applications in the global rebuttal which attempts to address this point. We will include more discussion on this point in the final version of the paper.\n\n> I am confused by the statement on lines 56-57: “facilitating linear separability of categories in a downstream classification task with a much better (5x) sample efficiency.” Isn’t it true that the increase in sample efficiency is for KNN classification but not in the linear separability? In that case this sentence is misleading.\n\nYou are correct. That sentence should be broken into two separate points. We will clarify that the linear separability claim is supported by the shattering dim results, while the sample efficiency claim is supported by the KNN results.\n\n> What’s the significance of the global prior parameter 𝑟? Have you tried ablation studies where the posteriors are directly parameterized?\n\nWe did not consider an independently parameterized posterior, because our current implementation has two desirable features:\n\n1. a direct connection to predictive coding; and,\n2. a nice factorization of the KL term.\n\nHowever, this question led us to investigate the learned global prior parameters, $r$, for which we thank the reviewer. We found that the prior rates learn an efficient representation of the natural scene statistics. We describe this exciting result in the global rebuttal (Fig. R1e).\n\n> What’s the takeaway for the \"𝑃-VAE learns sparse representations\" section? How should we view the amortization gap you mentioned on line 277?\n\nThere are two main takeaways. First, P-VAE produces sparser representations than continuous VAE counterparts, with a comparable reconstruction performance. Second, there remains a substantial amortization gap between P-VAE (which uses amortized inference) and true sparse coding, which employs iterative inference. In our experiments, LCA converged in typically hundreds of iterations, whereas P-VAE does inference in one shot. Developing an iterative P-VAE and using it to close the amortization gap is our top priority for future work.\n\n> There are multiple bold entries for some columns in tables 3 and 5, what do these mean?\n\nBold indicates values that passed a significance threshold using statistical tests. We will clarify in the final paper that the bolded models perform similarly, such that their performance is statistically indistinguishable.\n\n> Can you talk more about where you believe the 𝑃-VAE model should be applied and how it can be impactful to the scientific community?\n\nYes. Multiple reviewers have requested this. We hope our global rebuttal addresses this adequately and will include a discussion on the significance of the work in the final paper.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "kZYrhprbum", "reviewer": "Reviewer_tb4S", "summary": "This work introduces VAEs with Poisson-distributed latent variables and a Poisson reparameterization for efficient training. This approach has theoretical and empirical connections with sparse coding and behaves more similarly to biological networks that rely on discrete spike counts.", "questions": "See weaknesses.", "limitations": "Some limitations are addressed, but overall the presentation of the experimental results doesn't seem too balanced.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "To the best of my knowledge, the method is original and well motivated. The paper is well written and generally clear. Though I cannot speak to the significance of this work from a neuroscience perspective, introducing Poisson latents is interesting in itself from a probabilistic models perspective. The experimental results are relatively extensive and there are several ablations.", "weaknesses": "- Posterior collapse in VAEs (in the sense of \"some latent dimensions are not used\") is not necessarily an issue per se. If this is an issue in the specific scenarios considered in this paper, I think this should be clarified.\n- For the downstream classification tasks, I appreciate the experiments with different latent space sizes. However, using KNN as downstream classifier is a quite specific choice, and I would argue a simple linear probe is more common in the literature and seems like an intentionally missing baseline.\n- While my expertise is in machine learning rather than neuroscience, I appreciate the value of ANNs that mimic biological networks. However, from an ML perspective, it would be helpful to discuss any potential challenges or limitations of this method (besides the limitations currently mentioned). The paper currently highlights the advantages of using Poisson latents in various aspects, which may seem overly optimistic. Since the goal is to bridge ANNs with biological systems rather than outperform benchmarks, a balanced discussion of both strengths and weaknesses would provide a more comprehensive perspective and would IMHO add value to this work.\n- Since the main motivation is to learn brain-like representations, are there any datasets or experiments where this comparison with biological representations could actually be done? This might be the biggest missing part right now. If this is not feasible at all, it would also be fine to at least include a discussion for readers without a neuroscience background.\n- More broadly, connections to neuroscience, as well as jargon, should be made more explicit/clear for pure ML people (e.g. \"Gabor-like feature selectivity\")\n- Notation: at some point the authors start using $\\boldsymbol{r}$ and $\\boldsymbol{\\delta}$ without introducing them in the main text (though the notation is better explained in Appendix). I would even recommend using $\\boldsymbol{\\delta}_r$ instead of $\\boldsymbol{\\delta r}$ which looks like a product, and is especially confusing when writing $\\boldsymbol{r \\delta r}$.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and thoughtful comments.\n\n> Posterior collapse in VAEs (in the sense of \"some latent dimensions are not used\") is not necessarily an issue per se. If this is an issue in the specific scenarios considered in this paper, I think this should be clarified.\n\nWe agree that the concern over posterior collapse depends on the application and we will clarify this in the final text. Nevertheless, addressing the posterior collapse issue remains an active area of research ([He et al., 2019](https://openreview.net/forum?id=rylDfnCqF7), [Lucas et al., 2019](https://proceedings.neurips.cc/paper/2019/hash/7e3315fe390974fcf25e44a9445bd821-Abstract.html), [Razavi et al., 2019](https://openreview.net/forum?id=BJe0Gn0cY7), [Menon et al., 2022](https://openreview.net/forum?id=SrgIkwLjql9), just to name a few).\n\nIn our review of existing literature at the intersection of VAEs and sparse coding, we found that posterior collapse has been a recurring issue there as well. Often, sparse coding results are evaluated based on what features are learned by the dictionary elements. For example, both [Csikor et al., 2023](https://www.biorxiv.org/content/10.1101/2023.11.29.569262v2) and [Geadah et al., 2024](https://www.biorxiv.org/content/10.1101/399246v3), used Laplace-distributed latents in VAEs, aiming to learn Gabor-like feature selectivity. However, they did not show the full set of dictionary elements. Our experiments revealed that ~80% of latent dimensions collapse for the L-VAE, resulting in noisy and therefore useless basis elements. This significantly deviates from the classical sparse coding results.\n\nOne implication of our results for future work is that L-VAE should be avoided in favor of P-VAE if the goal is to develop a diverse set of dictionary elements that are reminiscent of classical sparse coding. We will add a discussion of why posterior collapse is relevant in our case to the final paper.\n\n> …using KNN as downstream classifier is a quite specific choice, and I would argue a simple linear probe is more common in the literature and seems like an intentionally missing baseline.\n\nWe reasoned that KNN is a good choice for evaluating the learned representations because it is a non-parametric method and its performance is directly influenced by the geometry of representations—which is what we were interested in evaluating.\n\nFor a more complete model evaluation, we performed simple logistic regression classification as part of the rebuttal results (Fig. R1c). We found that P-VAE achieves the best overall performance for latent dimensionally of $K = 100$. But for $K = 10$, both L-VAE and G-VAE outperform P-VAE. We plan to include these new results in Tables 3 and 5, alongside the KNN and shattering dim results.\n\n> …a balanced discussion of both strengths and weaknesses would provide a more comprehensive perspective and would IMHO add value to this work.\n\nWe agree. However, we thought our presentation was fairly balanced. For example, we show in Fig. 6 that both the quality of generated samples, and reconstruction performance, are higher for continuous VAEs compared to discrete ones. We also highlight a large amortization gap that remains between P-VAE and sparse coding models. That said, within discrete models, P-VAE still performs better than C-VAE and there are advantages to P-VAE, which we highlight throughout the paper. We will make sure that our discussion of strengths and limitations is more clear in the final paper.\n\n> Since the main motivation is to learn brain-like representations, are there any datasets or experiments where this comparison with biological representations could actually be done? This might be the biggest missing part right now.\n\nWe agree that a comparison with biological representations is an obvious next step and there are several datasets available for this. In the present work, we demonstrated the connection between Poisson VAEs and sparse coding and evaluated the representations learned by continuous and discrete VAEs. Previous work has evaluated VAEs for predicting neural activity ([Vafaii et al., 2023](https://openreview.net/forum?id=1wOkHN9JK8)), and has found that benchmarking, as is standardly done, does not discriminate well between different models. A more thorough evaluation of the learned representations is necessary and there isn’t room here to do that justice, although, this is an immediate plan of ours. Since both neurons in the brains and P-VAE encode information in firing rates, we believe the P-VAE will learn brain-like representations.\n\nWe have discussed some applications in the global rebuttal and will integrate some of the points made here into the discussion in the final paper.\n\n> …connections to neuroscience, as well as jargon, should be made more explicit/clear for pure ML people…\n\nThank you for pointing this out. We will amend the language to enhance clarity.\n\n> Notation: at some point the authors start using 𝑟 and 𝛿 without introducing them in the main text…\n\nWe have $\\log\\delta r \\in \\mathbb{R}^{K}$, such that $\\delta r \\in \\mathbb{R}_+^{K}$. Additionally, by $r \\delta r$, we mean $r \\odot \\delta r$, where $\\odot$ is the Hadamard or element-wise multiplication. We will clarify this in the text.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ubebyzFgy6", "reviewer": "Reviewer_gEE3", "summary": "The paper proposes a variation on variational auto-encoders with a Poisson distribution over the latents. To make the model differentiable they use a differentiable sampling of the latent variables where the indicator function is replaced with a continuous approximation. They take inspiration from biological neural networks to develop the model and make connections to sparse coding to argue that their model can be used to understand sensory processing in the brain. Their experiments examine various key aspects of their model and compare to other significant VAE baselines. They provide evidence that their model learns sparse basis vectors for image datasets when compared to popular sparse coding algorithms and their model learns latent representations that perform better than baselines on downstream classification tasks.", "questions": "- The outputs of the decoders in Fig. 1 should have noise added to them to illustrate that the likelihood isn't a delta distribution.\n- Fig 1b is a bit confusing as the sampling of the latents isn't shown whereas it was shown in Fig 1a.\n- Maybe Table 4 should be added to the main paper as it is the first experiment that is referred to and the reparameterization gradient is a crucial aspect of the model.\n- Can a discussion be added on the slight differences between the P-VAE basis vectors and those of LCA and ISTA in Figure 2? Also what is the gold standard here? Is there a quantitative metric or can this only be determined qualitatively.", "limitations": "The authors have sufficiently discussed limitations.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- A Poisson variational auto-encoder is a somewhat novel though not as original contribution to the VAE zoo that which can clearly be situated among other VAE models. The connection to sparse coding and brain representations also well motivates the utility of the model for studying representations in the brain.\n- The paper is very well written and was, for the most part, easy to follow. The model is well-explained and the figures served well to aid in understanding\n- The experimental evaluation was quite thorough with the results providing strong evidence to support the paper's claims. The systematic analysis of various of the model's aspects (sparseness of representations and utility for downstream tasks) was well done and made clear to the reviewer what the capabilities and limitations of the model are.", "weaknesses": "- A significant weakness is a lack of study of the effect of the temperature parameter in the Poisson re-parameterization. How does this need to be set and what consequences for different temperatures? The z variables on line 6 of Algorithm 1 aren't integers and so how does  changing temperature affect the gradient? Appendix A.3 does talk about how the temperature is annealed during training but motivation for this approach isn't provided even though the temperature is a key aspect of their method.\n- Also, a lack of discussion of the noise level that needs to be set for the likelihood $p(x | z)$ makes it unclear as to how the model can be re-configured to other datasets and what the consequences would be. What if we believed that the noise level in the data was much less than one (as is assumed with the Gaussian likelihood conditioned on the decoder output)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and thoughtful comments.\n\n> A significant weakness is a lack of study of the effect of the temperature parameter in the Poisson re-parameterization.\n\nWe agree. Please see our new rebuttal results, where we performed extensive experiments to address this point. We will include these new results in the paper.\n\n> The z variables on line 6 of Algorithm 1 aren't integers…\n\nGood point. This is true during training, but not during test time. We would like to emphasize that during validation, we always set $T = 0$, which results in integer samples drawn from a true Poisson distribution. We have included a discussion of this point and a figure in the global rebuttal and will integrate that into the final paper.\n\nRelated to this point, in Fig. R1b we plot the distribution of samples at various temperatures, which shows the samples are indeed not integers for $T > 0$, but they approach integer values when $T \\rightarrow 0$. Therefore, even during training, we can obtain \"almost-integer\" values if we anneal down to low temperatures. Interestingly, this does affect performance, and having these subtle non-integer values during training even improves model performance at test time with hard integers (Fig. R1a).\n\n> …how does changing temperature affect the gradient?\n\nWe weren't sure exactly how to address this. We find that the model performance is robust for a range of values of $T_\\mathrm{final} < 0.1$ (Fig. R1a), but we did not directly evaluate the effects on gradients. We are happy to address this more directly if you can suggest an analysis. Otherwise, we're inclined to interpret the performance results (Fig. R1a) instead of analyzing the gradients directly.\n\n> …a lack of discussion of the noise level that needs to be set for the likelihood $p(x \\vert z)$...\n\nWe construct our likelihood function, $p(x \\vert z)  = \\mathcal{N}(x; \\mathrm{dec}(z), \\sigma^2)$, by learning the Gaussian mean but using the same fixed variance for every pixel. Specifically, we chose $\\sigma = 1/\\sqrt{2}$, such that $-\\log p(x \\vert z) = ||x - \\mathrm{dec}(z)||_2^2$. This choice of fixed variance is fairly standard in the VAE literature; however, others have drawn attention to the \"disheartening\" limitations of this approach ([Arvantidis et al., 2018](https://openreview.net/forum?id=SJzRZ-WCZ)), and we share their sentiment. We will add these comments to the paper to further clarify our choices.\n\n> The outputs of the decoders in Fig. 1 should have noise added to them to illustrate that the likelihood isn't a delta distribution.\n\nThank you for pointing this out. We will edit the figure to highlight this.\n\n> Fig 1b is a bit confusing as the sampling of the latents isn't shown whereas it was shown in Fig 1a.\n\nThank you for pointing this out. We will edit the figure to make it more consistent.\n\n> Maybe Table 4 should be added to the main paper…\n\nWe plan to include a reduced version of Table 4 (or Fig. 4), along with Fig. R1b, in the main paper. The goal is to emphasize that, even though the approximate posterior is \"relaxed Poisson,\" the final performance is almost as good as using exact gradients for training.\n\n> Can a discussion be added on the slight differences between the P-VAE basis vectors and those of LCA and ISTA in Figure 2? Also what is the gold standard here? Is there a quantitative metric or can this only be determined qualitatively.\n\nThere is not really a gold standard here and, surprisingly, most of the sparse coding literature evaluates dictionaries qualitatively. One way to quantify the differences and similarities between pairs of dictionaries is by fitting parametric Gabor functions. These Gabor fits allow extracting and studying parameters such as orientation, spatial frequency, size, and location for each basis element. We briefly explored this direction for the rebuttal results. Please see Fig. R1e.\n\nWe believe this line of inquiry deserves a more thorough investigation. We are happy to add a comparison of the distributions of Gabor fits to the different dictionaries if the reviewer thinks that would be helpful. However, we are inclined to leave this direction as future work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their time and insightful feedback. We believe addressing the reviewers' comments will substantially improve our paper. We plan to include these changes in the form of two major components: **(i)** further discussion of the work significance; and, **(ii)** additional results, reported as a figure here (Fig. R1).

# (i) Work significance and contributions
## Utility for neuroscience
The neuroscience community increasingly uses ANNs to understand what biological neurons are selective for and why. ANNs have several advantages over real brains: all connections and activations are available to the investigator, in silico perturbations are inexpensive, and they do not require the use of animals.

Like biological neurons, the P-VAE generates spikes; therefore, its latents can be treated like neurons. This offers advantages over continuous unconstrained models such as L-VAE or G-VAE. Here, we focus on two examples, leaving a more comprehensive discussion to the final paper.

- **Example 1:** P-VAE supports direct comparison to causal perturbations in brains. Perturbation experiments selectively stimulate or silence neurons to assess the causal role of a group of neurons on perception. These types of experiments cannot be trivially compared to unconstrained VAEs, for which "stimulation" is complicated by the fact that the latents are signed. In contrast, the P-VAE can be readily used for designing and conducting in silico perturbation experiments, enabling an exciting potential transfer of insights to the in vivo setting.
- **Example 2:** "Maximally Exciting Inputs" (MEI; [Walker et al., 2019](https://www.nature.com/articles/s41593-019-0517-x)), which manipulate inputs following the gradients of a feed-forward network fit to biological neurons, have been used to understand what biological neurons are selective for. The concept of MEI requires that neuron activations be characterized by being more or less "excitable." Once again, this concept is readily applicable to P-VAE latents, but not the unconstrained VAEs.

There are just two examples out of many. We will expand upon this point in the final paper, mentioning more concrete examples of the application potential of the P-VAE and how it can help advance neuroscience research.

## Mechanistic interpretability
Sparse autoencoders (SAEs) have become popular for mechanistic interpretability ([Anthropic](https://transformer-circuits.pub/2023/monosemantic-features), [OpenAI](https://arxiv.org/abs/2406.04093)). Although not an initial motivation for our work, we have inadvertently built a probabilistic version of SAEs. We're excited about a hierarchical extension of P-VAE applied to both images and LLM activations to test if this approach extracts hierarchically organized semantic concepts.

## Reparameterization trick
We hope our reparameterization trick finds applications beyond VAEs, for example, in spiking neural networks (SNN). One of our new results shows that the "surrogate gradients" method—utilized heavily in the SNN literature—may be improved by relaxing the hard forward during training (see below).

## Hardware implementation
A key advantage of our architecture is its ability to learn discrete, sparse representations. The integer P-VAE representations eliminate the need for post hoc quantization, which is crucial for hardware implementation of models with float activations. This sparsity enhances memory efficiency and lowers energy use, highlighting P-VAE’s potential as a vision model implemented directly on hardware for robotics.

# (ii) New results (Figure R1)
## Temperature and performance (Fig. R1a)
Motivated by comments from Rev. gEE3, we performed additional experiments to quantify the effect of temperature (T) on the final model performance. Following standard practice ([Jang et al., 2017](https://openreview.net/forum?id=rkE3y85ee)), we annealed T from a large value ($T_\mathrm{start} = 1.0$) to a small value ($T_\mathrm{final} = 0.05$ in the main paper) during the first half of training. In Fig. R1a, we explore the effect of changing $T_\mathrm{final}$ on the van Hateren dataset, using two architectures (linear vs. convolutional encoders, linear decoder), and two annealing schedules (linear vs. exponential; Fig. R1e inset). We find values of $T_\mathrm{final} \leq 0.1$ work well, and both annealing schedules work well.

Importantly, all results were obtained using $T = 0$ during test time. We also experimented with the option of using a "hard forward" training scheme once the annealing is done (i.e., the last half of training), where we use non-zero temperatures only for the backward pass. This practice is known as "surrogate gradients." Somewhat surprisingly, we found that the surrogate gradients severely underperformed our "relaxed Poisson" approach. We anticipate this result will be highly interesting to the spiking neural network community, who rely mostly on surrogate gradients to train their networks. We plan to include this figure in the appendix and highlight the main takeaways in the main text.

## Natural image statistics learned in the prior rates (Fig. R1e)
Thanks to comments from Revs. gEE3 and w8Jn, we examined the properties of the P-VAE learned dictionary elements in conjunction with the global prior rates, r. See Fig. 1e and its caption.

We found that P-VAE prior rates are consistent with principles of efficient coding, which states brains should assign minimal neural resources to statistically dominant elements of the natural environment. Specifically, we found prior rates were lower for *cardinal* orientations, which are more common in natural image patches. This result mirrors biological brains and is another demonstration of the potential of P-VAE. We plan to explore this research direction more systematically in future work.

# Conclusion
P-VAE has shown promising results, opening up many exciting venues for future exploration at the intersection of machine learning and neuroscience.

Meta-review: There is very strong consensus among the reviewers that this paper should be accepted to NeurIPS. The Poisson VAE represents a reasonably novel model, and is buttressed by both the biological plausibility of the model and a strong experimental section. Ideas from this paper appear to be generalizable to other domains, further increasing the impact.

The authors also provided a very strong rebuttal to the various points of contention from the reviewers, eventually resulting in the clear consensus of acceptance.
Decision: Accept (spotlight)