Latent Diffusion for Neural Spiking Data

Authors: jaivardhan  kapoor
auguste  schulz
julius  vetter
felix  pei
richard  gao
jakob h macke
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 16:40:34

Abstract: Modern datasets in neuroscience enable unprecedented inquiries into the relationship between complex behaviors and the activity of many simultaneously recorded neurons. While latent variable models can successfully extract low-dimensional embeddings from such recordings, using them to generate realistic spiking data, especially in a behavior-dependent manner, still poses a challenge. Here, we present Latent Diffusion for Neural Spiking data (LDNS), a diffusion-based generative model with a low-dimensional latent space: LDNS employs an autoencoder with structured state-space (S4) layers to project discrete high-dimensional spiking data into continuous time-aligned latents. On these inferred latents, we train expressive (conditional) diffusion models, enabling us to sample neural activity with realistic single-neuron and population spiking statistics. We validate LDNS on synthetic data, accurately recovering latent structure, firing rates, and spiking statistics. Next, we demonstrate its flexibility by generating variable-length data that mimics human cortical activity during attempted speech. We show how to equip LDNS with an expressive observation model that accounts for single-neuron dynamics not mediated by the latent state, further increasing the realism of generated samples. Finally, conditional LDNS trained on motor cortical activity during diverse reaching behaviors can generate realistic spiking data given reach direction or unseen reach trajectories. In summary, LDNS simultaneously enables inference of low-dimensional latents and realistic conditional generation of neural spiking datasets, opening up further possibilities for simulating experimentally testable hypotheses.

Main paper: 
 introduction: 
 Modern datasets in neuroscience are becoming increasingly high-dimensional with fast-paced innovations in measurement technology [1, 48, 23] , granting access to hundreds to thousands of simultaneously recorded neurons. At the same time, the types of animal behaviors and sensory stimuli under investigation have become more naturalistic and complex, resulting in experimental setups with heterogeneous trials of varying length, or lacking trial structure altogether [35, 33, 57] . Therefore, a key target in systems neuroscience has shifted towards understanding the relationship between high-dimensional neural activity and complex behaviors.For high-dimensional neural recordings, analyses that infer low-dimensional structures have been very useful for making sense of such data [11] . For example, latent variable models (LVMs) are often used to identify neural population dynamics not apparent at the level of single neurons [61, 32, 40] . More recently, deep learning-based approaches based on variational autoencoders (VAEs) [26, 44, 51, 36, 63, 22, 6, 16] have become particularly popular for inferring latent neural representations due to their expressiveness and ability to scale to large, heterogeneous neural recordings with behavioral covariates [36, 63, 16] .However, in addition to learning latent representations, another important consideration is the ability to act as faithful generative models of the data. In other words, models should be able to produce diverse, realistic samples of the neural activity they were trained on, ideally in a behavior-or stimulus-dependent manner. Models with such capabilities not only afford better interpretability analyses and diagnoses for whether structures underlying the data are accurately learned, but have a variety of downstream applications surrounding the design of closed-loop in silico experiments. For example, with faithful generative models, one can simulate population responses to hypothetical sensory, electrical, or optogenetic stimuli, as well as possible neural activity underlying hypothetical movement patterns. Most VAE-based approaches focus on the interpretability of the inferred latents, but not the ability to generate realistic and diverse samples when conditioning on external covariates, while sample-realistic models (e.g., based on generative adversarial networks (GANs) [17] ) do not provide access to underlying low-dimensional representations. As such, there is a need for models of neural population spiking activity that both provide low-dimensional latent representations and can (conditionally) generate realistic neural activity.Here, we propose Latent Diffusion for Neural Spiking data (LDNS), which combines the ability of autoencoders to extract low-dimensional representations of discrete neural population activity, with the ability of (conditional) denoising diffusion probabilistic models (or, diffusion models) to generate realistic neural spiking data by modeling the inferred low-dimensional continuous representations. LDNS allows for (un)conditional generation of neural spiking data through combining a regularized autoencoder with diffusion models that act on the low-dimensional latent time series underlying neural population activity.Diffusion models [49, 20, 50] have been highly successful for conditional and unconditional data generation in several domains, including images [20] , molecules [59] , and audio spectrograms [27] and have demonstrated samplingfidelity that outperforms that of VAEs and GANs [20] . A key strength of diffusion models that makes them particularly attractive in the context of modeling neural datasets is the ability to flexibly condition the generation on various (potentially complex) covariates, such as to simulate neural activity given certain behaviors. Recently, diffusion models have been extended to continuous neural time series such as local field potentials (LFPs) and electroencephalography (EEG) recordings [53] . However, due to the discrete nature of spiking data, standard diffusion models cannot be easily applied, thus excluding their use on many datasets in systems neuroscience.To bypass these limitations, LDNS employs a regularized autoencoder using structured state-space (S4) layers [18] to project the high-dimensional discrete spiking data into smooth, low-dimensional latents without making assumptions about the trial structure. We then train a diffusion model with S4 layers as a generative model of the inferred latents-akin to latent diffusion for images [45] , where generation can be flexibly conditioned on behavioral covariates or task conditions.A fundamental assumption of most low-dimensional latent variable models is that all statistical dependencies between observations are mediated by the latent space. However, in neural spiking data, there are prominent statistical dependencies that ought to persist conditional on the latent state, e.g., single-neuron dynamics such as refractory periods, burstiness, firing rate adaptation, or potential direct synaptic interactions. We show how such additional structure can be accounted for in LDNS, by equipping it with an expressive observation model [41, 32, 54, 62] : We use a Poisson model for spike generation with autoregressive couplings which are optimized post hoc to capture the temporal structure of single-neuron activity. This allows LDNS to capture a wide range of biological neural dynamics [55] , with only a small additional computational cost.Main contributions In summary, LDNS is a flexible method that allows for both high-fidelity diffusion-based sampling of neural population activity and access to time-aligned low-dimensional representations, which we validate on a synthetic dataset. Next, we show the utility and flexibility of this approach on complex real datasets: First, LDNS can handle variable-length spiking recordings from the human cortex. Second, LDNS can unconditionally generate faithful neural spiking activity recorded from monkeys performing a reach task. We demonstrate how LDNS can be equipped with an expressive autoregressive observation model that accounts for additional dependencies between data points (e.g., single neuron dynamics), increasing the realism of generated samples. Third, LDNS can generate realistic neural activity while conditioning on either reach direction or full reach trajectories (time series), including unseen behaviors that are then accurately decoded from the simulated neural data. Overall, LDNS enables simultaneous inference of low-dimensional latent representations for single-trial data interpretation and high-fidelity diffusion-based (conditional) generation of diverse neural spiking datasets, which will allow for closed-loop in silico experiments and hypothesis testing. 
latent diffusion for neural spiking data (ldns): 
 We consider a dataset recorded from a population of n neurons, consisting of trials with spiking data s ∈ N n×T 0 (sorted into bins of fixed length resulting in spike counts over time), and optional simultaneously recorded behavioral covariates y ∈ R n (that can also be time-varying y ∈ R n×T ). A dataset of M such trials D can be written as D = {s (i) , y (i) }, possibly with varying trial lengths T 1 . . . T M . We make the assumption that a large fraction of the variability in this dataset can be captured with a few underlying latent variables z ∈ R d×T , where d < n.Our goal is to generate realistic spiking data s * that faithfully capture both population-level and singleneuron dynamics of s 1...T with the ability to optionally condition the generation on behavior y cond . To this end, we propose a new method, LDNS, that combines the strength of neural dimensionality reduction approaches with that of diffusion-based generation.LDNS uses a two-stage training framework, adopted from the highly successful family of latent diffusion models (LDMs) [45, 8, 59] . To train LDNS, we first train a regularized autoencoder [14] to compress the spiking data into a low-dimensional continuous latent space (Fig. 1 ). Concretely, we focus on two objects of interest for the LDNS autoencoder: (1) inferring a time-aligned, lowdimensional smooth representation z ∈ R d×T that preserves the shared variability of the spiking data, and (2) predicting smooth firing rates λ that are most likely to give rise to the observed spiking data.In the second stage, we train a diffusion model in latent space, possibly employing conditioning to make generation contingent on external (e.g., behavioral) covariates (Fig. 1 ). For the diffusion model, our main objective is the generation of z * ∈ R d×T that captures the distribution of inferred autoencoder latents. We also want the ability to sample latent trajectories of varying length.In both stages, we use structured state-space (S4) [18] layers for modeling temporal dependencies. S4 layers consist of state-space transition matrices that can be unrolled into arbitrary-length convolution kernels, allowing sequence modeling of varying lengths. For details on network architectures and S4 layers, see appendix A1.For the spiking data, we choose a Poisson observation model, and train autoencoders by minimizing the Poisson negative log-likelihood of the input spikes s given the predicted rates λ = decoder(z). To enforce smoothness in the latent space, where z = encoder(s), we add an L 2 regularization along with a temporal smoothness regularizer over z, resulting in the combined lossEQUATIONCode available at https://github.com/mackelab/LDNS.To prevent the autoencoder from predicting highly localized Poisson rates, which have sharp peaks at input spike locations, we further regularize training using coordinated dropout [24] , i.e., we randomly mask input spikes and compute the loss on the predicted rates at the masked locations (details in appendix A1.2).Accounting for single-neuron dynamics with an expressive observation model So far, LDNS (like most latent variable models for neural data) uses a Poisson observation model, which assumes that all statistical dependencies are mediated by the latent state. To address this limitation and to capture dynamics and variability, which are "private" to individual neurons (such as refractory periods or burstiness), we propose to learn an autoregressive observation model. We make the predicted Poisson rates for each neuron i dependent also on recent spiking history, by including additional spike history couplings h i [41, 32] , resulting in the observation modelEQUATIONwhere T ′ corresponds to the time-lagged window length. This modification is learned post hoc, and the parameters h i are fit with a maximum-likelihood objective (details in appendix A1.3). This approach does not alter the latent dynamics, while augmenting the model with single-neuron autoregressive dynamics. We observe that including spike history increases the realism of generated data and enables us to accurately capture single-neuron autocorrelation structures (Sec. 3.4).In the second stage of training, we train diffusion models [20] to generate (conditional) samples from the distribution of inferred latents. The training dataset therefore contains autoencoder-derived latents of each trial, and optionally, additional conditioning information such as the corresponding behavior, i.e., D z = {z (i) = encoder(s (i) ), y (i) }.Diffusion models aim to approximate the data distribution q(z) through an iterative denoising process starting from standard Gaussian noise. For latent z (denoted as z 0 for diffusion timestep 0), we first produce a noised version at step t by adding Gaussian noise of the form q(z t |z 0 ) = N ( √ ᾱt z 0 , (1 -ᾱt )I). Here, ᾱt = t k=1 α k , where the noise scaling factors α 1 . . . α T follow a fixed linear schedule. We then train a neural network to approximate the reverse process p θ (z t-1 |z t ) for each diffusion timestep. The true (denoising) reverse transition q(z t-1 |z t ) is intractable-however, we can apply variational inference to learn the conditional reverse transition q(z t-1 |z t , z 0 ), which has a closed form written asEQUATIONEQUATIONWe train the neural network µ θ (z t , t) to approximate the mean of this distribution by optimizing the loss E z0∼Dz,ϵ0,t ∥ϵ θ (z t , t) -ϵ 0 ∥ 2 , where ϵ 0 is the noise used to generate z t from z 0 , and ϵ θ (z t , t) is the equivalent reparameterization for µ θ (z t , t). At test time, we sequentially sample z t-1 given z t using the learned transition p θ (z t-1 |z t ), starting from standard Gaussian noise. Using S4 layers in the denoising network allows us to generate latents with varying lengths. This is achieved by unrolling the state transition matrix in the S4 layers to the desired length for each denoising step.Diffusion models may be conditioned on fixed-length and time-varying covariates y, in which case we learn the approximate reverse transition p θ (z t-1 |z t , y). Details on the conditioning mechanisms in appendix A1.4.3 Experiments and ResultsQuestion: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: We discuss the limitations of our work in a paragraph in the last section of the main paper.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. 
datasets and tasks: 
 We first evaluate the performance of LDNS on a synthetic spiking dataset where we have access to the ground-truth firing rates and latents. We choose the Lorenz attractor [29] as a low-dimensional, non-linear dynamical system commonly used in neuroscience [7, 36] . We simulate rates as an affine mapping from the 3-dimensional system to a 128-dimensional neural space, and sample from a Poisson distribution to generate spiking data. Next, we showcase the applicability of LDNS on two neural datasets: We apply our method on a highly complex dataset of human neural activity (128 units) recorded during attempted speech [56] . This dataset poses a challenge to many modeling approaches due to the different imagined sentences, resulting in variable lengths of the neural time series (between 2-10 seconds with a sampling rate of 50 Hz). Finally, we apply LDNS to model premotor cortical activity (182 units) recorded from monkeys performing a delayed center-out reach task with barriers [9, 38] . The multi-modal nature of the dataset allows us to assess both unconditional as well as conditional generation of neural spiking activity given monkey reach directions and entire velocity profiles of the performed reaches. See appendix A2,A3 for data and training details.For the unconditional generation of monkey reach recordings (Sec. 3.4), we train both a Poisson observation model as well as a spike history-dependent autoregressive observation model. For all other experiments, we only train a Poisson observation model.We compare LDNS to the most commonly known VAE-based latent variable model: Latent Factor Analysis via Dynamical Systems (LFADS [51, 36, 47] ), which has been shown to outperform various classical latent variable models on a variety of tasks ( [38] , details in appendix A4).To ensure that we use optimal hyperparameters for LFADS, we follow the auto-ML pipeline proposed by Keshtkaran et al. [25] . This approach, termed AutoLFADS, has been shown to perform better than the original LFADS on benchmark tasks [38] . For the unconditional generation of monkey reach recordings, we further compared to additional VAE baselines [21, 62] (appendix A5).Metrics For all experiments, we assess how well LDNS-generated samples match the spiking data in the training distribution. Concretely, we compare population-level statistics by computing 1) the distribution over the population spike count, which sums up all spikes co-occurring in the population in a single time bin (i.e., spike count histogram), and 2) pairwise correlations of LDNS samples and the spiking data for each pair of neurons. For single-neuron statistics, we compare 3) the mean and 4) standard deviation of the inter-spike-interval distribution for each neuron (mean isi and std isi). When multiple spikes occur in a single time bin, the spike times are distributed equally in this bin [12] . To further evaluate population dynamics, we compare the principal components of smoothed spikes.We simulate trials of length 256 timesteps from the three-dimensional (chaotic) Lorenz system (Fig. 2a ). The regularized autoencoder extracts smooth latent time series (eight latent dimensions) from the 128-dimensional spiking data, resulting in smooth firing rate predictions that closely match the ground-truth rates (Fig. 2b , Supp. Fig. A2, A3 ). We then train a diffusion model on the extracted autoencoder latents. Latents sampled from the diffusion model (red) preserve the attractor geometry of the Lorenz system (Fig. 2c , left, three of the eight latent dimensions), indicating that LDNS preserves a meaningful latent space. The architectural choice of S4 layers allows for length generalization: although we train on time segments of 256-time steps, we can sample and successfully generate latent trajectories that are much longer, but still accurately reflect the Lorenz dynamics (Fig. 2c , middle, 16× longer generation). In comparison, LFADS exhibits instabilities when generating such longer sequences (appendix A6.1). Overall, the latent time series distribution is captured well by the diffusion model, with matching power spectral densities (PSD) per latent dimension (Fig. 2c , right, other dimensions in Supp. Fig. A3 ).To assess the sampling fidelity of the generated synthetic neural activity, we compute a variety of spike statistics frequently used in neuroscience. LDNS captures both population-level statistics, such as the population spike count histogram and pairwise correlations between neurons (Fig. 2d ), as well as single-neuron statistics, quantified by the mean and standard deviation of inter-spike-intervals (Fig. 2e ). LDNS also captures the temporal correlation structure of the data (Supp. Fig. A4 ). These results demonstrate that LDNS can both perform inference of low-dimensional latents and provide high-fidelity diffusion-based generation that perfectly captures the statistics of the ground-truth synthetic data.Next, we assess whether LDNS is capable of capturing real electrophysiological data, applying it to neural recordings from human cortex during attempted speech (Fig. 3a , top, Willett et al. [56] ).A participant with a degenerative disease who is unable to produce intelligible speech attempts to vocalize sentences prompted on a screen, while neural population activity is recorded from the ventral premotor cortex. Since there is a large variation in the length of prompted sentences (Fig. 3a , bottom), this dataset allows us to evaluate the performance of LDNS on real data in naturalistic settings with variable-length and highly heterogeneous dynamics.To account for varying trial length during autoencoder training, we pad all trials to a maximum length of 512 bins and compute the reconstruction loss only on the observed time bins. For the diffusion model, we indicate the target trial length with a binary mask as a conditioning variable.This approach allows us to infer time-aligned latents underlying the cortical activity of the participants, compressing the population activity by a factor of four before training an unconditional diffusion model on these latents. Resulting samples of LDNS, mimicking human cortical activity, are visually indistinguishable from the real data (Fig. 3b, c , additional samples in Supp. Fig. A7 ). This is reflected in closely matched population spike count histograms (Fig. 3d , left), and single neuron statistics such as mean and standard deviation of the inter-spike interval (Fig. 3d , right). Additionally, real and LDNS-sampled spikes display similar population dynamics, as reflected in the top principal components (Supp. Fig. A8 ). While LDNS tends to overestimate some pairwise correlations, it captures prominent features of the correlation structure in the data (Fig. 3e , Pearson correlation coefficient r = 0.47), and our analysis indicates that this slight mismatch already arises at the autoencoder stage (Supp. Fig. A9 ).LDNS allows for both inferring latent representations and generating variable-length trial data, making it applicable to complex real neural datasets without a fixed trial structure.We further evaluate LDNS in a different setting by applying it to model sparse spiking data recorded from a monkey performing a reaching task constrained by barriers that form a maze (Fig. 4a , left). The variety of different maze architectures leads to diverse reach movements of both curved and straight reaches (Fig. 4a , right). We again infer low-dimensional latent trajectories that capture the shared variability of the neural population and then train an unconditional diffusion model on these latents. Sampled spikes from LDNS closely resemble the true, sparse population data (Fig. 4b , additional samples in Supp. Fig. A11 ), and closely match population-level spike statistics (Fig. 4c ). Single neuron statistics in this low spike count regime (a maximum of three spikes per neuron in 5 ms bins) are also captured well (Fig. 4d ), and are on par with or better than LFADS [36, 25] (see Table 1 for summary of main comparisons, and appendix A5 for additional baselines [21, 62] ). Beyond spiking statistics, we observe that LDNS also preserves the temporal structure of population dynamics, as reflected in the top principal components of smoothed spikes (Supp. Fig. A15 ). Thus, LDNS can generate spiking data that is faithful at the level of both single-neuron and population dynamics. Both LFADS and the LDNS autoencoder are optimized by maximizing the Poisson log-likelihood, and thus cannot capture single-neuron dynamics such as refractoriness [41, 55] , which can have a strong influence on the observed autocorrelation structure. Given the overall sparsity of the spiking data and resulting low correlations (Supp. Fig. A10 ), we focus on the temporal structure of autocorrelations averaged within groups of neurons (≈ 45 neurons per group) split by their instantaneous correlation strength [32] : darker colors correspond to the highest correlated group of four, lighter colors correspond to the group with the second highest correlations (Fig. 4e ). We then compare these auto-correlations to those of grouped LDNS samples with a Poisson observation model (red). As expected, LDNS with Poisson observations is unable to capture the dip in the data auto-correlation at 5 ms lags (one time bin) (Fig. 4e , left).To overcome this mismatch, we train an additional spike history-dependent autoregressive observation model on top of the inferred rates (LDNSsh, for spike history). In contrast to the Poisson samples, autoregressive samples can capture this aspect of neural spiking data very accurately while also improving the overall fit to the empirical auto-correlation (Fig. 4e , right). Moreover, the post hoc optimization of these filters also improves modeling of other single-neuron, as well as population-level statistics, such as the population spike count histogram or the mean of the isi (Table 1 , Supp. Fig. A13 ).We view this post-hoc augmentation as a key modular contribution, which can be flexibly applied to other generative models. To this end, we extend AutoLFADS with spike history dependence (LFADSsh), improving its performance across metrics. The augmented LFADSsh also captures the dip in autocorrelation at 5 ms lags (Supp. Fig. A12 ). Still, in both observation model variants, LDNS maintains superior or comparable performance (Table 1 ).Thus, Poisson LDNS allows for the generation of spiking data that is on par or better in terms of sampling fidelity than previous approaches. Incorporating spike-history dependence and sampling spikes autoregressively allows us to further increase the realism of generated spike trains, leading to a large improvement on several of the considered metrics.Lastly, we assess the ability of conditional LDNS to generate realistic neural activity conditioned on behavioral covariates of varying complexity: the reach angle or entire velocity time series (Fig. 5a ). We first validate that the autoencoder predicts firing rates that allow us to linearly decode the behavior following the ridge-regression approach proposed in [38] . Decoded behavior from autoencoder reconstructed rates matches the true trajectories of unseen test trials (Fig. 5b ).Given that the autoencoder performs adequately, we then test the ability to generate neural time series conditioned on the initial reach angle performed by the monkey θ reach . Indeed, from the generated samples of neural activity, we can decode-using the same linear decoder-realistic reach kinematics that are consistent with the conditioning angle θ reach and overall reach kinematics (Fig. 5c ). This indicates that LDNS can generate realistic neural activity consistent with a queried reach direction.An even more challenging task that is intriguing for hypothesis generation is the ability to mimic an entire experiment and ask what the neural activity would have looked like if the monkey had performed a particular hypothetical movement. To this end, we train a diffusion model on the same autoencoder-inferred latents but now condition on entire velocity traces (Fig. 5d ). Velocityconditioned LDNS is able to produce different samples of neural activity that are consistent with, but not exact copies of, the reach trajectories of the held-out trials given as the conditioning covariate. Such closed-loop conditioning experiments open the possibility of making predictions about neural activity during desired unseen behaviors, and thus make experimentally testable predictions.Finally, to understand how LDNS incorporates behavioral information, we analyzed latent trajectories that were conditionally sampled based on straight reach movements in different directions (Fig. 5e ). Individual samples of latent trajectories vary smoothly within a trial (Fig. 5f ), while reach direction varies smoothly across samples in the first principal component (PC1) of the latents (Fig. 5g, left ). Projection onto the first two PCs of latent trajectories shows clear clustering by reach direction (Fig. 5g , right), and we show that such clustering arises already at the autoencoder stage (Supp. Fig. A17 ).In summary, LDNS not only produces faithful spiking samples but also allows for flexible conditioning. Furthermore, LDNS learns an interpretable latent space with behaviorally-relevant structure.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
related work: 
 Latent variable models of neural population dynamics LDNS builds on previous LVMs in neuroscience, which have been extensively applied to infer low-dimensional latent representations of neural spiking data [61, 32, 39, 58, 62, 13, 28, 63] (see [38] for a comprehensive list.) In addition to capturing shared population-level dynamics and dependence on external stimuli [4] , LVMs have been extended to allow autoregressive neuron-level (non-Poisson) dynamics [32, 13, 62] or even direct neural interactions [54] . While these methods often have useful inductive biases (e.g., linear dynamical systems [32, 28] or Gaussian process priors [61] ), these models are typically not expressive enough to yield realistic neural samples across a range of conditions.Deep LVMs and other deep learning-based approaches Variational autoencoders (VAEs) [26] are particularly popular in neuroscience as they allow us to infer low-dimensional dynamics underlying high-dimensional discrete data [63, 16, 46] , especially when combined with nonlinear recurrent neural networks [36, 21] . VAEs have been used to infer identifiable low-dimensional latent representations conditioned on behavior [63, 21] and have incorporated smoothness priors using Gaussian Processes to regularize the latent space [16] . However, the generation performance of VAEs is rarely explored in neuroscience. Besides VAEs, generative adversarial networks (GANs [17] ) have been proposed to synthesize spiking neural population activity [34, 42] . While GANs produce high-fidelity samples, they are challenging to train reliably and lack a low-dimensional latent space. More recently, transformer-based architectures have also been adapted to model neural activity [5, 60] , though often with the focus of accurate decoding of behavior instead of generation of realistic spiking samples, while also lacking an explicit latent space [3] . Lastly, deterministic approaches utilizing RNNs for dynamical systems reconstruction also target low-dimensional latent dynamics underlying neural data [19] , but they do not act as probabilistic generative models.Diffusion models LDNS leverages recent advances in diffusion models, which have become state-of-the-art for high-fidelity generation in several domains [20, 27] , including continuous-valued neural signals such as EEG [53] , as well as in time series forecasting and imputation tasks [2, 52, 43] . Similar to the LDNS architecture, Alcaraz and Strodthoff [2] also use an S4-based denoiser for imputation. More specifically, LDNS is inspired by latent diffusion models [45, 27, 59, 15] , which benefit from operating on the latent space of an autoencoder and flexible conditioning mechanisms to generate samples based on a given covariate, as is done with text-to-image [45] and other cross-modality scenarios. Conveniently, this allows LDNS to bypass the challenges of directly modeling discrete-valued spiking data, by instead transforming spikes into the continuous latent space. 
summary and discussion: 
 We here proposed LDNS, a flexible generative model of neural spiking recordings that simultaneously infers low-dimensional latent representations and generates realistic neural activity conditioned on behavioral covariates. We apply LDNS to model three different datasets: synthetic data simulated from chaotic Lorenz dynamics, human cortical recordings with heterogeneous and variable-length trials, and finally, neural recordings in monkeys performing reach actions in a maze. Through our experiments, we demonstrate how several features of LDNS are beneficial for modeling complex datasets in neuroscience:First, following other LDMs in the literature, LDNS decouples latent inference and probabilistic modeling of the data, offering flexibility in reusing the trained autoencoder and diffusion model. For the monkey recordings, all diffusion models (unconditional, conditioned on reach angle, and conditioned on hand velocities) operate in the latent space of the same autoencoder, in contrast to existing approaches that require end-to-end retraining for each type of conditioning variable. LDNS is also faster to train than AutoLFADS, which requires population-based training to optimize hyperparameters (appendix A3.1). Second, we show that LDNS autoencoders can be augmented with per-neuron autoregressive dynamics to capture single-neuron temporal dynamics (e.g., refractoriness), which otherwise cannot be captured with population-level shared dynamics. Third, as a result of the length-generalizable autoencoders and diffusion models using S4 layers, LDNS can generate variable-length trials in both the Lorenz example and human cortical recordings-a feature that will be particularly useful in modeling datasets recorded during naturalistic stimuli or behavior.Altogether, these features enable LDNS to generate realistic neural activity, especially when conditioned on behavioral covariates. In our experiments, we demonstrate that unseen movement trajectories can be used to conditionally generate samples of neural activity, from which we can decode these hypothetical behaviors. These generated latent trajectories reflect behavioral information in an interpretable way. Our methodology is general and can be applied to recordings from any brain region, beyond the motor and speech cortex examples shown here. Thus, LDNS opens up further possibilities for hypothesis generation and testing in silico, potentially enabling stronger links between experimental and computational works.Limitations In real neural data, the latent dimensionality of the system is not known, and as with all LVMs (which often assume that population dynamics are intrinsically low-dimensional), choosing an appropriate latent dimension can be challenging. Furthermore, any modeling errors at the encoding and decoding stage of the autoencoder will affect the overall performance of the latent diffusion approach. Nevertheless, in our experiments, we found that autoencoder training is fast, stable, and reasonably robust to hyperparameter configurations. While LDNS was still able to model the data well under relatively severe compression (e.g., 182-to-16 for the monkey recordings), optimizing latent dimensionality to balance expressiveness and interpretability remains a goal for future research.Broader impact Realistic spike generation capabilities increase the risk of research manipulation by generating synthetic data that may be difficult to detect. On the other hand, LDNS could be useful for the dissemination of privatized clinical data, though we acknowledge the critical importance of protecting data privacy when working with sensitive human participant data. Finally, synthetically generated data (conditioned on unseen behavioral conditions) could be useful for augmenting the training of brain-computer interface decoding models.by the AI4Med-BW graduate program. We thank Chethan Pandarinath for providing access to their compute cluster to train AutoLFADS. We thank Christian F. Baumgartner and all Mackelab members for feedback and discussions. We would like to also thank our reviewers for their insightful comments which improved our paper. The parameters of the autoregressive observation model in Eq. ( 2) are learned by maximizing the Poisson log-likelihood. Training is performed jointly for all neurons with a history length of T ′ = 20, corresponding to 100 ms, using the AdamW optimizer [30] (learning rate 0.1, weight decay 0.01).In our implementation, we use the Softplus function given by Softplus(x) = log(1 + exp(x)) as an approximation to the exponential function in Eq. ( 2), which is accurate for the low-count regime while increasing numerical stability. During autoregressive sampling, we limit the maximum possible spike count to 5 spikes, which corresponds to the biological maximum, limited by the refractory period for 5 ms time bins.We consider four variants of our proposed diffusion model with time-mixing and channel-mixing layers for four different tasks. In all cases, except for the conditioning mechanism, the internal architecture remains the same (Supp. Fig. A1 ). For diffusion timestep and fixed-length conditioning vector, we shift and scale the inputs and outputs to the time mixing and channel mixing blocks using adaptive instance normalization, as done in Peebles and Xie [37] .1. Unconditional generation for synthetic spiking data with Lorenz Dynamics and cortical spiking data in monkeys -we use only time conditioning.2. Angle-conditioned generation for cortical monkey spiking data -we add an embedding MLP([cos θ, sin θ]) of the reach angle θ to the timestep embedding output.3. Trajectory conditioned generation for cortical monkey spiking data -we concatenate the hand velocities v x , v y of the monkey with the input as two additional channels.4. Unconditional variable-length generation for cortical human spiking data -we concatenate the desired length (with a maximum sequence length of 512) as a centered binary mask channel in the input. We only backpropagate through the central section of the output corresponding to the binary mask.We use a DDPM scheduler with 1000 timesteps and ϵ-parameterization. To stabilize and speed up training, we train all diffusion models using a smooth L 1 loss, written asEQUATIONwith δ = 0.05. Table 2 : Training details for autoencoder models on Lorenz, Monkey reach, and Human BCI datasets. We used the AdamW [31] optimizer, whose learning rate was linearly increased over in the initial period and then decayed to 10% of the max value with a cosine schedule. Mean firing rate for Lorenz was 0.3. In all cases, we used K = 5 for the temporal smoothness loss in Eq. 1. We performed all training and evaluation of LDNS on the Lorenz and Monkey reach datasets on an NVIDIA RTX 3090 GPU with 24GB RAM. For the Human BCI data, we used an NVIDIA A100 40GB GPU.The autoencoder for the Lorenz dataset is trained in ≈ 6 minutes, and the diffusion model in ≈ 20 minutes. For the evaluation, all sampling is performed on the GPU in 5 minutes. The effective GPU wallclock time (time when the GPU is utilized) for the entire training and evaluation run is within 30 minutes. Optimizing the autoregressive observation model took less than 1 minute.AutoLFADS, the baseline used for unconditional sampling for the Monkey reach dataset, was trained on a cluster of 8 NVIDIA RTX 2080TI GPUs for one day. As it requires automated hyperparameter tuning to achieve the best accuracy using population-based training (PBT, [24] ), AutoLFADS is significantly more compute-expensive to train than LDNS.For the Human BCI dataset, due to larger trial lengths, more data points, and more heterogeneous temporal dynamics, we trained a slightly larger autoencoder and diffusion model than in Monkey reach. The autoencoder took 50 minutes to train, and the diffusion model took 10 hours to train.Sampling from the trained model took 9 minutes, resulting in a total of under 12 hours of effective GPU wallclock time.We ran several preliminary experiments for LDNS to optimize the architecture and hyperparameters, as well as for designing appropriate evaluations. We estimate the total effective GPU wallclock time to be ≈ 10× that of the final model runs. The AutoLFADS baseline was only trained once with PBT, as this framework automatically optimizes the model hyperparameters.We implemented all training and evaluation code using the Pytorch framework2 , and used Weights & Biases3 to log metrics during training.Latent Factor Analysis via Dynamical Systems (LFADS) is a sequential variational autoencoder used to infer latent dynamical systems from neural population spiking activity [51, 36] . LFADS consists of an encoder, a generator, and optionally, a controller, all of which are RNNs. The generator RNN implements the learned latent dynamical system, given an initial condition and time-varying inputs.The internal states of the generator are mapped through affine transformations to lower-dimensional latent factors and single-neuron Poisson firing rates. The encoder RNN maps the neural population activity into an approximate posterior over the generator's initial condition.At each timestep, the controller RNN receives both encoded neural activity and the latent factors from the previous timestep and outputs an approximate posterior over the input to the generator. The entire model is trained end-to-end to maximize the ELBO, as is done in VAEs. To address the difficulty of hyperparameter optimization for LFADS, Population-Based Training (PBT) has been proposed to automate hyperparameter selection, termed AutoLFADS [25] .In our experiments with the monkey reach dataset, we use the PyTorch implementation of AutoLFADS [47] . We use the hyperparameters and search ranges from Pei et al. [38] , but omit the controller RNN to simplify generation from prior samples. Although this might limit the model's expressiveness, prior research indicates that the monkey reach data can be well-modeled as autonomous, without external inputs from the controller [10] . LFADS has previously performed well on this data without the controller [36] .We generate samples from LFADS by sampling initial conditions from the Gaussian prior, running the generator RNN forward, and Poisson-sampling spikes from the resulting firing rates. For inclusion of spike history in the observation model of LFADS, we used the same training method and hyperparameter settings as in LDNSsh (appendix A1.3).For an extended baseline comparison, we implemented two additional methods for the task of unconditional generation on the Monkey dataset (Sec. 3.4) -Targeted Neural Dynamical Modeling (TNDM, [21] ) and Poisson-identifiable VAE (pi-VAE, [63] ). It is important to note that while both TNDM and pi-VAE have demonstrated success in analyzing neural and behavioral data, neither was specifically designed for realistic spike train generation. The architectural choices in our implementation of these methods reflect their original intended applications in neural data analysis rather than generation of neural spiking data. Nevertheless, our comparisons show that LDNS, especially with spike-history, is superior or on par with all other methods (Table 4 ).TNDM [21] is a VAE-based model designed to jointly model neural activity and behavior. TNDM extends LFADS by using an RNN to generate latent dynamics that are mapped to both neural activity and behavioral variables. TNDM separates the latent space into behavior-specific and behaviorindependent subspaces to disentangle task-relevant and intrinsic neural dynamics.For our comparison on the unconditional monkey reach task, we trained TNDM using the architecture and hyperparameters proposed in the original implementation of the paper 4 . We used 64-dimensional latent dynamics for each of the two sets. These project to a total of 10 latent factors z (5 behaviorspecific z r and 5 behavior-independent z i ), which is the maximum number demonstrated in the original work.To generate unconditional samples, we sampled initial generator states from a standard normal prior N (0, I), then generated the latent dynamics and projected into neuron rates the same way as in LFADS. TNDM, performs well in matching real data in spike statistics (Supp. Fig. A13 cyan, Supp. Fig. A16e ) and temporal dynamics (Supp. Fig. A15 ). Overall, we observe that LDNS captures spike statistics better than TNDM, except for the the population spike history count. In all metrics, LDNS augmented with spike history outperforms TNDM on spike statistics.Poisson-identifiable VAE (pi-VAE) [63] is a VAE-based model for count data that ensures identifiability in the latent space. pi-VAE does not model temporal dependencies, instead treating each time point as an independent sample.We trained pi-VAE on the monkey dataset using the original architecture and hyperparameters 5 . We used a General Incompressible-flow Network as a decoder, with 2 behaviorally relevant dimensions and 2 independent dimensions in the latent space. However, our evaluation context differs significantly from the original paper's demonstrations: while pi-VAE was initially evaluated on 50ms time bins and straight reaches only, our comparison uses 5ms bins and conditions on angles across all reaches at both middle and end trajectory points. The "label", or behavior, is presented as a 4-dimensional vector containing the cosine and sine of initial and final reach angles. Since this has a conditional latent space, sampling is performed by sampling angles randomly.Importantly, sampling from pi-VAE does not introduce any temporal dependence between spike bins within a trial -pi-VAE was not intended to be a generative model of neural spiking data. The lack of temporal modeling in pi-VAE's is a fundamental limitation for generating realistic spike trains, as evident in our empirical comparisons (Supp. Fig. A13 in yellow, Supp. Fig. A15, A16f ). Note that this failure cannot be diagnosed simply from looking at the sampled spiking data (Supp. Fig. A14 ).• LDNS is designed specifically for the purpose of accurately generating neural spiking data (unconditionally or conditionally)-a task often ignored by other LVMs designed for neural data analysis such as LFADS, pi-VAE, and TNDM. • The S4 autoencoder and diffusion model in LDNS are trained in separate stages, offering modularity, while both components naturally account for temporal dependencies (unlike pi-VAE). • S4 is autoregressive, similar to other RNN-based models, but empirically we found it to perform better when extending past the training trial length (compared to LFADS, see Sec. A6.1). • One feature provided by some neural-behavioral analysis models (such as pi-VAE and TNDM) is an explicit disentangling of neural vs. behavior-relevant latents. While we found LDNS latents contain relevant behavioural information (Fig. 5e -g, Supp. Fig. A17 ), we did not explicitly supervise the latent space to induce this property. • Finally, the spike history-dependent observation model in LDNS is modular and can be optimized post-hoc using rate predictions of any model to improve spike generation quality.We observed this with LDNS as well as LFADS (Table 1 ). Method D KL psch RMSE pairwise corr RMSE mean isi RMSE std isi pi-VAE 0.0063 ± 2.9e-4 0.0031 ± 1.08e-5 0.064 ± 0.002 0.034 ± 0.001 TNDM 0.0028 ± 6.6e-5 0.0027 ± 1.17e-5 0.057 ± 0.004 0.029 ± 0.001 AutoLFADS 0.0040 ± 2.2e-4 0.0026 ± 1.25e-5 0.039 ± 0.003 0.029 ± 0.001 LDNS 0.0039 ± 3.9e-4 0.0025 ± 1.1e-5 0.037 ± 0.001 0.023 ± 0.001 AutoLFADSsh 0.0036 ± 2.1e-4 0.0026 ± 1.8e-5 0.034 ± 0.002 0.023 ± 0.0001 LDNSsh 0.0016 ± 6.2e-4 0.0025 ± 1.07e-5 0.024 ± 0.002 0.023 ± 0.001 To analyze whether LFADS [36] exhibits similar length generalization properties as LDNS, we trained an AutoLFADS model on the Lorenz dataset (256 bins). We used the same architecture as the Monkey dataset, with 40-dimensional latent dynamics. We sampled initial conditions from the LFADS prior, then generated dynamics for both the original 256 steps and for an extended duration of 16× the training length. Qualitatively, we observed that while LFADS produced trajectories resembling the attractor dynamics of the ground truth Lorenz system (Supp. Fig. A5a , across various dimension combinations), these trajectories often diverged when run for longer intervals (Supp. Fig. A5b ). However, the system eventually returned to typical dynamics. Furthermore, when generating for extended durations, we observed that the mean population firing rates sometimes reached extreme values in some samples (Supp. Fig. A6 ), though they eventually returned to typical ranges. This behavior was not observed in LDNS samples, suggesting that bidirectional generation in the diffusion model provides more stability in variable-length generation. , TNDM [21] , LFADS [36] , LFADS with spike history (LFADSsh), LDNS and LDNS with spike history (LDNSsh). Mean and standard deviation across 5 folds sampled with replacement. Justification: Our work is methodological and empirical, and we justify our claims by evaluating our method's performance on synthetic and real-world data.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 be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: We provide a table with all relevant hyperparameters in the Appendix. We also provide yaml configuration files in the zipped code directory.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 to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
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?Answer: [Yes]Justification: For all results that require statistical significance (Table 1 in main paper, Table 4 in appendix), we report performance and standard deviation over 5 sets of samples.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence 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). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. Justification: We list the hardware configurations used for training and evaluation, as well as estimates on wallclock time for this hardware in the Appendix.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). 
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?Justification: We have read and confirmed that we conform to all the points mentioned in the NeurIPS Code of Ethics, in every respect.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). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss the broader impact of our work in the last section of the main paper. We do not foresee any negative impact of our work, apart from possibly synthesizing scientific or clinical data with dishonest intent.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: In our opinion, we do not foresee a high risk of misuse of our work.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?Answer: [Yes] Justification: We cite the real-world datasets that we use in the main paper (CC-BY-4.0 for the Monkey reach dataset, CC0 1.0 for the Human BCI dataset.) We have made sure to not use any code or data whose license prevents us from using it. (Apache-2.0 for code for the S4 layers, MIT license for Weights & Biases, BSD-3 Clause for Pytorch, Apache-2.0 license for the baseline implementation LFADSTorch.)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: We did not interact with human subjects for any experiments or data collection in this work. 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.Github link: https://github.com/state-spaces/s4/blob/main/models/s4/s4.pyPaszke et. al. PyTorch: An Imperative Style, High-Performance Deep Learning Library (2019)Lukas Beiwald. Experiment Tracking with Weights and Biases (2022)Code adapted from Github respository: https://github.com/HennigLab/tndm.Code adapted from article by Lyndon Duong (2021) -https://www.lyndonduong.com/pivae. 
 
 Appendix:  


Decision: There was broad consensus among the reviewers that this paper is technically solid and interesting for the NeurIPS community. The additional empirical validation done in the authors' rebuttal was found to be particularly useful.

Reviews: review: paper_summary: paper_summary: The authors here propose a new autoencoder style latent variable model for neuroscience which flexibly adapts to variable time-series using S4 encoders and decoders. They then train diffusion based models with option behavioral covariates to generate realistic neural spiking data. Additionally, they use a more flexible spiking likelihood with history-filters which more accurately capture within neuron statistics. They evaluate their model on a synthetic lorenz-attractor example as well as multiple real-work datasets.
strength_weakness: strength_weakness: I think the authors could do a better job communicating the disparate contributions and capabilities of this model more clearly to the reader, especially in contrast to existing approaches. 

For example, I am having a hard time understanding to what extent the spike-history likelihood is important for this model. Many of the reported results in figures 2 and 3 are on the statistics of the spiking. Are these figures using the likelihood with the spike-history filter? I am assuming they are not, as the spike-history filter is explicitly specified as the likelihood used for figure 4, but not in figures 2 and 3. If they are not used here, why not? Further, can the authors demonstrate that this spike-history filter leads to a different scientific result? I.e. are the conditioned latents any more accurate or do they provide any more insight if one likelihood is used compared to another. The evaluations in the supplement demonstrate that the spiking statistics is better captured with this more sophisticated likelihood, but further discussion as to why this is important for scientific insight would be appreciated. 

Overall, the authors primarily use spiking statistics as their measure of accuracy in this model throughout the paper. While the ability of a model to capture single-level neural statistics is important, it seems to me that the core use this model would have to an experimental practitioner is to visualize latents in an informative way in trial-varying data, potentially conditioned on behavior. Because of this, it would be nice to see more evaluations of the latent space, and not focus as much on the spiking statistics. Similar to the point above, it would be interesting to take a couple existing autoencoding LVMs used in neuroscience (like LFADs and others) and augment them with the same history-based Poisson likelihood. Then we would be able to see more clearly the contributions of the other features of the model such as S4 and diffusion, which I think are quite interesting, and their roles in identifying interpretable scientific latents where other models fail or cannot do so, such as in the trial-varying case or conditioned on behavioral covariates. 

Lastly, there are many models that use autoencoder based approaches other than LFADS that could be compared to here but are not. It is of course not possible to compare to all of them, but some more thorough evaluation and discussion alongside existing approaches I think would significantly improve the paper. See for example, *Poisson Interpretable VAE, **Targeted neural dynamical modeling, ***PSID and Duncker and Sahani 2018. LFADs alone, especially one without the same history-dependent observation likelihood, does not provide a complete picture of the capabilities of latent identification of this model in relationship to other other approaches. 

*) "Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE"
**) "Targeted Neural Dynamical Modeling"
***) "Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification"
questions: questions: Are the results reported in figures 2 and 3 using the spike history filter or just Poisson likelihoods?

What is the 'kde' in figure 3d?
limitations: limitations: See above. Further discussion and potential evaluation of this model alongside other approaches, and the dissociating the particular role of the likelihood as compared to the other model features, especially concerning scientific purposes, would be helpful for this paper.
paper_summary: paper_summary: This paper proposes a new generative model for neural spiking datasets. The model consists of a deterministic, deep SSM (S4) autoencoder paired with a diffusion model of the learned autoencoder latent sequences. This enables generating accurate neural time series traces across variable length trials lasting up to 10ms and conditional generation given behavioral covariates. The approach is applied to a synthetic dataset and two different neural datasets, where the authors investigate a variety of uses of the model.
strength_weakness: strength_weakness: While much of the modeling approach is well-motivated, I do not find that to be the case for the specific implementation of the spike-history component. Additionally, it appears that the spike-history component is responsible for much of the improvement over LFADS, and not the alternative underlying generative model that is the primary novel contribution of this paper (one could imagine also training LFADS with a spike history filter). In particular, it is not well-motivated why the filters are trained post-hoc and why the softplus approximation is preferred over using the exponential function.
questions: questions: - Why is the spike-history filter trained posthoc instead of during the autoencoder training? Additionally, can the authors provide more details about why they choice the softplus approximation? I do not necessarily agree that the approximation is accurate, and its not clear why this approximation is even necessary during training or why it is used over a numerically stable exponential function. Is this primarily an issue when the model is run in generative mode? 

- The LDNS overestimates some of the pairwise neural correlations in the attempted speech dataset. Is it typical for LDNS to overestimate pairwise correlations? Could this be improved by increasing the latent dimensionality of the model, to allow for more uncorrelated latent dimensions?

- Could the authors comment on the use of S4 as compared to other deep state space models like Mamba? Does one appear to work better than the other?

- Have the authors considered using the approach to do conditional generation of the attempted speech neural recordings given the cued sentence? This seems like a challenging but direct application of the proposed method.
limitations: limitations: Yes.
paper_summary: paper_summary: This paper introduces the Latent Diffusion Model for Neural Spiking data (LDNS). LDNS combines the capacity of autoencoders to extract low-dimensional representations of discrete neural population activity with the capability of denoising diffusion probabilistic models to generate realistic neural spiking data. It achieves this by modeling the inferred low-dimensional continuous representations. Through experiments on three different datasets, LDNS has been proven to achieve low-dimensional latent variable inference and realistic conditional generation of neural spiking datasets, providing possibilities for simulating experimentally testable hypotheses.
strength_weakness: strength_weakness: 1. There is only one comparison method in this paper.
questions: questions: 1. What is the significance of studying the generation of spiking data? Is this generated data of practical value?
2. Have you tried to compare it with other VAE-based methods? Despite the claim of the authors that LFADS is the most successful VAE-based method available, it is clear that no method performs best in all tasks, especially in more cutting-edge tasks like this paper.
3. Can you explain how the stochastic operations in Figure 1 are performed?
4. In the conditional generation experiment of neural activity given the reach direction, the experiment is conducted under conditions of only two directions. Is it possible to conduct the experiment under conditions of three directions?
limitations: limitations: This paper still has certain limitations. Firstly, the exploration of LDNS in simulating neural activity is currently confined to the abdominal cortex during speech, which leaves the simulation of neural activities under more complex behavioral patterns in other parts of the body unexplored. Secondly, the processing of human-related data in the research raises concerns about privacy protection. In response to these limitations, the following suggestions are offered:
1. This paper has initiated an investigation into the effects of LDNS in simulating neural activities in the abdominal cortex during human speech. To further verify its potential, future research could expand to simulate cortical neural activities in other parts of the body under more complex behavioral patterns, aiming for a comprehensive evaluation of LDNS's applicability and efficacy in the field of neuroscience.
2. Regarding the use of human-related data in the paper, it is suggested that in subsequent research, greater emphasis should be placed on protecting data privacy to avoid negative social impacts.
paper_summary: paper_summary: In their study "Latent Diffusion for Neural Spiking Data", the authors introduce the titular LDNS model for generating realistic neural population activity, and apply it to three datasets: a synthetic dataset with true latents generated from the three-dimensional Lorenz system, and two previously published neuroscientific datasets. 

The LDNS model is an instance of the latent diffusion model (LDM) introduced by Rombach et al. (2022), which is adapted here by the authors to support generation of variable-length discrete-valued time series. The autoencoder of the LDM in this case is trained with a Poisson loss to account for the discrete nature of neural spiking data. The diffusion model of the LDM operating in latent space here is constructed with layers of structured state space sequence (S4) models of Gu et al. (2021). The diffusion model can be trained to generate samples conditional on certain external covariates by conditioning the reverse mode on the covariates. The split of LDMs between the autoencoder stage and the diffusion stage is particularly convenient for the task of generating spike trains, as it avoids any adjustment to the diffusion model for generating discrete variables. A variant of the LDNS model has additional parameters trained post-hoc to incorporate the spike-time history of individual output neurons to account for refractoriness and other single-cell auto-regressive features. 

According to the considered neural population statistics, the LDNS model gives an excellent fit to the synthetic data generated from the three-dimensional Lorenz system. For the two neuroscientific datasets, the LDNS performs solid in terms of matching the considered neural population statistics, outperforming the important reference model LFADS (Sussillo, 2016) in one case. On the first neuroscientific dataset comprised of human neural population data sampled during attempted speech, the LDNS is shown to handle variable sequence lengths in the training data and during sampling. The second neuroscientific dataset consists of monkey neural population data recorded during a maze reach task. Here, the LDNS variant with additional spike-history parameters fitted post-hoc is shown to also capture the temporal auto-correlation structure of neural subpopulations. The authors furthermore showcase the ability to generate samples conditionally on experimental covariates, by conditioning sampling on the initial reach angle, respectively the full velocity trajectory.
strength_weakness: strength_weakness: The authors have adapted the latent diffusion model to the task of generating neural population activity. While I believe they did good work on that, I am less convinced of the evaluation and comparison of their model, which is important to judge its overall usefulness. 

For a generative model of multivariate time-series, I find the evaluation somewhat lacking. I understand that judging spike train generation is still a much harder task than judging image generation, and that one could argue that to date we don't understand all the relevant aspects of neural population activity. But firing rates, pairwise correlations and population spike counts are features that don't take dynamics into account. Average interspike interval and standard deviations of interspike intervals are still very local temporal features. Computational neuroscience has come up with a host of methods for analysing neural population dynamics -- in particular those with low-dimensional structure as assumed by the LDNS model -- which could be used here for comparing the sampled spike trains against the data. The authors already cited both LFADS and Gaussian Process Factor Analysis, but did not compare e.g. the low-dimensional trajectories extracted from their sampled data against the low-dim. trajectories extracted from the neural recordings. If nothing else, a supplementary figure or two with a handful of additional sampled spike trains against real recordings could help readers judge if the model captures the overall neural dynamics of the data. 

Another aspect that makes it difficult to evaluate the quality of their suggested model is the lack of comparison against direct competitors. The authors do an excellent job of explaining LFADS and its relevance as a comparison, but from what I understand only compare against LFADS on one of the three numerical experiments (I take that figure A4 which isn't mentioned in the text is also about the monkey data?). I understand that autoLFADS is significantly more computationally expensive to run than the LDNS model,  (which I would encourage the authors to state more prominently also in the main text!) but especially for the second experiment with human data, the LDNS model seems to perform the worst of the three experiments, and I could find no alternative model for comparison.
questions: questions: As explained above, analysis of the full population dynamics of their sampled neural population data would be very interesting. The authors could, but do not have to follow my suggestion of analysing the sampled neural population with one of the more commonly used dynamical dimensionality reduction methods. 

Can we get some comments or additional analysis of the LDNS model in cases where it doesn't work perfectly? There are several interesting supplementary figures for the Lorenz synthetic data, but for the other two experiments I am left wondering whether the non-negligible model errors arise primarily from the autoencoder or from the diffusion model. Those are the cases where I would most appreciate a comparison of the pairwise correlations from autoencoder against those of the full LDNS model, as in Fig A7 for the Lorenz data.
limitations: limitations: I think the authors have adequately addressed limitations and possible negative societal impacts of their work.

Comments: comment: Thank you for the positive response, detailed engagement with our work, and further increasing the score.
I thank the authors for their in-depth response and the additional analyses they added to address my concerns. 
I am happy to raise my rating by a point.
Thank you for taking the time to conduct our ethics review.
Thank you for taking the time to conduct our ethics review.
Thank you for the positive response, detailed engagement with our work, and further increasing the score. 

We really appreciate the thorough questions about the spike history term and agree it would be interesting to see the effect of different nonlinearities in future work and how joint training would interact with coordinated dropout and alternative regularization schemes. Thank you for pointing us toward this paper; combining such approaches is indeed very promising.
Thank you so much for the prompt and kind response, and now recommending acceptance. 

We appreciate your thorough engagement with our work and detailed comments, which allowed us to clarify our contributions. We agree that adding these analyses will strengthen the paper.
Thank you for your thorough response and for running the additional experiments. I especially appreciate the new comparisons of the spike history filter with LFADS, and for the additional baselines and visualizations. Overall, I have decided to increase my score and recommend accepting this paper.

Nonetheless, I still remain unconvinced that the spike-history filter should be fit post-hoc and find the need for the softplus approximation unsatisfactory. Perhaps this implies the model should use a softplus nonlinearity rather than exponential during the initial training phase. Alternatively, joint training may alleviate the stability issues. 

To fit the model jointly with the spike history filter, I suggest using alternative masking schemes -- see e.g. a very recent paper [1] with alternative options that should work with the spike history filter. 

[1] Towards a "universal translator" for neural dynamics at single-cell, single-spike resolution. Zhang et al., arXiv:2407.14668
This is a very thorough and impressive response. I think these additional figures paint a much more complete picture of the model's capabilities in a way that now clarifies it's disparate contributions compared to competing approaches. I believe these will substantially increase the work's impact.  

Particularly, R2 and R3 clearly demonstrate the LDNS addresses something that is lacking in modern generative LVMs for neural data, and I believe they should be highlighted in the main manuscript. R1 also adds important baselines that would be nice to prominently highlight as well (particularly including the LFADS sh).  The additional latent space characterizations (R6-R8) helps future model practitioners get an idea of the model's scientific utility. It would be nice if possible to include some of this in the main body of the paper, but I don't believe it is necessary.

I have updated my score accordingly and am excited to hopefully see this work at this year's NeurIPs.