Latent Diffusion for Neural Spiking Data
Authors: Jaivardhan Kapoor
Auguste Schulz
Julius Vetter
Felix C Pei
Richard Gao
Jakob H. Macke
Venue: NeurIPS 2024 spotlight (2024)
Keywords: neural population
diffusion models
latent variable models
electrophysiology
brain-computer interfaces
Primary area: neuroscience_and_cognitive_science
TL;DR: We propose a generative model for (conditional) generation of neural spiking activity using latent diffusion.

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.

Latent Diffusion for Neural Spiking Data
Jaivardhan Kapoor'
Auguste Schulz!
!*
Julius Vetter!
Felix Pei
Richard Gao!! f
Jakob H. Macke
Machine Learning in Science, University of Tübingen & Tübingen AI Center, Tübingen, Germany DDpartmeent Empirical Inference, Max Planck Institute for Intelligent Systems, Tübingen, Germany 'Equal contribution, order determined by coin toss. tEqual supervision. (firstname. laseeellstaaann
de)
Abstract
Modern datasets in neuroscience enable unprecedented inquiries into the relation- ship between complex behaviors and the activity of many simultancously recorded neurons. While latent variable models can successfully extract low-dimensional embeddings from such recordings, using them to generate realistic spiking data, es- pecially in a behavior-dependent manner, still poses a challenge. Here, we present Latent Diffusion for Neural Spiking data (LDNS), 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 to sample neural activity with realistic single-neurd 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 equip LDNS with pressive observation model that accounts for single-neuron dynamics not mediated by the latent state, further increasing the realism of generated samples. Finally, con- ditional LDNS trained on motor cortical activity during diverse reaching behaviors can generate realistic spiking data given reach direction unseen reach trajectories. In summary, LDNS simultaneously enables inference of low- oow-dimesiional latents and realistic conditional generation of neural spiking datasets, opening up further possibilities for simulating experimentally testable hypotheses.
1 Introduction
Modern datasets in neuroscience are becoming increasingly high-dimensional with fast-paced in- novations in measurement technology [I
granting access to hundreds thousands of simultaneously recorded neurons. At the same time, the types of animal behaviors and sensory stimuli naturalistic and complex, resulting in experimental setups with
48 23
under investigation have become heterogeneous trials of varying length, or lacking trial structure altogether (35] 33 57 Therefore.
key target in systems neuroscience has shifted towards understanding the relationship between high-dimensional neural activity and complex behaviors.
For -igh-dimensional neural recordings, analyses that infer low-dimensional structures have been very useful for making sense of such data [[]. For example, latent variable models (LVMs) are often used identify neural population dynamics not apparent the level of single neurons [6] 32 40)
More recently, deep learning-band basedd approaches based on variational autoencoders (VAEs) 26 44
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

51 36 63
have become particularly popular for inferring latent neural representations due to their expressiveness and ability scale to large, heterogeneous neural recordings with behavioral covariates 33 63 [6].
22 6 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 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.
Diffusion models 44 20 50) have been highly successful for conditional and unconditional data generation in sev- eral domains, including images [20]. molecules [59], and audio spectrograms 27 and have demonstrated sampling- fidelity that outperforms that of VAEs and GANs 20). A key strength of dif- fusion models that makes them particu- larly attractive in the context of model- ing neural datasets is the ability to flex- ibly condition the generation various (potentially complex) covariates, such as to simulate neural activity given certain behaviors. Recently, diffusion models have been extended to continuous neu- ral time series such as local field poten- tials (LFPs) and clectroenccphalograaty
autoencoder
spikes
Ialents
stochaelic
spikes
dicaner
affirion
evalustion and tasks
(optional) behovioral conditoring
noise
diffusion model
Figure Latent Diffusion for Neural Spiking data. 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.
(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 [8] 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 $4 layers as 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, neural spiking data, there are prominent statistical dependencies that ought to persist conditional on the latent state, e.g., single-neuron dynamies 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, 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 small additional computational cost.
by equipping it with an expressive observation model [41 32 (54 62

---

Main contributions In summary, LDNS is a flexible method that allows for both high-fideling diffusion-basee sampling of neural population activity and access time-aligned -iwdeinnsionall representations, which we validate on a synthetic dataset. Next, we show the utility and flexibility of this approach on complex real datasets: First, L.DNS 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 full reach trajectories (time series), including unseen behaviors that are then accurately decoded from the simulated neural data. Overall, LDNS enables simultancous inference of Now-dimensiones latent representations for single-trial data interpretation and high-fidelity diffusion- (conditional) generation of diverse neural spiking datasets, which will allow for closed-looy loop in silico experiments and hypothesis testing.
2 Methods
2.1 Latent Diffusion for Neural Spiking Data (LDNS)
We consider dataset recorded from a population of n neurons, consisting of trials with spiking data E NoxT (sorted into bins of fixed length resulting in spike counts over time), and optional simultaneously recorded behavioral covariates y E R" (that can also be time-varying y E B=xT). A dataset of M such trials D can be written as D = {s(()) .y("), possibly with varying trial lengths TM- We make the assumption that a large fraction of the variability in this dataset can be captured with few underlying latent variables
RR where d <
Our goal is to generate realistic spiking data s* that faithfully capture both population-leve level and single- neuron dynamics of St... with the ability to optionally condition the generation on behavior ycond- To this end, we propose a new method, LDNS, that combines the strength of neural dimensionality reduction approaches with that of diffusion-bian based generation.
LDNS uses two-stage training framework, adopted from the highly successful family of latent diffusion models (LDMs) 45 59). To train LDNS, we first train regularized autoencoder [14] to compress the spiking data into a now-dimensional continuous latent space (Fig. Concretely, we focus on two objects of interest for the LDNS autoencoder: (1) inferring a time-aligned, low- dimensional smooth representation z E R that preserves the shared variability of the spiking data, and (2) predicting smooth firing rates À that are most likely give rise to the observed spiking data.
In the second stage, we train diffusion model in latent space, possibly employing conditioning to make generation contingent on external (e.g.. behavioral) covariates (Fig. D For the diffusion model, our main objective is the generation of zz E <R<x 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. $4 layers consist of state-space transition matrices that can be unrolled into arbitrary-length convolution kemels,, allowing sequence modeling of varying lengths. For details on network architectures and $4 layers, see appendix A1
2.2 Regularized autoencoder for neural spiking data
For the spiking data, choose Poisson observation model, and train autoencoders by minimizing the Poisson negative log-likelihood of the input spikes s given the predicted rates À = (zcoder(z). To enforce smoothness in the latent space, where encoder(s), we add an La regularization along with a temporal smoothness regularizer over z, resulting in the combined loss
T. (Ai(t) - s. (t) (t) +31 | +B2
K.T
[2(1|) - 2(1 - 12
CAB E
(I)
(1 + k)
Poisson NLL
La reg.
temporal smoothness
Code available https:/ggthhb.ccm//aacckla/////d
3

---

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 [2], 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; [1] 32), resulting in the observation model
T
log (e(t) + h4, +
Σ si () ) T=1
st) ~ exp
where 7" corresponds to the time-lagged window length. This modification is learned post hoc, and the parameters h; are fit with a maximun-likelihodd 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 accurately capture single-neurg autocorrelation structures (Sec. 3.4).
2.3 Denoising Diffusion Probabilistic Models
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 autoeneoder-derivee latents of each trial, and optionally, additional conditioning information such as the corresponding behavior, ie., Dz = {z = encoder(s()) (.)).
Diffusion models aim to approximate the data distribution (zz) through an iterative denois- ing process starting from standard Gaussian noise. For latent z (denoted as Zu for diffusion timestep 0), we first produce noised version at step t by adding Gaussian noise of the form q(2,)2 (2)) = . ( (yā,zo, (1 -))))) 7). Here, a = I Okk, where the noise scaling factors follow a fixed linear schedule. We then train neural network to approximate the reverse pro- cess po (2,-1/2,) (2.) for each diffusion timestep. The true (denoising) reverse transition g( (z,) is intractable however, we can apply variational inference learn the conditional reverse transition 4(2;- |2, 2o). which has a closed form written
or
q(2-11z4, 20) =N Vara(1 (i:-1)
Vā,-i(1 ar) 1 āt 2o:
(1- (1- -111-))))7)
1 24
(3)
We train the neural network 14g (z,2), to approximate the mean of this distribution by optimizing the -olla. where eo is the noise used to generate 2, from %o, and c0 ()
loss Fxo=-D..o..:lleo the equivalent reparameterization for ss (z, L). test time, we sequentially sample 2 given z; using the learned transition ps (z(-1//zz) starting from standard Gaussian noise. Using S4 layers in the denoising network allows us to generate latents with varying lengths. This 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 po (z Ze, y). Details on the conditioning mechanisms appendix A1.4
Experiments and Results
3.1 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 low-dimensional, non-linear dynamical system commonly used neuroscience [7] (6). We simulate rates as an affine
mapping from the 3-dimensional system to 128-dimensionall neural space, and sample from

---

a
chactic Lorenz
aufcencoder prediction
C
difiusion sampled latents of varying length
distribution match
log pod latents
rato: spikes
diffusion
10
10
time
100
0
time
100
freq.
synthetic spikes
d
population level statistics
single neuron statistics
128
0.5
15
Idns meen
neuron id
dota kde Ions lide
Idns
Iche
0.0
5
258
0
100
200
300
0.0
10
15
spike counts
gl painwise correlation
gl mean
gt
Figure 2: Realistic generation of spiking data with underlying chaotic dynamics. a) Synthetic spiking data from an underlying Lorenz system with a Poisson observation model. b) Accurate, smooth rate predictions of the autoencoder for held-out spiking data. e) Plotted trace of sampled latents (256 bins training length, left) and 16x the original training length (middle). The sampled latent distribution matches the PSD of the autoencoder latents (right; median, 10%, and 90% percentiles). d) LDNS population spike count histogram (kde: kernel density estimate) and pairwise cross- coss-correlations match the training distribution. LDNS single neuron statistics, i.e., mean inter-spike interval (isi) and std isi, match the training distribution.
Poisson distribution to generate spiking data. Next, we showease the applicability of LDNS on two neural datasets: We apply our method on 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 delayed center-out reach task with barriers [8 The multi-modal nature of the dataset allows us to assess both unconditional as well 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 Poisson observation model as well as a spike history-dependent autoregressive observation model. For all other experiments, we only train Poisson observation model.
Baselines We compare LDNS to the most commonly known VAE-based latent variable model:
Latent Factor Analysis via Dynamical Systems (LFADS [5]] 36 47). which has been shown to outperform various classical latent variable models on a variety of tasks ([8), details in appendix [4]]. To ensure that we use optimal hyperparameters for LFADS, we follow the auto-ML pipeline proposed by Keshtkaran 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 (appendix (A5).
recordings, we further compared additional VAF baselines 21 62
Metrics For all experiments, we assess how well LDNS- generated samples match the spiking data statistics by computing 1) the
in the training distribution. Concretely, we compare population- distribution over the population spike count, which sums up all spikes co-occurring in the population in 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 werr--pikk-iitevve distribution for each neuron (mean isi and std isi). When multiple spikes occur in 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.
3.2 LDNS captures the true spiking data distribution with an underlying Lorenz system
We simulate trials of length 256 timesteps from the three-diment dinensiamelional (chaotic) Lorenz system (Fig. 24). The regularized autoencoder extracts smooth latent time series (cight latent dimensions)
5

---

a
ettempled speech lask
b
variable-enggt trials data
population and single-urit stalistics
0.10 maan
0.15
0.06
15
05 0.10 dota nean ki
0
0.15
pop. spike count
data dtd
Idns samples
pairvise correlations
Uni
0.0-
2
trial length (8)
10
dide
kira
0.0
0.5
Figure 3: Unconditional generation of variable-length trials of human spiking data during attempted speech. Multi-unit activity is recorded from speech production- regions of the brain (top) during attempted vocalization of variable-length sentences (bottom). b) Neural activity during sentences of different lengths. c) LDNS unconditionally sampled trials with different lengths, using the Poisson observation model. d) LDNS population spike count histogram, and mean and std of the isi match those of the data. e) Correlation matrices of the data (left) and LDNS samples (middle), and scatterplot of the pairwise correlations of data vs. LDNS samples (right).
from the 128-dimensional spiking data, resulting in smooth firing rate predictions that closely match the ground-truth rates (Fig. Zp. Supp. Fig. A3). We then train diffusion model on the extracted autoencoder latents. Latents sampled from the diffusion model (red) preserve the attractor geometry of the Lorenz system (Fig. Zp. 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 time segments of 256-time steps, can sample and successfully generate latent trajectories that are much longer, but still accurately reflect the Lorenz dynamics (Fig. 2 middle, 16x 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. A3).
To assess the sampling fidelity of the generated synthetic neural activity, we compute variety of spike statistics frequently used in neuroscience. LDNS captures both population- statistics, such as the population spike count histogram and pairwise correlations between neurons (Fig. 1)), as well as single-neuron statistics, quantified by the mean and standard deviation of inter-spike-intes 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 statisties of the ground-truth synthetic data.
3.3 Modeling variable-length trials of neural activity recorded in human cortex
Next, we assess whether LDNS is capable of capturing real electrophysiological data, applying to neural recordings from human cortex during attempted speech (Fig. (3.. top, Willett al. (56). A participant with a degenerative disease who is unable to produce intelligible speech attempts to vocalize sentences prompted on sereen, while neural population activity is recorded from the ventral premotor cortex. Since there is a large variation in the length of prompted sentences (3, bottom), this dataset allows us evaluate the performance of LDNS real data in naturalistic settings with variable-length and highly heterogeneous dynamics.
To account for varying trial length during autoencoder training, pad all trials to 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 binary mask as a conditioning variable.
This approach allows us to infer time-all aligned latents underlying the cortical activity of the participants, compressing the population activity by 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. D.c. additional samples in Supp. Fig. A7p. This is reflected in closely matched population spike count histograms (Fig. 31, left), and single neuron statistics such as mean and standard deviation of the inter-spike interval (Fig. H. right). Additionally, real

---

mace reach task
b
dataa
lons semple
population level stabisicc
d11
ione
data kde
100
Iona kdo
-0.00
0.5
0.5
-0.05
time
time
spike bins
data painwise correlation
single neuron
temporal auto-correlation
Idha mean ta
0.2
Idns Polsson
data 5
Knn 166
Idne Paissonn
with spike history
daaa 5 wh
Idne esd
0. 0.0000
0.1
-0 00.5
0.1
0.2
0.1
0.2
-126
125
-125
125
mean
lag (ms)
lag (ms)
Figure 4: Realistic generation of spiking data in monkey performing reach fasks. a) monkey performs diverse reach movements in different mazes. b) Neural activity during a reach trial and sampled trial from LDNS with a Poisson observation model. с) The LDNS population spike count histogram, and pairwise correlations match those of the data. d) LDNS mean- and std isi match the monkey data distribution. e) Auto-correlation of data, LDNS samples with Poisson observations (left), and LDNS samples with spike history, grouped according to correlation strength.
and LDNNS-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 the data (Fig. 33,, Pearson correlation coefficient 0.47), and 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 fixed trial structure.
3.4 Realistic generation of spiking data from monkey performing reach tasks
We further evaluate LDNS in a different setting by applying model sparse spiking data recorded from a monkey performing a reaching task constrained by barriers that form maze (Fig.4.. left). The variety of different maze architectures leads diverse reach movements of both curved and straight reaches (Fig. Pa, 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. 4p. additional samples in Supp. Fig. All), and closely match population-level spike statistics (Fig. 4b). Single
neuron statistics in this low spike count regime (a maximum of three spikes neuron in 5 ms bins)
are also captured well (Fig. fi), and are on par with or better than LFADS [36 25 (see Table [ for
summary of main comparisons, and appendix AS 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, L.DNS can generate spiking data that is faithful at the level of both single-neuron and population dynamics.
Table 1: Model metries comparison. D for the population spike count histogram and RMSE comparisons. Mean and standard deviation across folds sampled with replacement. sh represents observation models with spike history. Bolded entries represent best-performing values for Poisson and spike-history observation models.
Method
DEL psch
RMSE pairwise
RMSE mean isi
RMSE std isi
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.le-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 6.2c-4
0.0025 1.07e-5
0.024 = 0.002
0.023 + 0.001
7

---

b
behavioral conditioning
decoded
conditioning
%%% conditioning
25
data
condttin samples
predicted behavior
rdgge wook!
y possson
0 condition
diffusion
-25
25
-25
25
-25
0
25
position
pasition
x position
g
straight test conditions
latent dim O
latent dim
PC1
PC1:PC2
20
tait reaches
sampled latent
PC of Intents
positinn
-20
x tion
[8]
[6]
PC1
Figure 5: Generation conditioned on monkey reach directions and velocity traces. a) Closed loop assessment: do conditionally generated latents translate to neural activity consistent with the desired direction reach movement? b) Unseen reach movements (data) and corresponding movements decoded from the rates predicted by the autoencoder (ae). e) Decoded reach directions LDNS samples conditioned conditioned on 3 unseen reach movements (velocities Uz, Uy). e) Straight reaches from the test set used for velocity conditioning. f) LDNS sampled latents conditioned on trajectories shown in e) vary smoothly over time and reflect information about reach angles. g) PCs of sampled LDNS latents shown in f) reveal meaningful and separable information about behavior.
initial reach angles 6. d) Decoded reach directions of LDNS samples
Both LFADS and the LDNS autoencoder are optimized by maximizing the Poisson log-likelihood, and thus cannot capture single-neuron dynamics such as refractoriness [1 (55). which can have strong influence on the observed autocorrelation structure. Given the overall sparsity of the spiking data and resulting low correlations (Supp. Fig. ATO, focus on the temporal structure of auto- correlations averaged within groups of neurons (z 45 neurons per group) split by their instantaneous correlation strength [3]] darker colors correspond
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 Poisson observation model (red). As expected, LDNS with Poisson observations is unable to capture the dip in the data auto-correlation at ms lags (one time bin) (Fig. f.. 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 the Poisson samples, autoregressive samples can capture this aspect of neural spiking data very accurately while also improving the overall fit the empirical coteeccorrlltion (Fig. ec. right). Moreover, the post hoc optmizatiion of these filters also improves modeling of other single-neuron, as well as population-Ievel statistics, such the population spike count histogram or the mean of the isi (Table 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, extend AutoLFADS with spike history dependence (LFADSsh), improving its performance across metrics. The augmented LFADSsh also captures the dip in autocorrelation ms lags (Supp. Fig. A12). Still, in both observation model variants, LDNS maintains superior or comparable performance (Table (.)
Thus, Poisson LDNS allows for the generation of spiking data that is par better in terms of sampling fidelity than previous approaches. Incorporating spike-histong dependence and sampling spikes autoregressively allows to further increase the realism of generated spike trains, leading to large improvement on several of the considered metrics.
3.5 Conditional generation of neural activity given reach directions or velocity profiles
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 (p). We first validate that the autoencoder predicts firing rates that allow to linearly decode the behavior
8

---

following the ridge-reg regge-regession approach proposed in [38.. Decoded behavior from autoencoder reconstructed rates matches the true trajectories of unseen test trials (Fig. (Fig.Ep).
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 freach- Indeed, from the generated samples of neural activity, we can decode-using the same linear realstr-rcalistic reach kinematics that are consistent with the conditioning angle Greach and overall reach kinematics (Fig. 55)). 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 diffusion model the same perteencoderinfferre
latents but now condition on entire velocity traces (Fig. 5td. Velocity- conditioned LDDSS is able 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. 5 Individual samples of latent trajectories vary smoothly within trial (Fig. 5D, while reach direction
varies smoothly across samples in the first principal component (PC1) of the latents (Fig. [3g. left).
Projection onto the first two PCs of latent trajectories shows clear clustering by reach direction (Fig. 5
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 condition- ing. Furthermore, LDNS learns an interpretable latent space with behaviorally-relevan structure.
4 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 [6] 32 39 58
62
13
28,
63) (see (38 for a comprehensive list.) In addition
to capturing shared population-Ievel dynamies and dependence on external stimuli 4 LVMs have
or even direct neural interactions (44) While these methods often have useful inductive biases (e.g.. linear
been extended to allow autoregressive neuron-level (non-Poisson) dynamics 33
13 62
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 cow-dimensional dynamics underlying nigh-dimensional discrete data 663 16 46). especially when combined with nonlinear recurrent neural networks 21) VAES have been used infer identifiable low-dimensional latent representations
conditioned on behavior 63 21 and have incorporated smoothness priors using Gaussian Processes
to regularize the latent space [6.. 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 (2). 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 of for high- gig-fiiceiity generation in several domains 200 277 including continuous- valued neural signals such as EEG [53], as well as in time series forecasting and imputation tasks 52 43 Similar to the LDNS architecture, Alcaraz and Strodthoff (2] also use an $4-based denoiser for imputation. More specifically, LDNS is inspired by latent diffusion models [5] 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, 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 simultancously 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 maze. Through 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 autocncoder, 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-basee training to optimize hyperparameters (appendix A3.1p. 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-leve shared dynamics. Third, as a result of the generaaaaeel 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 behavior.
Altogether, these features enable LDNS to generate realistic neural activity, especially when con- ditioned 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 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 still able to model the data well under relatively severe compression (e.g.. 182-to-16 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.
Acknowledgements
This work was supported by the German Research Foundation (DFG) through Germany's 's Excellence Strategy (EXC-Number 2064/1, PN 390727645) and SF31 (PN 276693517), SFB 1089 (PN 227953431), SPP2041 (PN 34721065), the German Federal Ministry of Education and Research (Tübingen AI Center, FKZ: 011818039), the Human Frontier Science Program (HFSP), and the European Union (ERC, DeepCoMech Tome, 101089288). We utilized the Tübingen Machine Learning Cloud, supported by DFG FKZ. INST 37/1057-1 FUGG. JK, AS, and JV are members of the International Max Planck Research School for Intelligent Systems (IMPRS-IS) and JV is supported
10

---

by the AI4Med-Bil BW graduate program. We thank Chethan Pandarinath for providing access to their compute cluster to train FADD. 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.
References
[1] Misha B, Ahrens, Michael Orger. Drew N. Robson, Jennifer M. Li. Philipp Keller. Whole-brain functional imaging at cellular resolution using light-sheet microscopy. Nature Methods, May 2013.
[2] Juan Lopez Alcaraz and Nils Strodthoff. Diffusion-hase hased time series imputation and forecasting with structured state space models. Transaction on Machine Learning Research, 2022.
[3] Antonis Antoniades, Yiyi Yu, Joe S. Canzano, William Yang Wang, and Spencer Smith. Neuroformer: Multimodal and Multitask Generative Pretraining for Brain Data. In International Conference on Learning Representations, October 2024.
[4] Evan W Archer, Urs Koster, Jonathan Pillow, and Jakob Macke. Low-dimensional models neural population activity in sensory cortical circuits. Advances in neural information processing systems, 27, 2014.
[5] Mchdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards, Matthew Perich, Guillaume Lajoic, and Eva Dyer. A Unified, Scalable Framework for Neural Population Decoding. Advances in Neural Information Processing Syxtems, 2023.
[6] Mohammad Bashiri, Edgar Walker, Konstantin- Monstantin-Klemens Lurz, Akshay Jagadish, Taliah Muhammad, Zhiwei Ding, Zhuokun Ding, Andreas Tolias, and Fabian Sinz. A flow-based latent state generative model of neural population responses to natural images. In Advances in Neural Information Processing Systems, 2021.
[7] Manuel Brenner, Florian Hess, Georgia Koppe, and Daniel Durstewitz. Integrating multimodal data for joint generative modeling complex dynamics, 2024.
[8] Zijiao Chen, Jiaxin Qing. Tiange Xiang, Wan Lin Yue, and Juan Helen Zhou. Seeing beyond the brain: Conditional diffusion model with sparse masked modeling for vision decoding. Computer Vision and Pattern Recognition, 2023.
[9] Mark Churchland and Matthew Kaufman. Memaze: macaque primary motor and dorsal premotor cortex spiking activity during delayed reaching (version 0.2201 13.0400) [data set], 2022. Version 0.2201 13.0400.
[10] Mark M. Churchland, John P. Cunningham, Matthew T. Kaufman, Justin D. Foster, Paul Nuyujukian, Stephen I. Ryu, and Krishna V. Shenoy, Neural population dynamics during reaching. Nature, 2012.
[11] John P Cunningham and Byron M Yu. Dimensionality reduction for large-scale neural recordings. Nature neuroscience, 2014.
[12] Kevin Doran, Marvin Seifert, Carola A. M. Yovanovich, and Tom Baden. Distance function for spike prediction, 2023.
[13] Lea Duncker and Mancesh Sahani. Temporal alignment and latent gaussian process factor inference in population spike trains. In Advances in Neural Information Processing Systems, 2018.
[14] Partha Ghosh, Mehdi S. M. Sajjadi, Antonio Vergari, Michael Black, and Bernhard Scholkopf. From variational to deterministic autoencoders. In International Conference Learning Representations, 2020.
[15] Karan Goel, Albert Gu, Chris Donahue, and Christopher R'c. It's raw! audio generation with statc-space models. International Conference on Machine Learning, 2022.
[16] Rabia Gondur, Usama Bin Sikandar, Evan Schaffer, Mikio Christian Aoi, and Stephen L. Keeley. Multi- modal Gaussian Process Variational Autoencoders for Neural and Behavioral Data. International Conference Learning Representations, 2024.

---

[17] lan Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. Advances in Neural Information Processing Systems, 2014.
[18] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces, International Conference on Learning Representations, 2022.
[19] Florian Hess, Zahra Monfared, Manuel Brenner, and Daniel Durstewitz. Generalized teacher forcing for learning chaotic dynamics. International Conference for Machine Learning. 2023.
[20] Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems, 2020.
[21] Cole Hurwitz, Akash Srivastava, Kai Xu, Justin Jude, Matthew Perich, Lee Miller, and Matthias Hennig. Targeted neural dynamical modeling, Advances in Neural Information Processing Systems, 2021.
[22] Kristopher Jensen, Ta- Fa-Chu Kao, Jasmine Stone, and Guillaume Hennequin. Scalable Bayesian GPFA with automatic relevance determination discrete noise models. In Advances Neural Information Processing Systems, 2021.
[23] James J. Jun, Nicholas A. Steinmetz, Joshua H. Siegle, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature, 2017.
[24] Mohammad Reza Keshtkaran and Chethan Pandarinath. Enabling hyperparameter optimization se- quential autoencoders for spiking neural data, Advances in neural information processing systems, 32, 2019.
[25] Mohammad Reza Keshtkaran, Andrew R. Sedler, Raeed H. Chowdhury, Raghav Tandon, Diya Basrai, Sarah L. Nguyen, Hansem Sohn, Mehrdad Jazayeri, Lee E. Miller, and Chethan Pandarinath. A large-scale neural network training framework for generalized estimation of single-trial population dynamics. Nature Methods, 2022.
[26] Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations, 2014. doi: 115350Warriv.3312.61121
[27] Zhifeng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: A versatile diffusion model for audio synthesis. International Conference on Learning Representations, 2021.
[28] Scott Linderman, Mathhw Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian Learning and Inference Recurrent Switching I Linear Dynamical Systems. Proceedings of the 20th International Conference on Artificial Intelligence and Statistics, 2017.
[29] Edward N. Lorenz. Deterministic Nonperiodic Flow. Journal of the Atmospheric Sciences, March 1963.
[30] Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Conference on Learning Representations, 2019.
[31] Tlya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. In International Conference on Learning Representations, 2019. doi: 1=Xiv:1711.0511
[32] Jakob H Macke, Lars Buesing, John Cunningham, Byron M Yu, Krishna Shenoy, and Maneesh Sahani. Empirical models of spiking in neural populations. In Advances in Neural Information Processing Systems, 2011.
[33] Bartul Mimica, Tuçe Tombaz, Claudia Battistin, Jingyi Guo Fuglstad, Benjamin A. Dunn, and Jonathan Whitlock. Behavioral decomposition reveals rich encoding structure employed across neocortex in rats. Nature Communications, 2023.
[34] Manuel Molano-Mazon, Arno Onken, Fugenio Piasini, and Stefano Panzeri. Synthesizing realistic neural population activity patterns using generative adversarial networks. International Conference Learning Representations, 2018.
[35] Joseph O' O'Doherty, Mariana M. B. Cardoso, Joseph G. Makin, and Philip Sabes, Nonhuman Primate Reaching with Multichannel Sensorimotor Cortex Electrophysiology.22 2017.
12

---

[36] Chethan Pandarinath. Daniel J. Ohea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen Ryu, Leigh R. Hochberg, Jaimie M. Henderson, Krishma V. Shenoy, I., F. Abbott, and David Sussillo. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature Methods, 15, 2018.
[37] William S. Peebles and Saining Xie. Scalable diffusion models with transformers. IEEE International Conference on Computer Vision, 2022.
[38] Felix C Pei, Joel Ye, David M. Zoltowski, Anqi Wu, Raced Hasan Chowdhury, Hansem Sohn, Joseph O' O'Doherty, Krishna V. Shenoy, Matthew Kaufman, Mark M Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan W. Pillow, Memming Park, Eva L Dyer, and Chethan Pandarinath. Neural latents benchmark '21: Evaluating latent variable models of neural population activity. In Neural Information Processing Systems, Dotasets and Benchmarks Track (Round 2), 2021.
[39] Biljana Petreska, Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. Dynamical segmentation single trials from population neural data, In Advances in Neural Information Processing Systems, 2011.
[40] David Pfau, Eftychios A Pnevmatikakis, and Liam Paninski. Robust learning of fow-dimensional dynamies from large neural ensembles. In Advances in Neural Information Processing Systems, 2013.
[41] Jonathan W. Pillow, Jonathon Shlens, Liam Paninski, Alexander Sher, Alan M. Litke, E. J. Chichilnisky, and Eero P. Simoncelli. Spatio-temporal correlations and visual signalling in a complete neuronal population. Nature, 454:995- 999, 2008.
[42] Poornima Ramesh, Mohamad Atayi, and Jakob H Macke, Adversarial training of neural encoding models population spike trains. Real Neurons & Hidden Units: Future directions of the intersection of neuvoscience and artificial intelligence @ NeurlPS 2019, 2019.
[43] Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. International Conference Machine Learning, 2021.
[44] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approxi- mate inference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning. 2014.
[45] Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High-resolution image synthesis with latent diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022.
[46] Auguste Schulz, Julius Vetter, Richard Gao, Daniel Morales, Victor Lobato-Rios, Pavan Ramdya, Pedro J. Gonçalves, and Jakob H. Macke. Modeling conditional distributions of neural and behavioral data with masked variational autoencoders. bioRxiv, 2024.
[47] Andrew R. Sedier and Chethan Pandarinath. latent factor analysis via dynamical systems, arXiv preprint arXiv:2309.01230, 2023.
A modular and extensible implementation of
[48] Nicholas James Sofroniew, Daniel Flickinger, Jonathan King, and Karel Svoboda. A large field of view two-photon mesoscope with subcellular resolution for in vivo imaging. eLife, 2016.
[49] Jascha Sohl-Dickstein. Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, 2015.
[50] Yang Song. Jascha Sohl-Dicksteen, Diederik Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-hased generative modeling through stochastic differential equations. International Conference on Learning Representations, 2021.
[51] David Sussillo, Rafal Jazefowicz, L. F. Abbott, and Chethan Pandarinath. Lfads latent factor analysis via dynamical systems, 2016.
[52] Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score-hased diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems, 2021.
13

---

[53] Julius Vetter, Jakob Н Macke, and Richard Gao. Generating realistic neurophysiological time series with denoising diffusion probabilistic models. bioRxiv, 2023.
[54] Michael Vidne, Yashar Ahmadian, Jonathon Shlens, Jonathan W. Pillow, Jayant Kulkarni, Alan M. Litke,
E. J. Chichilnisky, Eero Simoncelli, and Liam Paninski. Modeling the impact of common noise inputs the network activity of retinal ganglion cells. Journal of Computational Neuroxcience, 2012.
[55] Alison I. Weber and Jonathan W. Pillow, Capturing the Dynamical Repertoire of Single Neurons with Generalized Linear Models. Neural Computation, 2017.
[56] Francis Willett et Data for: A high-performance speech neuroprosthesis [dataset], 2023. URL /ttpa://doi. org/ 5061/dryad x69p8czpp
[57] Francis R. Willett, Erin M. Kunz, Chaofei Fan, Donald T. Avansino, Guy H. Wilson, Eun Young Choi, Foram Kamdar, Matthew F. Glasser, Leigh R. Hochberg, Shaul Druckmann, Krishna Shenoy, and Jaimie M. Henderson. A high-performance speech neuroprosthesis, Nature, 2023. Publisher: Nature Publishing Group.
[58] Anqi Wu, Nicholas A Roy. Stephen Keeley, and Jonathan W Pillow. Gaussian process based nonlinear latent structure discovery in multivariate spike train data. Advances in neural information processing systems, 30, 2017.
[59] Minkai Xu, Alexander Powers, R. Dror, Stefano Ermon, and Leskovee. Geometric latent diffusion models for molecule generation. International Conference on Machine Learning, 2023.
[60] Joel Ye, Jennifer Collinger, Leila Wehbe, and Robert Gaunt. Neural Data Transformer 2: Multi-context Pretraining for Neural Spiking Activity. Advances in Neural Information Processing Systems, 2023.
[61] Byron Yu, John Cunningham, Gopal Santhanam, Stephen I. Ryu, Krishna V. Shenoy, and Maneesh Sahani. Gaussian-Proodt Factor Analysis for lww-Dimensional Single-Trial Analysis of Neural Population Activity, Journal of Neurophysiology. 2009.
[62] Yuan Zhao and Il Memming Park. Variational latent gaussian process for recovering single-trial dynamics from population spike trains. Neural Computation, 2017.
[63] Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAr VAE. In Advances in Neural Information Processing Systems, 2020.
14

---

Appendix
A1
LDNS architecture
Here we describe the exact network components and architecture for the autoencoder and diffusion model.
A1.1 Structured State-Space Layers (S4)
central component of our autoencoder architecture is the recently introduced structured state space models (1)) With input sequence [x 101...27 E TR. and corresponding output y (V1... yTl E R..." an S4 layer applies the following operation for each timestep
As + Bag
(4)
3r Css.
where the discretized state and input matrices A, B given continuous analogues A, B and step size A are computed as
- A/2
(5)
B = (1- AAB
When the state 8, is not required, this recurrent computation of the output y given input sequence y can be unrolled into parallelizable convolution operation
y= = K * x, with unrolled kernel
(6)
K = (CB, CAB CA' B).
We used the S4 implementation provided by Gu al. [8] that stably initializes the state transition matrix A using a diagonal plus low-rank approximation. For a multivariate input-output pair x, y e RDXT, we apply D separate univariate S4 layers for each dimension and then mix them in the channel-mixing layer using an MIP (see next section). Each univariate input-output mapping consists of H separate $4 "head" that expanded and contracted from and into single dimension.
Due to its recurrent nature, $4 a causal layer, enabling variable-length training and inference, To enable bidirectionality, we flip the input signal in time, apply H 2 S4 heads each for the flipped and unflipped signal, and then combine these at the end into univariate signal, y. This allows bidirectional flow of information from front-to-back and back-to-front of the signal.
A1.2 Autoencoder
We include temporal information only in the encoder and model the decoder as lightweight pointwise MLP for the autoencoder (Supp. Fig. AT). This allows us to temporally align the latents with the signal, and ensure that no further temporal dynamics are introduced when mapping the latents back into Poisson rates.
We use causal S4 layers, allowing length generalization and handling of variable-length signals. During training, we pad the input spiking data with zeros into fixed length and only backpropagate through the unpadded output rates.
Furthermore, to infer smooth rates and avoid spiking behavior, we use coordinated dropout [2]]. For each time bin independently, up the remaining spikes by T-p (this preserves the firing statistics of the spiking data). We then backpropagate through the Poisson NLL loss only over the masked positions, effectively preventing the network from collapsing to spiking prediction of the Poisson rates.
mask the input spikes to zero with random probability p and scale
'Github link: http:///shhhllloooo..oo////////s
15

---

M(B.0.)
LayerNorm
MLP MLPo
AE block
AE block
MLP -
Pointwise MLP -c,
mangg (S4)
Layurhoor
Spikes
Latents
Rates
roiing (Lnear
S4 Encoder
Xptc.11
Kp.ca
conditioning
MLP
Timm maiing judtinational
t
MLPa
denoiser block
denoiser
MLP GG
Scole
block
Scale Shhit
Charral moing inner
Latents,
Latents,1
S4 Denoiser network
Figure A1: (Top left) The S4 autoencoder architecture. (Top right) Architecture for the autoencoder blocks used in the encoder. (Bottom left) The S4 diffusion model architecture. (Bottom right) Architecture for the diffusion blocks.
A1.3 Spike--bistor--augeented Poisson model
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 history length of 20, corresponding to 100 ms, using the AdamW optimizer (30) (learning rate 0.1, weight decay 0.01). log(1 + exp(x)) as
In our implementation, an approximation to the exponential function Eq. (2). which is accurate for the low-count regime while increasing numerical stability. During autoregressive sampling, 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.
use the Softplus function given by Softplus(z)
A1.4 Diffusion model
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. AT. For diffusion timestep and exed-length conditioning vector, shift and scale the inputs and outputs to the time mixing and channel mixing blocks using adaptive instance normalization, as done Peebles and Xie (37].
Unconditional generation for synthetic spiking data with Lorenz Dynamics and cortical spiking data in monkeys use only time conditioning.
2. Angle-conditioned generation for cortical monkey spiking data we add an embedding MLP([cos0, sin 6]) of the reach angle 0 to the timestep embedding output.
3. Trajectory conditioned generation for cortical monkey spiking data we concatenate the hand velocities vz, Vy of the monkey with the input as two additional channels.
4. Unconditional variable-length generation cortical human spiking data - we concatenate the desired length (with maximum sequence length of 512) as centered binary mask channel the input. We only backpropagate through the central section the output corresponding the binary mask.
16

---

We use DDPM scheduler with 1000 timesteps and e-araaeterization To stabilize and speed up training. we train all diffusion models using smooth L1 loss, written
L(x; x² (26) if [x| < å | - otherwise,
(7)
with ā = 0.05.
A2 Dataset access information
All real- world datasets used in this work publicly available under open-access licenses. Our work does not involve the collection of new experimental data.
The human BCI dataset is available at htttps//ddatarry: /ttp///aaaadaad....a/ wwrg//ttash//oomloods///l
/5477699 under a CCO 1.0 Universal Public Domain Dedication license. This dataset was originally published in Willett et al. [57]. The data was collected under appropriate ethical oversight, with approval from the Institutional Review Board at Stanford University (protocol #20804).
The monkey reaching dataset (MC_Maze is available through the DANDI Archive (https:// dandiarchive wow.ooonooooooooooo..o.oo daadddiee/naiioo /00t/000022 IDD: 000128) under CC-BY-4.0 license. This dataset contains sorted unit spiking times and behavioral data from primary motor and dorsal premotor cortex during delayed reaching task.
A3 Hyperparameters and compute resources
Table 2: Training details for autoencoder models on Lorenz, Monkey reach, and Human BCI datasets. We used the AdamW [II optimizer, whose learning rate 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, used K 5 for the temporal smoothness loss in Eq. Human BCI
Parameter
Lorenz
Monkey Reach
Datasef Details
Num training trials
3500
2008
8417 512 (max)
Trial length (bins)
256
140
Data channels (neurons)
128
182
128
Model Details
Hidden layer channels
256 8
256
256
Latent channels
16
32
Num AE blocks
Spike history
No
Used in unconditional
No
Training Details
Max learning
0.001
0.001
0.001
Adam W weight decay
0.01
0.01
0.01
Num epochs
200
140 (early stop.)
400
Num Warmup Epochs
10
10
20
Batch size
512
512
256
L2 reg. 31
0.01
0.001
0.001
Temporal smoothness 82
0.01
0.2
0.1
CD mask prob.
0.2
0.5
0.2
A3.1
Computational Resources
We performed all training and evaluation of LDNS the Lorenz and Monkey reach datasets 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 trained R 6 minutes, and the diffusion model Rs 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 within 30 minutes.
17

---

Table 3: Training details for diffusion models on Lorenz, Monkey reach, and Human BCI datasets.
We used the same learning rate scheduler as for the autoencoder.
Parameter
Lorenz
Monkey Reach
Human BCI
Model Details
Latent channels
8
16
32
Hidden layer channels
64
256
384
Num diffusion blocks
4
6
Num denoising steps
1000
1000
1000
Training Details
Max learning rate
0.001
0.001
0.001
AdamW weight decay
0.01
0.01
0.01
Num epochs
1000
2000
2000
Num warmup epochs
50
50
100
Batch size
512
512
256
For the Monkey reach dataset, the autoencoder with the given hyperparameters trained in minutes, and the unconditional and conditional diffusion models in 40 minutes to 1 hour. With similar sampling times as in Lorenz, the effective GPU wallclock time approximately within one hour. Optimizing the autoregressive observation model took less than 1 minute.
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-bal 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 Monkey reach. The autoencoder took 50 minutes 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, well for designing appropriate evaluations. We estimate the total effective GPU wallclock time to be ss 10x that of the final model runs. The AutoLFADS bascline was only trained once with PBT, as this framework automatically optimizes the model hyperparameters.
We implemented all training and evaluation code using the Pytorch framework and used Weights Biase to log metrics during training.
A4 Baseline comparison: Latent Factor Analysis via Dynamical Systems- LFADS
Latent Factor Analysis via Dynamical Systems (LFADS) a sequential variational autoencoder used to infer latent dynamical systems from neural population spiking activity [5]]36] LFADS consists of an encoder, a generator, and optionally, 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 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 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- -aaetd Based Training (PBT) has been proposed to automate hyperparameter selection, termed AutoLFADS (25
Paszze et. al. Py PyTorch: An Imperative Style, High-Performance Deep learning I Library (2019)
Lukas Beiwald. Experiment Tracking with Weights and Biases (2022)
18

---

In our experiments with the monkey reach dataset, we use the Py Torch implementation of AutoLFADS 477. 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 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 poosson-samplee 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).
A5
Supplementary baseline comparisons
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, (2I) and 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
VAE (pi-VAE, [63 ). It is important to note that while both
A5.1 Targeted Neural Dynamical Modeling (TNDM)
TNDM [1 is a VAE-based model designed to jointly model neural activity and behavior. TNDM extends LFADS by using RNN generate latent dynamics that are mapped to both neural activity and behavioral variables. TNDM separates the latent space into behavior-specifte and behavior- independent subspaces to disentangle task-relevant and intrinsic neural dynamies.
For our comparison on the unconditional monkey reach task, we trained TNDM using the architecture and hyperparameters proposed in the original implementation of the pape We used 64-dimensional latent dynamics for each of the two sets. These project to a total 10 latent factors (5 behavior- specific 2, and 5 eeehaior-independeen 2;), which is the maximum number demonstrated in the original work.
To generate unconditional samples, we sampled initial generator states from a standard normal prior (O, 1). 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. AT6E) 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 statisties.
A5.2 Poisson-identifiable VAE (pi-VAE)
Poisson-identifrable VAE (pi-VAD) 633 is VAE-based model for count data that ensures identifia- bility in the latent space. 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 hyperparameter We used a General Incompressible-foot Network as a decoder, with 2 behaviorally relevant dimensions and independent dimensions in the latent space. However, our evaluation context differs significantly from the original paper's 's demonstrations: while VAE was initially evaluated on 50ms time bins and straight reaches only, our comparison uses 5ms bins and conditions on angles across all reaches both middle and end trajectory points. The "label", or behavior, 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.
I 'Code adapted from Github respository:
*Code adapted from article by Lyndon Duong (2021)
bttps:/www.lyydooddoog.oo//////
19

---

Importantl, sampling from pi- VAE does not introduce any temporal dependence between spike bins within a trial - pi- ViE was not intended be a generative model of neural spiking data. The lack of temporal modeling in pi- pi-VAE's is fundamental limitation for generating realistic spike trains, as evident in our empirical comparisons (Supp. Fig. AT3 in yellow, Supp. Fig. AT3 Note that this failure cannot be diagnosed simply from looking at the sampled spiking data (Supp. Fig. A14).
A5.3 Contributions of LDNS in context to LFADS, TNDM and Pi-VAE
LDNS is designed specifically for the purpose of accurately generating neural spiking data (unconditionally or conditionally)- task often ignored by other LVMs designed for neural data analysis such 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 LFADS, see Sec.[A6.]]
One feature provided by some neural-behavioral analysis models (such as pi-VAE and TNDM) is explicit disentangling of neural vs. behavior-relevant latents. While found LDNS latents contain relevant behavioural information (Fig. 55- Supp. Fig. A17 we did not explicitly supervise the latent space induce this property.
Finally, the spike sestory-dependett 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 well LFADS (Table
Table 4: Added Baselines metries comparison. DrL. for the population spike count histogram and RMSE comparisons. Mean and standard deviation across 5 folds sampled with replacement. Bolded entries represent best-performing values for Poisson observation and spike-history observation model, respectively.
Method
DE psch
RMSE pairwise corr
RMSE mean isi
RMSE std isi
Vi-VAE
0.0063 + 2.9e-4
0.0031 + 1.08e-5
0.064 + 0.002
0.034 + 0.001
TNDM
0.0028 + 6. 6.6e--
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
FAtSLF
0.0036 le-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
20

---

A6 Supplementary Figures Lorenz
real
(ale
0.8
0.4
100
200
100 200
1806 (H8)
0.4
(8)
0.2
0.0
0.0
time
Figure A2: The autoeneoder captures the gt Lorenz synthetic firing rate perfectly a) Autoencoder predictions (pred) and true rates from the test set, together with their difference (error, in red). b) Reconstructions sampled from the Poisson observation model (right) closely resemble the test sample (left). Both spiking activity is binarized for the visualization.
Intered tstents
500
200
100
200
Mant
-2.5
500
200
0.0
10
10'
10"
10"
16
10
10"
10
10"
Figure A3: The S4 autoencoder infers smooth latents from discrete spikes and samples from the diffusion model capture the latent distribution. Inferred autoencoder latents for test sample. b) Power spectral density for all eight latent dimensions for the inferred autoencoder training set and samples from the diffusion model
21

---

a
gt correlation
autbencoder
Idns
gt rates vs. model
1.0
neuron id
100
100
0.5
ae
Idns
model
0.0
50
50
50
-0.5
0
-1.0
100
neuron id
gt corr
C
temporal auto-corr
temporal cross-corr
0.075
auto-cort
0.10
0.050
gt
X-COFT
gt
0.025
Idns 1
Idns 2
0.00
0.000
Idns
4
-20
20
-20
lag
lag
Figure A4: LDNS captures the correlation structure of the Lorenz dataset a) Both the autoencoder and LDNS-sampled rates capture the ground truth instantaneous correlation structure of the synthetic rates. b) The auto-correlation structure of ground truth and sampled spiking activity matches perfectly in neuron groups, sorted according correlation strength. Synthetic Lorenz data group x denoted by gt x, LDNS samples as Idns æ. e) The time-lagged cross-correlational structure is also perfectly captured by LDNS all groups.
22

---

A6.1 Length generalization of LFADS on Lorenz
To analyze whether LFADS [36 exhibits similar length generalization properties LDNS, we trained an Autol.FADS 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 16x the training length.
Qualitatively, we observed that while LFADS produced trajectories resembling the attractor dynamics of the ground truth Lorenz system (Supp. Fig. ASJ,, 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, observed that the population firing rates sometimes reached extreme values in some samples (Supp. Fig. A6). though they eventually returned to typical ranges. This behavior not observed in LDNS samples, suggesting that bidirectional generation in the diffusion model provides more stability in variable-length generation.
Ifads latents original length
210
Ex
2
b
Ifads latents exceeding original length
Figure A5: Length generalization of LEADS on Lorenz Different projections of a 40-dimensional latent space from LFADS trained the Lorenz system. Trajectories compared between a) the original length and b) 16 times the original length using sampled initial conditions. For comparison with LDNS length generalization, see Fig.
sample
sample 2
sample 3
sample 4
mean rate
1.0
meen rate
meen rate
moon rate
28
0.5
20
0.0
2000
4000
2000
4000
2000
4000
2000
4000
time steps [a.u.]
Sime staps [a.u.]
time steps [a.u.]
time steps [a.u.]
sample
sample
semple
sample
1.0
mean rate
mean rate
moan rate
mean rate
0.5
0.5
20
20
0.0
2000
4000
2000
4000
2000
4000
2000
4000
time steps [a.u.]
time steps [a.u]]
5me steps [a.u.]
time steps [a.u.]
Figure A6: Mean population firing rates for eight different samples, shown for both the original length (pink) and 16x the original length (purple).
23

---

A7
Supplementary Figures Human BCI
real data spkes
Idns sampled spikes
PI num
di uaan
unit kd
Figure A7: Visual comparison of different sampled spiking data from LDNS, with five samples from the true dataset.
data
Idns
10 +
PC1
10
PC2
0.0
10.0 0.0
10.0
time (s)
time (s)
Figure A8: First two principal components (PCs) of smoothed spikes from true data (left) and model samples (right). Each line represents one sampled trial. Spikes were smoothed using a Gaussian filter with window of 160ms prior extracting the PCs.
24

---

gt neuron corr matrix
ae neuron corr matrix
diff neuron corr matrix
neuron vs neuron correlation
0.5
0.5
Idns
0.0
diff
0.0
0.5
gt
Figure A9: Correlation matrices for real spiking data from human and LDNS, comparing the nutoeeccoder-nneeeee (ae) correlation (sampled from reconstructed rates) and correlation of sampled spikes (diff). The deviations from the data (gt) already arise at the autoencoder stage.
25

---

A8
Supplementary Figures Monkey Reach task
correlation
autbencoder
gl model
150
150
150
neuron
100
100
100
0.00
-0.03
lons
0
0
0
-0.00
3.03
Correlation matrices for real spiking data and samples from Poisson LDNS, concatenated across trials, comparing the autoencoder-imfereee (ae) correlation (sampled from reconstructed rates) and correlation of sampled spikes (diffusion + ac). Most deviations from the data (gt) already arise at the autoencoder stage.
Figure A10:
semplee
1000
100
100
100
0.0
0.5
0.0
0.5
time (s)
time (s)
Figure A1l: Visual comparison different sampled spiking data from LDNS with five samples from the real dataset.
Ifads Paisson
Ifeds with spike history
0.0005
autc-corr
0.0000
-0.0005
-25
25
-25
25
lag
lag
Figure A12: Equipping LFADS with spike history Auto-correlation of data, LFADS samples with Poisson observations (left) and LFADSsh samples (with spike history), grouped according correlation strength.
26

---

hist
RMSE
RMSE
Vilee Soveewwern teban
0.000
0.00
I
3.0
2.04
0.04
3.00
0.001
300
3.00
ouR
our
Figure A13: Performance comparison with additional baselines pi-VAE 63 TNDM (2I), LFADS 36.. LFADS with spike history (LFADSsh), LDNS and LDNS with spike history (LDNSshh. Mean and standard deviation across 5 folds sampled with replacement.
pivae
Sime
Figure A14: Visual comparison of sampled spiking data from LDNS and all baselines with real data.
a
Idns
Idnssh
Ifads
Madssh
Indm
PC1
-5
0.0
0.5
0.0
0.5
0.0
0.5
0.0
0.0
0.5
0.0
0.5
PC2
-5
0.0
0.5
0.0
0.5
0.0
0.5
0.0
0.5
0.5
0.0
0.5
6
time [s]
PC power (dB)
PC power (8g)
PC power (B8]
PC power (dB)
PC power (dB)
PC power (BB)
10
data
10"
10
10
10
10"
10"
10"
10"
10"
10
10
frequency [Hz]
frequency
frequency
Trequency [Hz]
frequency [Hz]
frequency
MSE median ped
10"
PC1
10
PC2
MSE median
10"
Idna
Idnash Mads Maadsh trdm
pivae
Idns
Idnesh Ifads Ifadssh pivae
Figure A15: Comparing principal components smoothed sampled spikes a) First two principal components (PCs) of smoothed spikes from model samples. Each line represents one sampled trial. Spikes were smoothed using Gaussian filter with a window of 40ms prior to extracting the PCs. The PCs were fit using real data. Since pi- VAE does not account for temporal dynamics, it does not show any temporal structure in the PCs. b) Power spectral density (PSD) of PC1, plotted for model vs. data (in grey). Mean squared error of median PSD between model samples and data for PC1 (left) and PC2 (right). LDNS and LDNSsh perform the best here, with pi- VAE showing large errors.
27

---

dersity
dersity
Idhash
Knns
0.00
0.00
0.1
0.1
0.0
-0.05
0.0
-0.05
10
-0.05
0.00
10
-0.05
0.00
spiik count
gt
spiee count
gt
Idnosh sid tsi
lons moan
0.2
0.2
0.2
0.2
Kina std
0.1
Kdhosh
0.1
0.1
0.2
0.1
0.2
0.2
mean
gt
gt
gt
0.2
nota
danaty
Ibds
hoads
denatyy
cott kide Rechsh hode
Fidtchh
faassh
0.00
0.1
Meda Aole
0.1
-0.05
10
-3.06
0.00
10
-0.05
0.00
spice
0.2
Ifadzsh aod bai
Hadd moan
0.1
llads
0.1
0.1
0.1
0.2
mean isi
gt std isi
gt mean isi
gt std isi
02
0.2
dansity
Irdm
0.00
denety
pivae
0.00
0.1
0.0
-0.05
0.0
-0.05
10
-0.05
0.00
spike count
gt
spise count
0.8
0.3
piaae std m1
0.2
Erdin mean
0.2
0.1
0.1
0.2
0.1
0.25
0.50
0.0
0.2
0.1 0.2 0.3
0.0
0.2
g1 mean
gt std
gt meaa
gt std
Figure A16: Population-level and single neuron-level statistics of a) LDNS, b) LDNSsh (with spike history), e) LFADS, d) LFADSsh (with spike history), e) TNDM, and f) pi-VAE.
28

---

iraa Welent
kins Iatart
kina PC1
Idra PC2
kina PC1 PCC
camplet ottent
of Ittetts
stra gitt test conditions
0.0
0.5
0.0
0.5
0.0
0.5
0.0
0.5
time [4]
time
time
sim [+]
PC1
wutorrodder iatent
tatents
-25
0
25
0
tmm
time [s]
time
tmm
PC1
Figure A17: Comparison LDNS latent space trajectories of inferred and conditionally sampled latents a) Straight reaches from the Monkey reach test set. b) LDNS sampled latents (velocity- conditioned on reaches shown in a). e) PCs of LDNS sampled latents. d) Autoencoder-inferree latents of corresponding neural activity for the reaches shown in a). PCs of autoencoder--nfeeree latents.
29

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our work is methodological and empirical, and justify our claims by evaluating our method's '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 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 long is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
We discuss the limitations of our work in a paragraph in the last section of the main paper.
Justification:
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create 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 (c.g.. independence assumptions, noiseless settings, model sell-speciiiiatinn 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
30

---

Answer: [NA]
Justification:
We presented no theoretical results here, as our method is empirically motivated and evaluated.
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 cross- referenced.
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.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: In the Appendix, we provide the full architecture of the proposed LDNS model, along with hyperparameter configurations for the model, data, and training algo- rithm. We also make the training loss explicit for the autoencoder and diffusion model. Finally, we provide Github repository httt:://gthhub.com/maaklaaaaall implementation.
with our code
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, 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.
the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution be necessary to either make it possible for others to replicate the model with the same dataset, provide access 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 replicate the results, access 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.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the partieular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
31

---

some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
We have open sourced the code and instructions for reproducing our exper- iments on our Github repository https//ggiuub..oo/maacalaal access, and we provide links to downloading the data.
Justification:
The data is open
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed reproduce the results. See the NeurIPS code and data submission guidelines (https
//wips.cc/publlle wttpcc/pppppgppu.ieeeeeeep.op guides/deeeeee CodeSubmi ssioiioicyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide 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 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, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: For all results that require statistical significance (Table in main paper, Table appendix), we report performance and standard deviation over 5 sets of samples.
Guidelines:
The answer NA means that the paper does not include experiments.
32

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eπror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We list the hardware configurations used for training and evaluation, 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 compute required for each of the individual experimental rums well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: We have read and confirmed that we conform to all the points mentioned in the NeurlPS 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
33

---

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.
Justification:
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 partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point out that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: In our opinion, we do not foresee high risk of misuse of our work.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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, CCO 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 bascline implementation LFADSTorch.)
34

---

Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperssi thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: We provide zipped directory of the code use for training and evaluating our models and baselines. As previously mentioned questions 4,5 of checklist, we provide instructions to understand and run the code for reproducibility.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/codeemoodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: We did not perform any crowdsourcing experiments or research with human subjects in this work. Additionally, during the review of this paper, we had two ethics reviewers assigned to us regarding the use human BCI data in one experiment. Both ethics reviewers found issues with our paper, this data CCO 1.0 licenced, and therefore open access for us to use.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
35

---

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or 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 NeurlPS 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.
36


Reviews:
[{"review_id": "BMRdJIxRBM", "reviewer": "Reviewer_AuW4", "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.", "questions": "Are the results reported in figures 2 and 3 using the spike history filter or just Poisson likelihoods?\n\nWhat is the 'kde' in figure 3d?", "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.", "rating": 7, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "strengths": "I see three separate contributions of this paper --  1) The ability to handle time-varying inputs in VAE based neural models is rare, as the authors point out, yet important in many practical neuroscience applications.  2) Using diffusion-based methods to generate samples from the auto-encoder inferred latents, allowing for the generation of variable length, realistic spiking data on unseen behaviors. I think the results of figure 5 are particularly interesting in this regard. 3) the addition of a history-based poisson likelihood, which better captures individual spiking statistics. This is a somewhat superficial contribution, but is nonetheless important to many of the authors' reported improvements.", "weaknesses": "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. \n\nFor 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. \n\nOverall, 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. \n\nLastly, 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. \n\n*) \"Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-VAE\"\n**) \"Targeted Neural Dynamical Modeling\"\n***) \"Modeling behaviorally relevant neural dynamics enabled by preferential subspace identification\"", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for finding our work relevant, and our approach sound. Based on your suggestions, we have now performed new baseline experiments with additional VAE-based models. We also performed latent space analyses of sampled spikes from all models (using PCA and LFADS embeddings) to supplement the spike statistics. We additionally inspected the quality and interpretability of LDNS-derived latents as suggested. Lastly, we augmented LFADS with spike history dependence for a more clear comparison, and clarify how and why we account for it post-hoc.\n\n**In summary, LDNS remains the most performant model overall for spike generation, and its single-trial latents are informative of behavior**. We detail the experiments and analyses below, and clarify the disparate contributions and capabilities of LDNS more clearly, especially in contrast to existing approaches. We hope these additions address the reviewer’s concerns and enable them to raise their score. \n\nNew figures (**Fig. Rx**) are in the **PDF**.\n\n**New VAE-based baseline experiments**  \nWe implemented two new VAE-based models suggested by the reviewer, TNDM (Hurwitz et. al. 2021\\) and pi-VAE (Zhou et al, 2020), on the unconditional monkey reach task, and find that LDNS is superior in realistic spike generation (**Fig. R1, see PDF**).\n\n**TNDM**: We trained TNDM on the monkey dataset using the original proposed architecture and hyperparameters. We used 5+5 latent factors, the maximum in the original paper, and used the prior N(0,1) to sample the RNN initial states for unconditional generation.  \n**pi-VAE**: We trained pi-VAE using the original architecture (and model hyperparameters), which does not consider temporal dynamics and treats each time point as an independent sample. Thus, pi-VAE samples exhibit poor spiking statistics and no temporal structure (**Fig. R1, R7,** yellow).\n\n**Evaluation of LDNS latents**  \nWe agree that, while generation of realistic spike trains in terms of spike statistics is important, many neuroscientists are interested in the extracted latents. Thus, we analyzed LDNS-extracted latents of unseen test samples from the monkey reach dataset and colored them by reach angle (**Fig. R6,** straight reaches only for visualization). Note that behaviorally relevant information is clearly reflected in single-trial latents of LDNS (**Fig. R6**, bottom row). Furthermore, this latent structure is preserved when sampling from LDNS conditioned on unseen velocity trajectories (**Fig. R6**, top row). We will add this result to Fig. 5, highlighting the ability of LDNS to extract meaningful latents.\n\n**Contribution of spike history and extension to LFADS**  \nSpike trains exhibit autoregressive temporal dependencies not shared across the population, and therefore spike-history coupling is needed for any low-dimensional model to achieve realistic spike train generation.   \nThe spike history model in LDNS is fit post-hoc after autoencoder training, and therefore does not impact the inferred latents. This allows us to independently improve LDNS-generated spike trains and their autocorrelations, an important component towards accurate modeling of spiking data. Directly optimizing the spike history features during autoencoder training might interfere with the coordinated dropout regularization (Keshtkaran et al. 2019, see methods), which has been shown to be critical for other LVMs of neural dynamics.  \nLastly, we view the post-hoc augmentation of the observation model with spike history as a key modular contribution, which can be flexibly applied to other generative models. Based on the reviewer’s suggestion, we have additionally extended LFADS with spike history dependence, which improves its performance (**Fig. R1,4**). However, LDNS with spike history is still superior or on par on spike metrics against this extended version of LFADS. We also emphasize that, both with and without spike history, LDNS outperforms its LFADS counterpart, and we highlight the corresponding comparisons in Table 1 (see Reviewer y7Q4). We will include these results and clarifications in the revised paper.\n\n**Contributions of LDNS (in contrast to other approaches):** \n\n- LDNS combines S4 and diffusion models for the purpose of accurately modeling and generating neural spiking data—a task often ignored by other LVMs designed for neural data analysis (such as LFADS, pi-VAE, and TNDM).   \n- The S4 autoencoder and diffusion model are trained in separate stages, offering modularity and easier debugging, while both components naturally account for temporal dependencies (unlike pi-VAE).  \n- Furthermore, S4 is autoregressive, similar to other RNN-based models, but empirically we found it to perform better when extending past the training trial length and more readily handle variable length data (compared to LFADS), due to the masked training procedure we propose here.  \n- Finally, the spike history-dependent observation model is modular and can be optimized post-hoc using rate predictions of any model, offering flexibility while improving spike generation quality.  \n- One feature provided by other neural-behavioral analysis models (such as pi-VAE and TNDM) is an explicit disentangling of neural vs. behavior-relevant latents, which we did not consider but is a possible future extension for LNDS.\n\n**Are the results reported in figures 2 and 3 using the spike history filter or just Poisson likelihoods?** For the results in Figs. 2, 3, 4(b,c,d), 5, we only use the Poisson likelihood with no history terms. In the  Lorenz experiment, we know the ground-truth generative model uses Poisson emission, and for the human BCI dataset we omitted them due to the large bin size (20ms) used here. \n\n**What is the 'kde' in figure 3d?**  A Kernel Density Estimate (KDE) with a Gaussian kernel to estimate the population spike count distribution. We will clarify both points in the revision.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "mrQVlzKBVv", "reviewer": "Reviewer_y7Q4", "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.", "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? \n\n- 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?\n\n- 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?\n\n- 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": "Yes.", "rating": 7, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The proposed generative model of neural activity appears powerful and generally useful of fitting neural responses across a large variety of conditions. Both the unconditional and conditional generative performance of the model is impressive. Many components of the model and training process are well-motivated and clearly described. For these reasons, I find this to be a significant contribution.", "weaknesses": "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for finding the contributions presented in our work significant, and for noting the flexibility and performance of our model in generating unconditional and conditional neural spiking data in a wide variety of conditions. \n\nThe reviewer had questions and concerns over training and architectural details of the model, and the role of the spike history dependence model in the superior performance of LDNSsh. To disentangle the contribution of this observation model from the architectural contributions of our work, as suggested by the reviewer, we now trained an extended LFADS model by equipping it with the same spike history observation and found that it is still matched or outperformed by the equivalent LDNSsh model. We also discuss our reasoning for fitting these terms post-hoc, and why we used a softplus approximation. Finally, we address the reviewer’s questions on architectural choices, and on the attempted speech dataset.\n\nNew figures (**Fig. Rx**) are in the **PDF**.\n\n**Equipping LFADS with spike history**  \nWe agree that a better characterization of the contribution of spike history, relative to S4 and diffusion, would be beneficial. We have now extended LFADS with spike history dependence (LFADSsh) using the approach we introduced for LDNS, and found it to improve its performance (**Fig. R1,4**). However, LDNS with spike history (LDNSsh) is still superior or on par on spike generation metrics.\n\nFurthermore, even without spike history, LDNS was already on par or better compared to its counterpart in LFADS in our original evaluation, and we will emphasize the corresponding comparisons more prominently in Table 1 (LDNS vs. LFADS, with and without sh).  Thus, while spike history couplings are needed for realistic spike train generation (as neural data contains dynamics that are not shared across the population, and thus cannot be captured by a low-d latent state), the performance benefits of LDNS are not due to spike history couplings alone. \n\nYour suggestion also allowed us to show that our post-hoc fitting of spike-history couplings provides a way to increase the realism of generated spike data not just for LDNS but for a class of VAE-based methods.\n\n**Table 1**\n\n|Method|$\\mathbf{D_{KL},\\text{psch}}$|RMSE pairwise corr|RMSE mean isi|RMSE std isi|\n|--|--|--|--|--|\n|AutoLFADS|$0.0040\\pm2.2\\times10^{-4}$|$0.0026\\pm1.25\\times10^{-5}$|$0.039\\pm0.003$|$0.029\\pm0.001$|\n|LDNS|$0.0039\\pm3.9\\times 10^{-4}$|$\\mathbf{0.0025\\pm1.1\\times10^{-5}}$|$0.037\\pm0.001$|$\\mathbf{0.023\\pm0.001}$|\n| | | | | |\n|AutoLFADSsh|$0.0036\\pm2.1\\times10^{-4} $|$0.0026\\pm1.8\\times10^{-5}$|$0.034\\pm0.002$|$\\mathbf{0.023\\pm0.0001}$|\n|LDNSsh|$\\mathbf{0.0016\\pm6.2\\times10^{-4}}$|$\\mathbf{0.0025\\pm1.07\\times10^{-5}}$|$\\mathbf{0.024\\pm0.002}$|$\\mathbf{0.023\\pm0.001}$|\n\n**Post-hoc training for spike history, and choice of softplus vs. exponential**  \nTaking rate predictions from LDNS (or any model, see previous paragraph), we optimize the spike history parameters with respect to the ground-truth spiking data. As a result, this alternative observation model does not impact the latents inferred by the S4-autoencoder. It allows us to independently improve generated spike trains, in particular to capture realistic autocorrelations, an important component towards accurate modeling of spiking data that is missing in the deep LVM literature as a whole. Furthermore, this opens the possibility for replacing our current version of spike history dependence with more sophisticated observation models without re-incorporating it into autoencoder or diffusion model training.\n\nWe fit spike history dependence post-hoc since jointly optimizing with the autoencoder would likely interfere with the coordinated dropout regularization (Keshtkaran et al. 2019, see methods), which has been shown to be critical for other LVMs for neural dynamics. As we show in the above experiment, this post-hoc fitting process can be straightforwardly adapted to improve other models as well. \n\nEmpirically, we found that using the exponential function was less numerically stable and resulted in higher loss values and poorer data fit than using the softplus-approximation. Softplus resulted in faster training convergence and a lower final loss (**Fig. R5**).   Furthermore, we believe that in the low spike count regime this is a fair approximation, as 99.8% of the inferred rates (across time bins) were less than 0.2 (per bin).\n\n**Other questions**  \n**S4 vs. Mamba as an architectural choice**: S4 is parameterized as a time-invariant (stationary/autonomous dynamics) system, while Mamba and other context-selective models are parameterized as (and may be better suited for) non-stationary, input-driven dynamics, such as the human speech data. Both S4 and Mamba allow for time-parallelized training and length generalization. Exploring Mamba and other non-stationary linear recurrence models is an interesting idea of future work, but would go beyond the scope of this project.\n\n**Overestimation of correlation in attempted speech data**: Based on the reviewer’s suggestion, we increased the dimensionality of the latent space in LDNS from 32 to 48 and observed that the overestimation remains. We also analyzed the autoencoder and diffusion model separately and find that the inflated correlations are a result of the AE, not the diffusion part (**Fig. R9**). We do not yet fully understand the exact source of this overestimation, which we only observed for this particular dataset, and agree that further characterization would be beneficial.\n\n**Conditional generation of speech data**: While this application would be a natural extension for conditional generative modeling on complex behavioral data, in this work we focus on developing the methodology to enable such exciting applications in the future. Decoding speech from spiking data is an active research area, and consequently, evaluating the correctness of sampled data and decoded speech would be beyond the scope.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Pn4hicNV5k", "reviewer": "Reviewer_pMhd", "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.", "questions": "1. What is the significance of studying the generation of spiking data? Is this generated data of practical value?\n2. 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.\n3. Can you explain how the stochastic operations in Figure 1 are performed?\n4. 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": "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:\n1. 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.\n2. 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.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. Quality: The paper is comprehensive, with detailed explanations of the model's characteristics and corresponding experimental designs.\n2. Clarity: The paper is logically clear and well-structured.\n3. Importance: The paper significantly contributes to addressing the modeling challenges of complex datasets in neuroscience.", "weaknesses": "1. There is only one comparison method in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for finding our work comprehensive and our writing and presentation clear and well-structured. We also appreciate that the reviewer finds our contributions of addressing the modeling challenges of complex neuroscientific datasets significant.\n\n\n**Additional comparison with other VAE-based methods**  \nFollowing reviewer suggestions (as well as that of AuW4), we have now included two additional proposed VAE-based models as baselines:  Poisson-identifiable VAE (pi-VAE, Zhou et. al, 2020\\) and Targeted Neural Dynamical Modeling (TNDM, Hurwitz et. al, 2021\\) on the unconditional monkey reach task. We emphasize that while these LVMs have successfully been used for analyzing neural and behavioral data, they were not intended for realistic spike generation. \n\nNew figures (**Fig. Rx**) are in the **PDF**.\n\nCompared to these additional baselines, as well as an extended LFADS model augmented with spike history, we found that LDNS remains superior in realistic spike generation (**Fig. R1**), and recovers structured latents informative of behavior (**Fig. R6**)**.** Implementation details of these baselines are provided below, and will be included in the revised paper.\n\nWe also now trained LFADS on the Lorenz and human BCI tasks, in particular to evaluate its ability to length-generalize and handle variable length input data compared to LDNS. We find that LFADS had difficulties in extending the learned dynamics correctly (**e.g., Fig. R2,3**) – in contrast to LDNS, which can accurately generalize to 16 times the training length in the Lorenz task and can naturally deal with variable length trials in the human task.\n\n**TNDM**: We trained TNDM on the monkey dataset using the original proposed architecture and model hyperparameters. We used 5+5 latent factors, the maximum shown in the original paper. We used the prior N(0,1) to sample the initial generator states for unconditional sampling of spikes (**Fig. R1,** blue). \n\n**pi-VAE**: We trained pi-VAE on the monkey dataset using the original proposed architecture and model hyperparameters. pi-VAE’s architecture does not consider temporal dynamics, and treats each time point as an independent sample. Furthermore, while Zhou et al. 2020 evaluated pi-VAE on 50ms time bins and straight reaches only, we here use 5ms bins and condition on angles of all reaches in the middle and end of the trajectory. Sampled spiking data shows poor statistics mainly due to the lack of temporal dependence in the model (**Fig. R1, R7,** yellow).\n\n**Significance and practical value of realistic spiking data generation**  \nAccurate modeling and generation of neural spiking data has scientific, clinical, and practical value. If generated data were used for augmenting training for a downstream task, for example, introducing obvious artifacts (such as not capturing refractory periods) can introduce bias. Furthermore, when studying relations between variables in neuroscience with such emulator models, subtle changes such as spike times or phases when spikes occur in oscillations are known to make a difference. \n\n**Modeling activities from other brain regions**  \nOur work proposes a general methodology for modeling and generating spiking data, and is agnostic to the particular brain region where it may be recorded from. In our evaluation, we have considered the motor cortex of monkeys and the speech cortex of humans, but it can be straightforwardly applied to other datasets as well. We will discuss these possibilities in the revised paper.\n\n**Clarification on ethics concern and suggestions on data privacy**  \nThe human BCI dataset we used is publicly available under a CC0 1.0 Universal Public Domain Dedication license. It is from a peer-reviewed paper that was previously published in Nature (Willett et al. (2023). A high-performance speech neuroprosthesis. *Nature*.), and according to this paper, was cleared in ethical reviews by the Institutional Review Board at Stanford University (protocol \\#20804).  Our paper did not provide new sensitive data, nor does it provide a methodology to obtain such data. We will make this information clear in the appendix of the revised paper. \n\nWe agree that protecting data privacy is very important, especially when using sensitive data involving human participants, and will further acknowledge this in the Discussion section. \n\n**Other clarifications**  \n**Stochastic operation in Figure 1**: Unless specified otherwise, we use the Poisson observation model as the stochastic operation, going from inferred Poisson rates to spike counts. In specific scenarios (e.g., Fig. 4e), we extend this observation model to include spike history dependence in the LDNSsh model. This dependence allows us to capture e.g. refractory periods of neuron firing behavior by reducing the probability of a spike occurring directly after a previous spike, enabling LDNSsh to accurately capture biologically plausible spiking statistics. \n\n**3-axis reach conditioning**: Our architecture is agnostic to the number of reach axes. The dataset we consider involves monkeys performing 2-dimensional reach tasks, which we pass to the diffusion model as two additional channels or as a scalar reach angle (see Fig. 5 and A1). This can be extended straightforwardly to condition on higher-dimensional behavioral variables, allowing to also model reaches in 3-dimensions.\n\nWe hope these additions, in particular the inclusion of additional comparison methods, and clarifications, address the reviewer’s concerns and enable them to raise their score.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EmKb37PX7W", "reviewer": "Reviewer_mYEj", "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. \n\nThe 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. \n\nAccording 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.", "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. \n\nCan 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": "I think the authors have adequately addressed limitations and possible negative societal impacts of their work.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The submitted work is original, in that latent diffusion models (LDMs) to my knowledge have not yet been used to (conditionally) generate neural population activity. I agree that LDMs are indeed a nice idea here, since the separation into autoencoder stage and diffusion model stage means we can try to use denoising diffusion models without much changes.\n\nThe work seems technically solid to me. The authors make sensible changes to the LDM for the task of generating neural population activity. The adaptations for discreteness (Poisson loss, spike-history terms) show a good understanding of best practices in neural data analysis. The inclusion of S4 layers into autoencoder and diffusion models to handle and generate variable-length sequences is interesting, and could further spread usage of these tools in the computational neuroscience community.\n\nThe manuscript is well written, explaining not only the key ideas and experiments, but also summarizing the important model parts (autoencoders, denoising diffusion models) next to the non-standard adaptations (Poisson loss, S4 layers, spike-history terms) done by the authors for this study. The figures are high quality and way above what I would expect from a conference paper.\n\nThis study follows and expands a (small) recent trend of using modern machine learning algorithms for generation of neural population activity. LDNS models / LDMs clearly have some advantages over previously suggested models, in that they combine the low-dimensional representation of autoencoders with the generative fidelity of diffusion models.", "weaknesses": "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. \n\nFor 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. \n\nAnother 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed summary and evaluation and suggestions that will significantly improve our work, e.g. recommending to further assess the dynamics of the generated samples. We are also thankful for the generally positive response and describing our work as original.\n\n**Analyzing population dynamics to assess sample quality beyond spike statistics**  \nWe agree with the reviewer’s point that an evaluation of LDNS-generated samples in terms of dynamics would be valuable. We addressed this in two different ways:\n\n**1\\. Principal Component Analysis (PCA) on smoothed spikes:**   \nWe followed a common approach in neuroscience of performing dimensionality reduction using **PCA on smoothed spikes and then visualizing the resulting first n principal components** over time (**Fig. R7; please see attached PDF**, here, n=2). Overall, we see that PCs of unconditional LDNS samples closely match that of the real held-out data samples. This approach allowed us to clearly see inconsistencies in the dynamics, e.g. for an added baseline model pi-VAE (see Reviewer AuW4) that does not model dynamics, leading to flat PCs (**Fig. R7, left, yellow**) and higher MSE between the median power spectral density of pi-VAE samples to true data as compared to all other methods (**Fig. R7, right**).\n\n**2\\.  Comparing inferred latent dynamics of sampled spikes using LFADS:**  \nAs a more sophisticated dynamics analysis, we followed the reviewer’s suggestion and embedded both the true data and the unconditionally generated sampled spiking data from LDNS using **LFADS** (used here as an **analysis tool**). This allows us to compare latent dynamics by comparing the distributions of inferred initial states of the LFADS generator RNN (**Fig. R8**, visualized as PCs of $g\\_0$) across true data and generated spikes. This analysis revealed that LDNS spikes more closely capture the broad true data distributions than, e.g., spikes sampled from LFADS (as a generative model) (**Fig. R8**). Such analyses, however, require a well-fit state-space model, which in of itself poses a challenge for many datasets.\n\nWe agree on the **value of visualizing more spike raster plots and extracted principal components**, and will add them in the Appendix for all models and tasks. Similar to Appendix A6, we will also add and quantify power spectral densities of the principal components (**Fig. R7**) for the different methods against the real data.\n\n**Fitting the LFADS baseline to the human data and the Lorenz dataset**  \nFollowing the reviewer’s suggestion, we now also trained LFADS on the human dataset and the Lorenz dataset. \n\n**Human:** The human dataset is challenging as the trials are highly heterogeneous with different lengths. LDNS can handle these challenges due to our architectural choices and our masking training scheme even if, as pointed out, LDNS does not work as perfectly (see discussion below) as in the other datasets. Please note that this is a new dataset (released less than a year ago), and to date, no readily available generative modeling baseline exists for this data. To allow for a fair comparison, we now attempted to modify LFADS so that it can be fit to this data. \n\nTo be able to fit LFADS on this dataset, we had to cut the data into equal length segments of 140 time steps (2.8s), since without major modifications LFADS (with its bidirectional encoder architecture) is not well-suited to handling variable length inputs. Despite capturing the spiking statistics on the 2.8s length it was trained on \\-- LFADS failed to capture any realistic dynamics beyond the 2.8s cut-off during sampling (**Fig. R3**). The latents of LFADS decay to a fixed point when the RNN is run forward beyond 2.8s, indicating that it is ill-suited to variable-length generation of such heterogeneous data (**Fig. R3**), highlighting the need for flexible methods such as LDNS.  \nWe acknowledge that there may be variants of LFADS that can handle this case better, however, our point is that “vanilla” LFADS is unable to handle this task without significant modifications. \n\n**Lorenz:** While it is well established that LFADS fits the Lorenz dataset well (Sussillo et al. 2015, Pandarinath et al. 2018), LFADS struggles with length generalization: Running the LFADS generator on 16 times the original length results in inconsistent latent trajectories (**Fig. R2**) compared to LDNS (Fig. 2c).\n\nWe will additionally point out in the main text, not just in the Computational Resources Appendix, that LFADS is significantly more computationally expensive to run than LDNS, and we thank the reviewer for this suggestion.\n\n**Comparison against additional baselines on the monkey reach task**  \nIn addition, as also suggested by other reviewers, we include two additional baselines for comparison on the monkey dataset: TNDM (Hurwitz et. al, 2021\\) and pi-VAE (Zhou et al. 2020), are VAE-based models that were proposed as analysis tools to jointly study neural and behavioral data. We find that LDNS and its spike-history based extension outperforms these baselines on the monkey dataset (see **Fig. R1** and response to reviewers AuW4 and pMhd for implementation details).\n\n**Contribution of autoencoder vs. diffusion for correlations**  \nWe agree that disentangling the contribution of the S4 autoencoder and diffusion is useful, particularly in cases when LDNS does not work perfectly. The model errors arise primarily from the autoencoder, not from the diffusion model: The values of pairwise correlations both in monkey reach data and human are very similar for the autoencoder reconstructions and LDNS samples (**Fig. R9**)**,** i.e. mismatches arose at the autoencoder and not the diffusion stage. \n\nAnalogous to the Lorenz experiment (Fig. A7), we will add the analyses comparing the autoencoder performance and the diffusion performance in the Supplementary Material for the two other datasets (monkey, human).", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their constructive and detailed engagement with our work, resulting in many helpful comments and opportunities for clarification. We are especially grateful for several reviewers’ assessments that the work is “original” (mYEj), well written and clearly motivated (y7Q4, mYEj, pMhd), technically sound (mYEj, AuW4) with “impressive performance” (y7Q4), and represents a significant contribution to the field (pMhd, y7Q4).

We agree with many of the suggestions and questions raised, and respond individually to each review in detail. We summarize here the main points and new experiments. We hope that our response addresses their questions and concerns, allowing all reviewers to recommend our work for acceptance.

**Summary of original contribution**  
LDNS combines S4-based autoencoders and diffusion models for both extracting low-dimensional latent dynamics from neural population spiking data and generating realistic spike trains, an aspect of evaluation often ignored by other works in the literature. On a commonly used monkey reach task, LDNS outperforms or matches LFADS in all metrics. By equipping LDNS with a modular spike history-dependent observation model, it surpasses LFADS on all metrics and accurately captures spike autocorrelation. Moreover, LDNS works well with variable trial lengths and shows length generalization abilities. Finally, LDNS can conditionally generate realistic spiking data based on behavioral covariates for the monkey reach task. 

**Summary of new results**  
In response to the reviewers, we substantially expanded our evaluation with new experiments: On the monkey reach task, LDNS outperforms two newly suggested baseline models **(Fig. R1, see PDF**). We also augment LFADS with spike history and report increased performance, but it is still surpassed or matched by LDNS-sh on the same task (**Fig. R4, Table 1 in response to pMhd**). In addition, we evaluate LFADS on the other two datasets and find that unlike LDNS, LFADS failed at length generalization (**Fig. R2,3**). Moreover, we show that on the monkey task, LDNS latents contain behaviorally relevant information, suggesting their utility in analyzing and visualizing neural data (**Fig. R6**). Finally, we evaluate the dynamics of generated spikes sampled from LDNS (and other models) via PCA and LFADS embedding, and find that LDNS accurately captures the dynamics of real data (**Fig. R7,8**).

**1\. New baseline experiments (AuW4, pMhd, mYEj)**  
We conducted numerous additional baseline experiments to demonstrate the contributions of LDNS relative to existing (VAE-based) methods:  
We implemented and fit two other methods, pi-VAE (Zhou and Wei, 2020\) and TNDM (Hurwitz et al. 2021), to the monkey dataset (**Fig. R1**), and find that LDNS still consistently outperforms the others on spike generation. Details in individual responses to AuW4 and PMhd.   
We added our proposed spike history observation model (see point 3 below) to the LFADS baseline on the same task, which increases performance but still not to the level of LDNS with spike history. We note that LDNS without spike history is still superior or on par compared to all baselines, highlighting the contribution from the latent diffusion model.

**2\. Evaluating LFADS on the other datasets (mYEj)**  
We further applied LFADS on the Lorenz and human dataset, in particular, to assess the length-generalization ability of LDNS relative to existing models of neural dynamics.  
It is well established that LFADS fits the Lorenz dataset well (Pandarinath et al. 2018). However, LFADS struggles with length generalization (**Fig. R2**), while LDNS samples can be generated at 16x the training trial length (Fig. 2c), a key feature made possible by S4.  
We also fit LFADS to spike recordings from the human brain, a challenging task due to highly heterogeneous trials with variable length, which LDNS can handle thanks to its architecture and masking scheme. To fit LFADS, we cut the data into equal length segments. Despite successful training, LFADS failed to capture dynamics beyond the cut-off during sampling (**Fig. R3** vs. Fig. 3). This highlights the capability of LDNS to deal with variable-length data, which is highly relevant for many neuroscience datasets (as pointed out by AuW4). Further discussion in response to reviewer mYEj.

**3\. Contribution of, and clarifications on, spike-history coupling (AuW4, y7Q4)**  
We clarify that the spike-history observation model is applied **after** first training the autoencoder model with standard Poisson likelihood. The history model increases the quality of the sampled spikes, and can be flexibly applied to other models. Addressing questions raised by reviewers, we ran a new experiment augmenting LFADS with spike history, which increases the realism of generated samples from this baseline model (**Fig. R4**), though it still underperforms LDNS with spike history.

**4\. Analysis of LDNS latents (AuW4, mYEj)**  
We agree that single-trial latent analysis is of interest to the neuroscience community, and have therefore evaluated the quality of LDNS latents (**Fig. R6**). On the monkey reach dataset, we plot PCs of the inferred latents and show that behaviorally relevant information (colored by reach angle) is reflected in the latent space of LDNS, a desirable property for neuroscientists aiming to visualize their data.

**5\. Evaluating dynamics of LDNS samples (mYEj)**  
We evaluated the realism of LDNS samples in terms of underlying low dimensional dynamics using PCA (applied to smoothed, generated spikes vs. true data, **Fig. R7**). This suggestion clearly distinguished failure cases (e.g., no temporal dependencies in samples generated using pi-VAE). As suggested by reviewer mYEj, we also embedded the generated spikes using LFADS, and show that LFADS latents extracted from LDNS spikes more closely capture the broad distributions of the true data than, e.g., spikes sampled from LFADS (**Fig. R8**).

Meta-review: 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.
Decision: Accept (spotlight)