Poisson Variational Autoencoder

Authors: hadi  vafaii
dekel  galor
jacob l yates
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 06:23:40

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.

Main paper: 
 introduction: 
 The study of artificial neural networks (ANN) and neuroscience has always been closely linked, driving advancements in both fields [5] [6] [7] [8] [9] [10] . Despite the close proximity of the two fields, most ANN models deviate substantially from biological brains [11, 12] . A major challenge is designing models that not only perform well computationally but also exhibit "brain-like" structure and function. This is seen both as a goal for improving ANNs [13] [14] [15] , and better understanding biological brains [8, 9, [16] [17] [18] [19] , which has recently been referred to as the neuroconnectionist research programme [20] .Drawing from neuroscience, a major guiding idea is that perception is a process of inference [21, 22] , where the brain constructs a representation of the external world by inferring the causes of sensory inputs [23] [24] [25] [26] . This concept is mirrored in "generative AI" where models learn the generative process underlying their inputs [27] [28] [29] . However, in this vein, there is a tension between small well-understood models that are directly inspired by cortex, such as sparse coding [3] and predictive coding [30] , and deep generative models that perform well [31] [32] [33] [34] .The variational autoencoder (VAE; [35, 36] ) model family is a promising candidate for neuroconnectionist 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 as diffusion models, understood as special cases of hierarchical VAEs [37] [38] [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 in the way they encode information. Biological neurons fire all-or-none action potentials [44] , and are thought to represent information via firing rate [45] [46] [47] [48] [49] . These firing rates must be positive and generate discrete "spike" counts, which exhibit conditionally Poisson-like statistics in small counting windows [49] [50] [51] . 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. 1 ). 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. 2a ; [30, 53] ). We introduce a reparameterization trick for Poisson samples (Algorithm 1), and derive the evidence lower bound (ELBO) 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. ( 3)) naturally results in a metabolic cost term that penalizes high firing rates, such that P-VAE with a linear decoder implements amortized sparse coding (Fig. 2b ). 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 1 ), 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 (5×) sample efficiency.We evaluate these results on two natural image datasets and MNIST. The P-VAE paves the way for the future development of interpretable hierarchical models that perform "brain-like" 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 of the world results from the unconscious inference over the causes of the senses. In other words, the brain learns a generative model of the sensory inputs. This has led to fruitful theoretical work in neuroscience [23, [54] [55] [56] and machine learning [57, 58] , including VAEs [52] . See Marino [10] for a 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] .Predictive coding [30, 61, 62] postulates that the brain generates a 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 to capture a wide range of phenomena in biological neurons across the visual system [64, 65] . More broadly, prediction in time has emerged 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 [47, 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 [69] [70] [71] [72] [73] [74] [75] [76] [77] [78] [79] . SC is formalized with a generative model where neural activations z are sampled from a sparsity-inducing prior, z ∼ p(z), and the input image x is reconstructed as a linear combination of basis vectors Φ, plus additive Gaussian noise, x = Φz + ε. The SC loss is as follows:EQUATIONCommonly used algorithms for sparse coding include the locally competitive algorithm (LCA; [80] ), which is a biologically plausible algorithm to optimize eq. ( 1), and iterative shrinkage-thresholding algorithm (ISTA; [81, 82] ), which has shown robust performance in learning sparse codes given a fixed dictionary Φ.VAEs define a probabilistic generative model p(x, z), where x denotes the observed data and z are some latent variables. The generative process samples z from a prior distribution p(z) and then generates the observed data x from the conditional distribution p θ (x|z), also known as the "decoder". The "encoder", q ϕ (z|x), 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 A for the full set of derivations). The ELBO is given by:EQUATIONThe 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 Gaussians are typically used: q = N (z; µ(x), σ 2 (x)) and p = N (z; 0, 1). The likelihood, p θ (x|z), is also typically modeled as a Gaussian conditioned on a parameterized neural network dec θ (z).Amortized inference in VAEs. A major contribution of VAEs is the idea of amortizing inference over the latents z with a black box ANN [83, 84] . "Amortized" inference borrows a term from finance to capture the idea of spreading out costs-here, the cost of performing inference over multiple samples. In amortized inference, a neural network learns (during training) how to map a 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, 87] , VAEs have been found to learn representations that are predictive of single-neuron activity in both the ventral [1] 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 [1, 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] . times ← cumsum(∆t, dim=0) ▷ compute arrival times, same shape as ∆t 5:indicator ← sigmoid return z 8: end procedure VAEs connection to sparse coding. Previous work has attempted to connect sparse coding and VAEs directly [91] [92] [93] , with each approaching the problem differently. Geadah et al. [91] introduced sparsity-inducing priors (such as Laplace or Cauchy) and a linear decoder with an overcomplete latent space. Tonolini et al. [92] introduced a spike and slab prior into a modified ELBO, and Xiao et al. [93] added a sparse coding layer learned by ISTA to the latent space of a VQ-VAE. 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 [91, 92] . Because the exponential distribution is trivially reparameterized [35] , and PyTorch contains an implementation [97] , we need only to approximate the hard threshold for comparing cumulative wait times with the window size. We accomplish this by replacing the indicator function with a sigmoid as in refs. [89, 98] . Algorithm 1 demonstrates the steps: Given a matrix of rates λ, sample n_exp wait times t 1 , t 2 , ...t n_exp for each element of λ by sampling from an exponential distribution with mean 1/λ(i). We then calculate the cumulative event times S(n_exp) = n_exp j=1 t j , pass them through a sigmoid σ( 1-S temperature ), and sum over samples to get event counts, z. The temperature controls the sharpness of the thresholding. We adaptively scale the number of samples, n_exp, by keeping track of the maximum rate in each batch, λ max , and then use the inverse cumulative density function (cdf) for Poisson to find the number of samples, n_exp, such that cdf(n_exp; λ max ) = 0.99999.At non-zero temperatures, our parameterization algorithm provides a continuous relaxation of the Poisson distribution. Figure 3 shows histograms of samples drawn using Algorithm 1 for rate λ = 1 and temperatures T = 1.0, 0.1, 0.01, and 0. The latter case (T = 0, true Poisson) is equivalent to torch.poisson().P-VAE architecture. The architecture of P-VAE captures the interactions between feedforward and feedback connections that are present in all visual cortical areas [99] . Feedforward areas carry sensory information and feedback connections are thought to carry modulatory signals such as attention [53] or prediction [30] , which interact multiplicatively with feedforward inputs [53, 100] .P-VAE embodies this idea by having the posterior rates depend on the prior, such that r prior = r and r post. = r ⊙ δr(x), where ⊙ is the Hadamard (element-wise) product. The prior rates, r ∈ R K , are learnable parameters that capture expectations about the statistics of the input. The encoder outputs, δr(x) ∈ R K , 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 a posterior, q(z|x) = Pois(z; r ⊙ δr(x)), and prior, p(z) = Pois(z; r), where z is the spike count variable. Notably, this multiplicative relationship is maximally general, as any pair of positive variables, r prior , and r post. can be expressed as a base variable, r := r prior , multiplied by their relative ratio, δr := r post. /r. The general model architecture is shown in Fig. 2a . P-VAE loss function. For a comprehensive derivation of the P-VAE objective, see appendix A. Here, we report the final result:EQUATIONwhere dec(•) is the decoder neural network, and f (y) := 1 -y + y log y (see supplementary Fig. 6 ).P-VAE relationship to sparse coding. The KL term in eq. (3) penalizes firing rates. Both r and δr are positive by definition, and f (y) ≥ 0, strongly resembling the sparsity penalty in Olshausen and Field [3] . To make this connection more explicit, we make two additional assumptions (Fig. 2b ):1. The decoder is a linear generative model:x = Φz, with x ∈ R M and Φ ∈ R M ×K . 2.The latent space is overcomplete: K > M .Because both E z∼Pois(z;λ) [z i ] and E z∼Pois(z;λ) [z i z j ] have closed-form solutions (eq. ( 21)), the reconstruction term in eq. ( 3) can be computed analytically for a linear decoder, resulting in:EQUATIONwhere λ = r⊙δr(x) are the posterior firing rates, f (y) is defined as above, and β is a hyperparameter that scales the contribution of the KL term, and changes the sparsity penalty for the P-VAE. Gaussian VAE (G-VAE; [35, 36] ) Laplace VAE (L-VAE; [40, 91] )The relationship between the linear P-VAE loss (eq. ( 4)) and the sparse coding loss (eq. ( 1)) 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, and the only additional assumption was an overcomplete linear generative model. The inference is accomplished using a parameterized feed-forward neural network, δr(x), thus, it is amortized. We call this specific case of P-VAE "Amortized Sparse Coding" (Fig. 2b ).Note that a closed-form derivation of the reconstruction term is possible for any VAE with a linear decoder and a generating distribution that has a mean and variance (see eq. ( 22)).This closed-form expression of the loss given a 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. 
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 selfcontained in the paper (section 3) and supplement (appendix A).Guidelines:• The answer NA means that the paper does not include theoretical results.• All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced. • All assumptions should be clearly stated or referenced in the statement of any theorems.• The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a 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. 
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 ( 22)) 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-VAE with a 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 shrinkage-thresholding algorithm (ISTA; [81, 82] ) to see if P-VAE reproduces their results. Third, we test the P-VAE in a generic representation learning context and evaluate the geometry of learned representations for downstream tasks. For these experiments, both the encoder and decoder's architecture is a ResNet (see appendix B for full architecture and training details).Alternative models. We compare P-VAE to both discrete and continuous VAEs (Table 1 ). 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-Softmax trick [89, 98] . We use PyTorch's implementation which is based on Maddison et al. [98] .Finally, we test models where latents are Gaussian passed through an activation function before passing to the decoder. We call these models G-VAE +act , where act ∈ {relu, exp}, but they capture 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 [101] .Datasets. For sparse coding results, we use 101 natural images from the van Hateren dataset [102] . We tile the images to extract 16 × 16 patches and apply whitening and contrast normalization, as is typically done in sparse coding literature [3, 103] . To test the generalizability of our sparse coding results, we repeat these steps on CIFAR10 [104] , a dataset we call CIFAR 16×16 . For the general representation learning results, we use MNIST. See appendix B for additional details.Statistical tests. In the VAE literature, it is known that random seeds can have a large effect compared to architecture or regularization [105] . Therefore, we train each configuration using 5 different random initializations. We report 99% confidence intervals throughout, and perform paired t-tests, reporting significance for p < 0.01 (FDR corrected using the Benjamini-Hochberg method). Training details. Because we considered a variety of architectures, training time is variable. 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, 82] and LCA [80] with 270 hyperparameter combinations each. See appendix B for more details. Training all models took roughly a week on 8 RTX 6000 Ada GPUs.Evaluating Poisson reparameterization. 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 as well as Categorical and Laplace VAEs. Table 4 shows the results. Monte-Carlo sampling with Poisson reparameterization closely matches exact inference just like established methods for Gaussian and Laplace. In contrast, the straight-through (ST; [106] ) estimator performs poorly (see also supplementary Fig. 7 ).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] . Inspecting the dictionaries learned by different models demonstrates this is not trivial (Fig. 4 ). As expected from theoretical results [4] , G-VAE (top left) learn probabilistic PCA, but with many noisy elements. As demonstrated previously [40, 91] , L-VAE (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 priors [40, 91] . In contrast, P-VAE (top middle) learns Gabor-like filters that cover space, orientation, and spatial 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. 4 is the sheer number of noisy basis vectors for both continuous VAEs (G-VAE, L-VAE). We suspected this reflected dead neurons with vanishing KL, which is indicative of a collapsed latent dimension that's no longer encoding information. To quantify this, we binned the distribution of KL values and thresholded the resulting distribution at discontinuous points (see supplemental Fig. 8 ). Table 2 shows the results of this analysis for all VAEs 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-VAE had ∼ 2% dead neurons compared to ∼ 80% for G-VAE and L-VAE. Having a more expressive encoder slightly increases this percentage, but a dramatic difference between P-VAE and continuous VAEs (G-VAE, L-VAE) persists.The P-VAE learns sparse representations. To quantify whether P-VAE learns sparse representations, 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: EQUATION)where N is the number of images, and z ij is sampled from the posterior for the i-th image. Intuitively, s j = 1 whenever neuron j responds to a single stimulus out of the entire set (highly selective). In contrast, s j = 0 whenever the neuron responds equally well to all stimuli indiscriminately. Fig. 5a shows the reconstruction performance (MSE) compared to lifetime sparsity (s, eq. ( 5)) for all VAEs. The G-VAE finds good reconstructions (MSE = 71.49) but with low sparsity (s = 0.37).The P-VAE finds much sparser solutions (s = 0.94), but at a cost of reconstruction quality (MSE = 102.68). Because the P-VAE KL term explicitly penalizes rate (eq. ( 3)), we explored different β values for P-VAE with a linear encoder and decoder (Fig. 5a 1)). The curve is a sigmoid fit.Does P-VAE match the performance of traditional sparse coding trained with LCA or ISTA? Figure 5b compares P-VAE to sparse coding models that were trained using a wide range of hyperparameters, and the best models were selected for each class (appendix B). P-VAE 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 P-VAE (Fig. 5b curve/solid points). We explored a range of hyperparameters to determine whether the MSE improved for similar sparsity levels. Indeed, LCA inference using P-VAE dictionary, was able to nearly match the performance of sparse coding LCA for similar levels of sparsity. This confirms our hypothesis that a large amortization gap remains for the specific encoder architectures we tested, highlighting the need for improved inference algorithms/architectures [108] . The P-VAE is more sample efficient in downstream tasks. We test the unsupervised learned representations on a downstream classification task. We first trained all VAE models with K = 10 dimensional latent space and a convolutional encoder and decoder on the MNIST training set (see supplementary Fig. 9 for the quality of generated samples and reconstruction performance). Next, we used the trained encoders to extract representations and evaluated their performance on classifying MNIST digits. We define VAE representation to be the output of the encoder. Following conventions in the VAE literature [105] , we used the mean vectors µ for continuous VAEs (G-VAE, L-VAE). For the P-VAE, we used log δr, and for the C-VAE, we used logits.We split the MNIST validation set into two 5,000 sample sets, used as train/test sets for this task. We train a K-nearest neighbors (KNN) classifier with a varying number of limited supervised samples (N = 200, 1000, 5000) drawn without replacement from the first set (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 [109] . 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 3 ). By this measure, P-VAE is 5× more sample efficient. But from Alleman et al. [110] , 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 posterior. This biological constraint improved G-VAE, but it still underperformed P-VAE (Table 3 ). We also found this result held for higher dimensional latent spaces (supplementary Table 5 ).The 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-VAE compared to other VAE families (Table 3 ). To test this more explicitly, we calculated the "shattering dimensionality" of the latent space [111] [112] [113] . Shattering dim measures the average accuracy over all possible pairwise classification tasks. This is called "shattering" because if the model shatters data points around into a high dimensional space, they will become more linearly separable. For MNIST with 10 classes, there are 10 5 = 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 3 shows 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-VAE strongly outperformed alternative models (Table 5 ).Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it 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. ( 3)), architecture and training details (appendix B), and pseudo-code for Poisson reparameterized sampling (Algorithm 1). In addition, we intend to release all code and data needed for replicating our work. 
conclusions: 
 In this paper, we describe the P-VAE, which performs posterior inference encoded in discrete spike counts. We introduce a Poisson reparameterization trick and derive the ELBO for Poisson-distributed VAEs. The P-VAE loss results in a KL term that penalizes firing rates like sparse coding [3] . We show that P-VAE with a linear generative model reduces to amortized sparse coding.When trained on natural image patches, P-VAE with a linear decoder learns sparse solutions with Gabor-like basis vectors resembling sparse coding, both in the form of the learned basis as well as the lifetime sparsity of the latents. We evaluated the representations on downstream classification tasks and found that P-VAE encodes its inputs in a higher dimensional space that enabled good linear separability between classes.Limitations and future directions. 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] . Extending P-VAE to hierarchical architectures [33] will make the latents conditionally Poisson, but not marginally Poisson (as they are modulated by top-down rates). Further extensions could implement doubly-stochastic spike generation [51, 114] . A second limitation is the amortization gap we observed between our current implementation of P-VAE and traditional sparse coding. This could likely be closed with more expressive encoders [115] or through iterative inference [116] , but it is an open area of research [108] .Overall, the P-VAE is a promising step towards learning brain-like representations in deep hierarchical generative models. 
code & data: 
 Our code, data, and model checkpoints are available here: https://github.com/hadivafaii/PoissonVAE. 
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 Graduate Research Fellowship Program under Grant No. DGE-1752814 (DG). Any opinions, findings, conclusions, or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation. We thank our anonymous reviewers for their helpful comments, and the developers of the software packages used in this project, including PyTorch [97] , NumPy [117] , SciPy [118] , scikit-learn [119] , pandas [120] , matplotlib [121] , and seaborn [122] .Gaussian. Let q(z|x) = N (z; µ q (x), σ q (x)) and p(z) = N (z; µ p , σ p ), where the mean and variance are either outputs of the encoder network (red) or parameters of the decoder network (blue). Now, let 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:EQUATIONWith these modifications, the Gaussians KL term becomes:EQUATIONIn standard Gaussian VAEs, the prior has no learnable parameter. Instead, we have µ → 0 and σ → 1. Therefore, the final form of the KL term for a standard Gaussian VAE is:EQUATION)We observe that the KL term vanishes when δµ → 0 and δσ → 1. This happens whenever no new information is propagated through the encoder, a phenomenon known as posterior collapse.Other than this trivial observation, eq. ( 12) does not really lend itself to interpretation. In contrast, will show below that a Poisson parameterization of VAEs leads to a much more interpretable outcome for the KL term.Poisson. Now suppose q(z|x) = Pois(z; rδr(x)), and p(z) = Pois(z; r), where z is literally the spike count of a single latent dimension-or 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 z, given rate λ ∈ R >0 , is given by:EQUATION)Plug this expressions into the KL divergence definition to get: EQUATIONwhere we have define f (y) := 1 -y + y log y.To examine the behavior of the Poisson KL term, we assume δr = 1 + ϵ, where ϵ ≪ 1, then Taylor expand f . Calculating the first and second derivatives of f (y) = 1 -y + y log y gives f ′ (y) = log y and f ′′ (y) = 1/y. Thus:EQUATIONPlug this back into eq. ( 14) to get:EQUATION)For small deviations ϵ, the KL term simplifies to the product of the prior firing rate, r, and ϵ 2 . See Fig. 6 for a visualization of the full function, f (δr) = 1 -δr + δr log δr, along with its quadratic approximation near δr = 1. 14). Right, quadratic approximation of f from eq. ( 15).In general, there are two ways to minimize the KL term:• r → 0: dead prior neurons.• δr → 1: posterior collapse.Together with the reconstruction loss, the NELBO for a 1-dimensional P-VAE reads:EQUATION)Finally, it is easy to show that for K-dimensional latent space, eq. ( 14) generalizes to:EQUATIONwhere ⊙ and • denote the Hadamard (element-wise) and vector products, respectively. 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines?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 not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.• The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in 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 or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness 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 a 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 as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] 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 we encourage authors to take this into account and make a 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?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 an equivalent approval/review based on the requirements of your country or institution) were obtained? 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. • 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 we 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 as the institution conducting the review. 
 
 Appendix: Equation ( 17) mirrors sparse coding due to 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 two major themes in theoretical neuroscience. Let's explore this curious connection to sparse coding more closely below.Suppose our P-VAE has K statistically independent neurons, and z ∈ Z K ≥0 is the spike count variable, where Z ≥0 = {0, 1, 2, . . .} is the set of non-negative integers. Let us use bold font r and δr to 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 to construct the posterior rates, λ i (x) = r i δr i (x). More explicitly, we have:EQUATION)Note that, unlike a standard Gaussian the prior in P-VAE is parameterized using r, which is learned from data along with the other parameters. Similar to standard Gaussian VAEs, δr(x) is parameterized as a neural network.Following the sparse coding literature, we will now assume our decoder generates the input image x ∈ R M as a linear sum of K basis elements, Φ ∈ R M ×K . We approximate the reconstruction loss as the mean squared error between the input x, and its reconstruction Φz. Given these assumptions, the VAE reconstruction loss becomes:EQUATION)For a linear decoder, the reconstruction term ∥x -Φz∥ 2 2 contains only the first and second moments of z. Consequently, the expectation in eq. ( 20) can be analytically resolved. This yields the loss function, and consequently, its gradients in closed form. Specifically, for the Poisson case, we only need to know the following expectation values:EQUATION)More generally, whenever the VAE decoder is linear, the following result holds:EQUATION)Note that a linear decoder is the only assumption we needed to obtain this closed-form solution.There are no restrictions on the form of the encoder: it can be linear, or as complicated as we want.We only have to compute the mean and variance of the posterior. Here are the reconstruction losses for both Poisson and Gaussian VAEs with linear decoders, put side-by-side for comparison:Poisson:EQUATION)Given these assumptions, the NELBO (eq. ( 17)) for P-VAE with a linear decoder becomes:EQUATION)Recall that we have f (y) = 1 -y + y log y (see Fig. 6 ). We introduced the β term here to control the trade-off between the reconstruction and the KL term. Additionally, we dropped the explicit dependence of δr(x) on the input image x to increase readability.We can further simplify the P-VAE architecture by making the encoder also linear. Let's use W ∈ R K×M to denote encoder weights, and assume an exponential link function mapping the input into the residual firing rates. In other words, we have δr = exp(W x). To obtain the final form of the loss function, we start from eq. ( 24), plug in log(δr) = W x, and rearrange some terms to find:EQUATIONB Architecture, training, and hyperparameter detailsWe consider three datasets in this paper. We tile up the van Hateren dataset of natural images [102] and CIFAR10 into 16 × 16 patches and apply whitening and contrast normalization using the code made available by Boutin et al. [103] . This operation results in the following total number of samples:• van Hateren: #train = 107,520, #validation = 28,224,• CIFAR 16×16 : #train = 200,000, #validation = 40,000.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 3 and 5): 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 3 and 5, 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).For sparse coding results, we focused on models with linear decoders. For the fully linear models (Figs. 4 and 8 ) 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 16×16 datasets, the encoders had 5 layers (2 × conv each). The decoders had 8 convolutional layers (1 × conv each). For the MNIST dataset, the encoders had 7 layers (2 × conv each). The decoders had 10 convolutional layers (1 × conv each). For all convolutional encoders, the output from ResNet was followed by a learned pooling layer. The pooled output was then fed into a feed-forward layer inspired by Transformers [123] , which includes a layer norm as the final operation, the output of which was 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.We experimented with both leaky_relu and swish activation functions [124, 125] , and found that swish consistently outperformed leaky_relu in all our experiments across datasets and VAE models.Please see our code for the full architecture details. There are n = 5 circles in each condition, corresponding to 5 random initializations. Using Monte-Carlo samples [132] and our Poisson reparameterization trick (Algorithm 1) to estimate gradients performs comparably to the situation where exact gradients are available (see eqs. ( 22) to ( 24)). EX, exact, MC, Monte-Carlo, ST, straight-through [106] . Figure 8 : Identifying dead neurons using a histogram-based method. We bin the KL 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-VAE, and L-VAE. The KL term is a single number for the C-VAE because its latent space consists of a 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 in Fig. 4 , where both encoder and decoder are linear. Table 2 uses this method to quantify the proportion of active neurons for VAEs across different datasets and choice of encoder architectures. NeurIPS Paper Checklist Justification: 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 and 5. Specifically, we have a dedicated paragraph for limitations in section 5. We evaluated our claims using multiple well-known datasets such as the van Hateren natural images [102] , CIFAR10, and MNIST, on tasks such as reconstruction, sparse coding, and classification.Guidelines:• The answer NA means that the paper has no 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 a 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 on the scope of the claims made, e.g., if the approach was 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.• 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 are provided or not. • If the contribution is a 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 is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. Justification: As stated in section 4, the paper reports confidence intervals and t-test significance 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. Justification: We provide details about our compute resources (GPUs), and duration of training in section 4.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 as well as 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 it into the paper). 


Decision: 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.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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:

- How do we know that the approximate posterior is Poisson? 
- 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?

Experiments:
- How does the ELBO vary with the dimension of the overcomplete latent space considered?
- 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.
questions: 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: limitations: The authors adequately discuss limitations of their work, and its potential societal impact.

I 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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.

To 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.

Despite the above points, I still think that this is a solid paper with potential contributions to the computational neuroscience community.
questions: 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.
- What’s the significance of the global prior parameter $r$? Have you tried ablation studies where the posteriors are directly parameterized?
- 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?
- There are multiple bold entries for some columns in tables 3 and 5, what do these mean?
- 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: limitations: The authors have adequately addressed the limitations of their work.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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.
- 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.
- 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.
- 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")
- 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}$.
questions: questions: See weaknesses.
limitations: limitations: Some limitations are addressed, but overall the presentation of the experimental results doesn't seem too balanced.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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.
- 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)
questions: 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.
- Fig 1b is a bit confusing as the sampling of the latents isn't shown whereas it was shown in Fig 1a.
- 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.
- 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: limitations: The authors have sufficiently discussed limitations.

Comments: title: reply to authors
comment: Thank you for your clarifications. I think my main concerns are well addressed and I now better understand why this work is important and interesting. I am hereby raising my score to a 7.
I thank the authors for addressing my questions and concerns. Given the lack of more illustrative experiments (in my opinion), I maintain my fairly very positive score for the paper.
I would like to thank the authors for addressing my comments.

The additions and changes proposed for the paper would certainly strengthen it and so I will raise my score to a 7
Thank you for your thorough rebuttal and for carefully addressing all the concerns raised by myself and the other reviewers. If the final version of the paper incorporates the promised revisions, I believe it will be a strong contribution that merits acceptance.

I have updated my score from 5 to 7.