Title: Latent Diffusion for Neural Spiking Data
Venue: NeurIPS 2024
Accepted: yes
Keywords: neural population

diffusion models

latent variable models

electrophysiology

brain-computer interfaces
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

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.

1 Introduction

Modern datasets in neuroscience are becoming increasingly high-dimensional with fast-paced innovations in measurement technology [1, 48, 23], granting access to hundreds to thousands of simultaneously recorded neurons. At the same time, the types of animal behaviors and sensory stimuli under investigation have become more naturalistic and complex, resulting in experimental setups with heterogeneous trials of varying length, or lacking trial structure altogether [35, 33, 57]. Therefore, a key target in systems neuroscience has shifted towards understanding the relationship between high-dimensional neural activity and complex behaviors.

For high-dimensional neural recordings, analyses that infer low-dimensional structures have been very useful for making sense of such data [11]. For example, latent variable models (LVMs) are often used to identify neural population dynamics not apparent at the level of single neurons [61, 32, 40]. More recently, deep learning-based approaches based on variational autoencoders (VAEs) [26, 44,

51, 36, 63, 22, 6, 16] have become particularly popular for inferring latent neural representations due to their expressiveness and ability to scale to large, heterogeneous neural recordings with behavioral covariates [36, 63, 16].

However, in addition to learning latent representations, another important consideration is the ability to act as faithful generative models of the data. In other words, models should be able to produce diverse, realistic samples of the neural activity they were trained on, ideally in a behavior- or stimulus-dependent manner. Models with such capabilities not only afford better interpretability analyses and diagnoses for whether structures underlying the data are accurately learned, but have a variety of downstream applications surrounding the design of closed-loop in silico experiments. For example, with faithful generative models, one can simulate population responses to hypothetical sensory, electrical, or optogenetic stimuli, as well as possible neural activity underlying hypothetical movement patterns. Most VAE-based approaches focus on the interpretability of the inferred latents, but not the ability to generate realistic and diverse samples when conditioning on external covariates, while sample-realistic models (e.g., based on generative adversarial networks (GANs) [17]) do not provide access to underlying low-dimensional representations. As such, there is a need for models of neural population spiking activity that both provide low-dimensional latent representations and can (conditionally) generate realistic neural activity.

Here, we propose Latent Diffusion for Neural Spiking data (LDNS) , which combines the ability of autoencoders to extract low-dimensional representations of discrete neural population activity, with the ability of (conditional) denoising diffusion probabilistic models (or, diffusion models) to generate realistic neural spiking data by modeling the inferred low-dimensional continuous representations.

Diffusion models [49, 20, 50] have been highly successful for conditional and unconditional data generation in several domains, including images [20], molecules [59], and audio spectrograms [27] and have demonstrated samplingfidelity that outperforms that of VAEs and GANs [20]. A key strength of diffusion models that makes them particularly attractive in the context of modeling neural datasets is the ability to flexibly condition the generation on various (potentially complex) covariates, such as to simulate neural activity given certain behaviors. Recently, diffusion models have been extended to continuous neural time series such as local field potentials (LFPs) and electroencephalography

autoencoder

noise

denoise

diffusion

spikes

spikes

rates

(optional) behavioral

conditioning

x

y

latents

encoder

decoder

evaluation

and

downstream tasks

stochastic

diffusion model

(EEG) recordings [53]. However, due to the discrete nature of spiking data, standard diffusion models cannot be easily applied, thus excluding their use on many datasets in systems neuroscience.

To bypass these limitations, LDNS employs a regularized autoencoder using structured state-space (S4) layers [18] to project the high-dimensional discrete spiking data into smooth, low-dimensional latents without making assumptions about the trial structure. We then train a diffusion model with S4 layers as a generative model of the inferred latents-akin to latent diffusion for images [45], where generation can be flexibly conditioned on behavioral covariates or task conditions.

A fundamental assumption of most low-dimensional latent variable models is that all statistical dependencies between observations are mediated by the latent space. However, in neural spiking data, there are prominent statistical dependencies that ought to persist conditional on the latent state, e.g., single-neuron dynamics such as refractory periods, burstiness, firing rate adaptation, or potential direct synaptic interactions. We show how such additional structure can be accounted for in LDNS, by equipping it with an expressive observation model [41, 32, 54, 62]: We use a Poisson model for spike generation with autoregressive couplings which are optimized post hoc to capture the temporal structure of single-neuron activity. This allows LDNS to capture a wide range of biological neural dynamics [55], with only a small additional computational cost.

Main contributions In summary, LDNS is a flexible method that allows for both high-fidelity diffusion-based sampling of neural population activity and access to time-aligned low-dimensional representations, which we validate on a synthetic dataset. Next, we show the utility and flexibility of this approach on complex real datasets: First, LDNS can handle variable-length spiking recordings from the human cortex. Second, LDNS can unconditionally generate faithful neural spiking activity recorded from monkeys performing a reach task. We demonstrate how LDNS can be equipped with an expressive autoregressive observation model that accounts for additional dependencies between data points (e.g., single neuron dynamics), increasing the realism of generated samples. Third, LDNS can generate realistic neural activity while conditioning on either reach direction or full reach trajectories (time series), including unseen behaviors that are then accurately decoded from the simulated neural data. Overall, LDNS enables simultaneous inference of low-dimensional latent representations for single-trial data interpretation and high-fidelity diffusion-based (conditional) generation of diverse neural spiking datasets, which will allow for closed-loop in silico experiments and hypothesis testing.

2 Methods

2.1 Latent Diffusion for Neural Spiking Data (LDNS)

We consider a dataset recorded from a population of n neurons, consisting of trials with spiking data s ∈ N n × T 0 (sorted into bins of fixed length resulting in spike counts over time), and optional simultaneously recorded behavioral covariates y ∈ R n (that can also be time-varying y ∈ R n × T ). A dataset of M such trials D can be written as D = { s ( i ) , y ( i ) } , possibly with varying trial lengths T 1 . . . T M . We make the assumption that a large fraction of the variability in this dataset can be captured with a few underlying latent variables z ∈ R d × T , where d < n .

Our goal is to generate realistic spiking data s ∗ that faithfully capture both population-level and singleneuron dynamics of s 1 ...T with the ability to optionally condition the generation on behavior y cond . To this end, we propose a new method, LDNS, that combines the strength of neural dimensionality reduction approaches with that of diffusion-based generation.

LDNS uses a two-stage training framework, adopted from the highly successful family of latent diffusion models (LDMs) [45, 8, 59]. To train LDNS, we first train a regularized autoencoder [14] to compress the spiking data into a low-dimensional continuous latent space (Fig. 1). Concretely, we focus on two objects of interest for the LDNS autoencoder: (1) inferring a time-aligned, lowdimensional smooth representation z ∈ R d × T that preserves the shared variability of the spiking data, and (2) predicting smooth firing rates λ that are most likely to give rise to the observed spiking data.

In the second stage, we train a diffusion model in latent space, possibly employing conditioning to make generation contingent on external (e.g., behavioral) covariates (Fig. 1). For the diffusion model, our main objective is the generation of z ∗ ∈ R d × T that captures the distribution of inferred autoencoder latents. We also want the ability to sample latent trajectories of varying length.

In both stages, we use structured state-space (S4) [18] layers for modeling temporal dependencies. S4 layers consist of state-space transition matrices that can be unrolled into arbitrary-length convolution kernels, allowing sequence modeling of varying lengths. For details on network architectures and S4 layers, see appendix A1.

2.2 Regularized autoencoder for neural spiking data

For the spiking data, we choose a Poisson observation model, and train autoencoders by minimizing the Poisson negative log-likelihood of the input spikes s given the predicted rates λ = decoder ( z ) . To enforce smoothness in the latent space, where z = encoder ( s ) , we add an L 2 regularization along with a temporal smoothness regularizer over z , resulting in the combined loss

To prevent the autoencoder from predicting highly localized Poisson rates, which have sharp peaks at input spike locations, we further regularize training using coordinated dropout [24], i.e., we randomly mask input spikes and compute the loss on the predicted rates at the masked locations (details in appendix A1.2).

Accounting for single-neuron dynamics with an expressive observation model So far, LDNS (like most latent variable models for neural data) uses a Poisson observation model, which assumes that all statistical dependencies are mediated by the latent state. To address this limitation and to capture dynamics and variability, which are 'private' to individual neurons (such as refractory periods or burstiness), we propose to learn an autoregressive observation model. We make the predicted Poisson rates for each neuron i dependent also on recent spiking history, by including additional spike history couplings h i [41, 32], resulting in the observation model

where T ′ corresponds to the time-lagged window length. This modification is learned post hoc, and the parameters h i are fit with a maximum-likelihood objective (details in appendix A1.3). This approach does not alter the latent dynamics, while augmenting the model with single-neuron autoregressive dynamics. We observe that including spike history increases the realism of generated data and enables us to accurately capture single-neuron autocorrelation structures (Sec. 3.4).

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 autoencoder-derived latents of each trial, and optionally, additional conditioning information such as the corresponding behavior, i.e., D z = { z ( i ) = encoder ( s ( i ) ) , y ( i ) } .

Diffusion models aim to approximate the data distribution q ( z ) through an iterative denoising process starting from standard Gaussian noise. For latent z (denoted as z 0 for diffusion timestep 0), we first produce a noised version at step t by adding Gaussian noise of the form q ( z t | z 0 ) = N ( √ ¯ α t z 0 , (1 -¯ α t ) I ) . Here, ¯ α t = ∏ t k =1 α k , where the noise scaling factors α 1 . . . α T follow a fixed linear schedule. We then train a neural network to approximate the reverse process p θ ( z t -1 | z t ) for each diffusion timestep. The true (denoising) reverse transition q ( z t -1 | z t ) is intractable-however, we can apply variational inference to learn the conditional reverse transition q ( z t -1 | z t , z 0 ) , which has a closed form written as

We train the neural network µ θ ( z t , t ) to approximate the mean of this distribution by optimizing the loss E z 0 ∼D z ,ϵ 0 ,t ∥ ϵ θ ( z t , t ) -ϵ 0 ∥ 2 , where ϵ 0 is the noise used to generate z t from z 0 , and ϵ θ ( z t , t ) is the equivalent reparameterization for µ θ ( z t , t ) . At test time, we sequentially sample z t -1 given z t using the learned transition p θ ( z t -1 | z t ) , starting from standard Gaussian noise. Using S4 layers in the denoising network allows us to generate latents with varying lengths. This is achieved by unrolling the state transition matrix in the S4 layers to the desired length for each denoising step.

Diffusion models may be conditioned on fixed-length and time-varying covariates y , in which case we learn the approximate reverse transition p θ ( z t -1 | z t , y ) . Details on the conditioning mechanisms in appendix A1.4.

3 Experiments and 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 a low-dimensional, non-linear dynamical system commonly used in neuroscience [7, 36]. We simulate rates as an affine mapping from the 3-dimensional system to a 128-dimensional neural space, and sample from a

a

chaoticLorenz

b

autoencoderprediction

C

diffusion sampled latents of varying length

distributionmatch

rate / spikes

4

-aerates

-spikes

log psd latents

0

4

ae

diffusion

0

10

2

10

0

time

100

0

time 100

synthetic spikes

d

populationlevelstatistics

e

128

Idns

data

0.5

Idns mean isi

15

Idns std isi

neuron id

Idns kde

datakde

supI

10

15

0.0

5

0

10

0

time

256

0

100

200

300

0.0

0.5

5

10

15

10

15

spike counts in 1 bin

gtpairwisecorrelation

gt mean isi

Poisson distribution to generate spiking data. Next, we showcase the applicability of LDNS on two neural datasets: We apply our method on a highly complex dataset of human neural activity (128 units) recorded during attempted speech [56]. This dataset poses a challenge to many modeling approaches due to the different imagined sentences, resulting in variable lengths of the neural time series (between 2-10 seconds with a sampling rate of 50 Hz). Finally, we apply LDNS to model premotor cortical activity (182 units) recorded from monkeys performing a delayed center-out reach task with barriers [9, 38]. The multi-modal nature of the dataset allows us to assess both unconditional as well as conditional generation of neural spiking activity given monkey reach directions and entire velocity profiles of the performed reaches. See appendix A2,A3 for data and training details.

For the unconditional generation of monkey reach recordings (Sec. 3.4), we train both a Poisson observation model as well as a spike history-dependent autoregressive observation model. For all other experiments, we only train a Poisson observation model.

Baselines We compare LDNS to the most commonly known VAE-based latent variable model: Latent Factor Analysis via Dynamical Systems (LFADS [51, 36, 47]), which has been shown to outperform various classical latent variable models on a variety of tasks ([38], details in appendix A4). To ensure that we use optimal hyperparameters for LFADS, we follow the auto-ML pipeline proposed by Keshtkaran et al. [25]. This approach, termed AutoLFADS, has been shown to perform better than the original LFADS on benchmark tasks [38]. For the unconditional generation of monkey reach recordings, we further compared to additional VAE baselines [21, 62] (appendix A5).

Metrics For all experiments, we assess how well LDNS-generated samples match the spiking data in the training distribution. Concretely, we compare population-level statistics by computing 1) the distribution over the population spike count, which sums up all spikes co-occurring in the population in a single time bin (i.e., spike count histogram), and 2) pairwise correlations of LDNS samples and the spiking data for each pair of neurons. For single-neuron statistics, we compare 3) the mean and 4) standard deviation of the inter-spike-interval distribution for each neuron (mean isi and std isi). When multiple spikes occur in a single time bin, the spike times are distributed equally in this bin [12]. To further evaluate population dynamics, we compare the principal components of smoothed spikes.

3.2 LDNS captures the true spiking data distribution with an underlying Lorenz system

We simulate trials of length 256 timesteps from the three-dimensional (chaotic) Lorenz system (Fig. 2a). The regularized autoencoder extracts smooth latent time series (eight latent dimensions)

aattempted speech task

b

variable-lengthtrialsindata

d

population and single-unit statistics

neuralrecording

Idns

data

Idns mean isi

0.10

Idns std isi

0.15

data kde

n!un

ldnskde

Hello

0.05

0.05

Hello

1s

50

100

150

0.050.10

0.05

0.15

。

pop. spike count

data mean isi

data std isi

C

e

pairwisecorrelations

0.5

0.5

Idns

nun

0.0

0.0

2

trial length (s)

10

data

Idns

0.0

data

0.5

from the 128-dimensional spiking data, resulting in smooth firing rate predictions that closely match the ground-truth rates (Fig. 2b, Supp. Fig. A2,A3). We then train a diffusion model on the extracted autoencoder latents. Latents sampled from the diffusion model (red) preserve the attractor geometry of the Lorenz system (Fig. 2c, left, three of the eight latent dimensions), indicating that LDNS preserves a meaningful latent space. The architectural choice of S4 layers allows for length generalization: although we train on time segments of 256-time steps, we can sample and successfully generate latent trajectories that are much longer, but still accurately reflect the Lorenz dynamics (Fig. 2c, middle, 16 × longer generation). In comparison, LFADS exhibits instabilities when generating such longer sequences (appendix A6.1). Overall, the latent time series distribution is captured well by the diffusion model, with matching power spectral densities (PSD) per latent dimension (Fig. 2c, right, other dimensions in Supp. Fig. A3).

To assess the sampling fidelity of the generated synthetic neural activity, we compute a variety of spike statistics frequently used in neuroscience. LDNS captures both population-level statistics, such as the population spike count histogram and pairwise correlations between neurons (Fig. 2d), as well as single-neuron statistics, quantified by the mean and standard deviation of inter-spike-intervals (Fig. 2e). LDNS also captures the temporal correlation structure of the data (Supp. Fig. A4). These results demonstrate that LDNS can both perform inference of low-dimensional latents and provide high-fidelity diffusion-based generation that perfectly captures the statistics of the ground-truth synthetic data.

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 it to neural recordings from human cortex during attempted speech (Fig. 3a, top, Willett et al. [56]). A participant with a degenerative disease who is unable to produce intelligible speech attempts to vocalize sentences prompted on a screen, while neural population activity is recorded from the ventral premotor cortex. Since there is a large variation in the length of prompted sentences (Fig. 3a, bottom), this dataset allows us to evaluate the performance of LDNS on real data in naturalistic settings with variable-length and highly heterogeneous dynamics.

To account for varying trial length during autoencoder training, we pad all trials to a maximum length of 512 bins and compute the reconstruction loss only on the observed time bins. For the diffusion model, we indicate the target trial length with a binary mask as a conditioning variable.

This approach allows us to infer time-aligned latents underlying the cortical activity of the participants, compressing the population activity by a factor of four before training an unconditional diffusion model on these latents. Resulting samples of LDNS, mimicking human cortical activity, are visually indistinguishable from the real data (Fig. 3b,c, additional samples in Supp. Fig. A7). This is reflected in closely matched population spike count histograms (Fig. 3d, left), and single neuron statistics such as mean and standard deviation of the inter-spike interval (Fig. 3d, right). Additionally, real

a

mazereachtask

b

data

Idns sample

C

populationlevelstatistics

0

data

neuronid

1

Idns

data kde

Idns

0

100

1

Idns kde

-TT

angle

TT

-0.05

0.0

0.5

0.0

0.5

0

10

-0.05

0

time (s)

time (s)

spikecountsin5msbins

datapairwisecorrelation

d

singleneuronstatistics

e

temporal structureauto-correlation

with spike history

Idns meani

0.2

0.2

0.0005

Idns 1 & 2

data 1& 2

data 1& 2

0.0000

0.1

0.1

-0.0005

T

T

0.1

0.2

0.1

0.2

-125

0

125

-125

0

125

datameanisi

data std isi

lag (ms)

lag (ms)

and LDNS-sampled spikes display similar population dynamics, as reflected in the top principal components (Supp. Fig. A8). While LDNS tends to overestimate some pairwise correlations, it captures prominent features of the correlation structure in the data (Fig. 3e, Pearson correlation coefficient r = 0 . 47 ), and our analysis indicates that this slight mismatch already arises at the autoencoder stage (Supp. Fig. A9).

LDNS allows for both inferring latent representations and generating variable-length trial data, making it applicable to complex real neural datasets without a fixed trial structure.

3.4 Realistic generation of spiking data from a monkey performing reach tasks

We further evaluate LDNS in a different setting by applying it to model sparse spiking data recorded from a monkey performing a reaching task constrained by barriers that form a maze (Fig. 4a, left). The variety of different maze architectures leads to diverse reach movements of both curved and straight reaches (Fig. 4a, right). We again infer low-dimensional latent trajectories that capture the shared variability of the neural population and then train an unconditional diffusion model on these latents. Sampled spikes from LDNS closely resemble the true, sparse population data (Fig. 4b, additional samples in Supp. Fig. A11), and closely match population-level spike statistics (Fig. 4c). Single neuron statistics in this low spike count regime (a maximum of three spikes per neuron in 5 ms bins) are also captured well (Fig. 4d), and are on par with or better than LFADS [36, 25] (see Table 1 for summary of main comparisons, and appendix A5 for additional baselines [21, 62]). Beyond spiking statistics, we observe that LDNS also preserves the temporal structure of population dynamics, as reflected in the top principal components of smoothed spikes (Supp. Fig. A15). Thus, LDNS can generate spiking data that is faithful at the level of both single-neuron and population dynamics.

a

behavioral conditioning

b

datavs.decoded

θ conditioning

d

Vx,V, conditioning

25

data

TT

datacondition

predicted

ae

θcondition

samples

behavior

ridge-regression

y position

decoding

0

0

rates

-25

diffusion

dec

TT

-25

0

25

-25

0

25

-25

0

25

xposition

x position

xposition

e

straighttestconditions

f

latentdim1

g

PC1

PC1:PC2

20

T

θtest reaches

sampledlatent

2

PC of latents

2

y position

PC2

0

0

0

0

0

-20

L

-2

-25

0

25

0.0

0.5

0.0

0.5

0.0

0.5

-2

0

2

x position

time [s]

time [s]

time[s]

PC1

Both LFADS and the LDNS autoencoder are optimized by maximizing the Poisson log-likelihood, and thus cannot capture single-neuron dynamics such as refractoriness [41, 55], which can have a strong influence on the observed autocorrelation structure. Given the overall sparsity of the spiking data and resulting low correlations (Supp. Fig. A10), we focus on the temporal structure of autocorrelations averaged within groups of neurons ( ≈ 45 neurons per group) split by their instantaneous correlation strength [32]: darker colors correspond to the highest correlated group of four, lighter colors correspond to the group with the second highest correlations (Fig. 4e). We then compare these auto-correlations to those of grouped LDNS samples with a Poisson observation model (red). As expected, LDNS with Poisson observations is unable to capture the dip in the data auto-correlation at 5 ms lags (one time bin) (Fig. 4e, left).

To overcome this mismatch, we train an additional spike history-dependent autoregressive observation model on top of the inferred rates (LDNSsh, for spike history). In contrast to the Poisson samples, autoregressive samples can capture this aspect of neural spiking data very accurately while also improving the overall fit to the empirical auto-correlation (Fig. 4e, right). Moreover, the post hoc optimization of these filters also improves modeling of other single-neuron, as well as population-level statistics, such as the population spike count histogram or the mean of the isi (Table 1, Supp. Fig. A13).

We view this post-hoc augmentation as a key modular contribution, which can be flexibly applied to other generative models. To this end, we extend AutoLFADS with spike history dependence (LFADSsh), improving its performance across metrics. The augmented LFADSsh also captures the dip in autocorrelation at 5 ms lags (Supp. Fig. A12). Still, in both observation model variants, LDNS maintains superior or comparable performance (Table 1).

Thus, Poisson LDNS allows for the generation of spiking data that is on par or better in terms of sampling fidelity than previous approaches. Incorporating spike-history dependence and sampling spikes autoregressively allows us to further increase the realism of generated spike trains, leading to a large improvement on several of the considered metrics.

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 (Fig. 5a). We first validate that the autoencoder predicts firing rates that allow us to linearly decode the behavior

following the ridge-regression approach proposed in [38]. Decoded behavior from autoencoder reconstructed rates matches the true trajectories of unseen test trials (Fig. 5b).

Given that the autoencoder performs adequately, we then test the ability to generate neural time series conditioned on the initial reach angle performed by the monkey θ reach . Indeed, from the generated samples of neural activity, we can decode-using the same linear decoder-realistic reach kinematics that are consistent with the conditioning angle θ reach and overall reach kinematics (Fig. 5c). This indicates that LDNS can generate realistic neural activity consistent with a queried reach direction.

An even more challenging task that is intriguing for hypothesis generation is the ability to mimic an entire experiment and ask what the neural activity would have looked like if the monkey had performed a particular hypothetical movement. To this end, we train a diffusion model on the same autoencoder-inferred latents but now condition on entire velocity traces (Fig. 5d). Velocityconditioned LDNS is able to produce different samples of neural activity that are consistent with, but not exact copies of, the reach trajectories of the held-out trials given as the conditioning covariate. Such closed-loop conditioning experiments open the possibility of making predictions about neural activity during desired unseen behaviors, and thus make experimentally testable predictions.

Finally, to understand how LDNS incorporates behavioral information, we analyzed latent trajectories that were conditionally sampled based on straight reach movements in different directions (Fig. 5e). Individual samples of latent trajectories vary smoothly within a trial (Fig. 5f), while reach direction varies smoothly across samples in the first principal component (PC1) of the latents (Fig. 5g, left). Projection onto the first two PCs of latent trajectories shows clear clustering by reach direction (Fig. 5g, right), and we show that such clustering arises already at the autoencoder stage (Supp. Fig. A17).

In summary, LDNS not only produces faithful spiking samples but also allows for flexible conditioning. Furthermore, LDNS learns an interpretable latent space with behaviorally-relevant structure.

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 [61, 32, 39, 58, 62, 13, 28, 63] (see [38] for a comprehensive list.) In addition to capturing shared population-level dynamics and dependence on external stimuli [4], LVMs have been extended to allow autoregressive neuron-level (non-Poisson) dynamics [32, 13, 62] or even direct neural interactions [54]. While these methods often have useful inductive biases (e.g., linear dynamical systems [32, 28] or Gaussian process priors [61]), these models are typically not expressive enough to yield realistic neural samples across a range of conditions.

Deep LVMs and other deep learning-based approaches Variational autoencoders (V AEs) [26] are particularly popular in neuroscience as they allow us to infer low-dimensional dynamics underlying high-dimensional discrete data [63, 16, 46], especially when combined with nonlinear recurrent neural networks [36, 21]. VAEs have been used to infer identifiable low-dimensional latent representations conditioned on behavior [63, 21] and have incorporated smoothness priors using Gaussian Processes to regularize the latent space [16]. However, the generation performance of V AEs is rarely explored in neuroscience. Besides VAEs, generative adversarial networks (GANs [17]) have been proposed to synthesize spiking neural population activity [34, 42]. While GANs produce high-fidelity samples, they are challenging to train reliably and lack a low-dimensional latent space. More recently, transformer-based architectures have also been adapted to model neural activity [5, 60], though often with the focus of accurate decoding of behavior instead of generation of realistic spiking samples, while also lacking an explicit latent space [3]. Lastly, deterministic approaches utilizing RNNs for dynamical systems reconstruction also target low-dimensional latent dynamics underlying neural data [19], but they do not act as probabilistic generative models.

Diffusion models LDNS leverages recent advances in diffusion models, which have become state-of-the-art for high-fidelity generation in several domains [20, 27], including continuous-valued neural signals such as EEG [53], as well as in time series forecasting and imputation tasks [2, 52, 43]. Similar to the LDNS architecture, Alcaraz and Strodthoff [2] also use an S4-based denoiser for imputation. More specifically, LDNS is inspired by latent diffusion models [45, 27, 59, 15], which benefit from operating on the latent space of an autoencoder and flexible conditioning mechanisms to generate samples based on a given covariate, as is done with text-to-image [45] and other cross-

modality scenarios. Conveniently, this allows LDNS to bypass the challenges of directly modeling discrete-valued spiking data, by instead transforming spikes into the continuous latent space.

5 Summary and discussion

We here proposed LDNS, a flexible generative model of neural spiking recordings that simultaneously infers low-dimensional latent representations and generates realistic neural activity conditioned on behavioral covariates. We apply LDNS to model three different datasets: synthetic data simulated from chaotic Lorenz dynamics, human cortical recordings with heterogeneous and variable-length trials, and finally, neural recordings in monkeys performing reach actions in a maze. Through our experiments, we demonstrate how several features of LDNS are beneficial for modeling complex datasets in neuroscience:

First, following other LDMs in the literature, LDNS decouples latent inference and probabilistic modeling of the data, offering flexibility in reusing the trained autoencoder and diffusion model. For the monkey recordings, all diffusion models (unconditional, conditioned on reach angle, and conditioned on hand velocities) operate in the latent space of the same autoencoder, in contrast to existing approaches that require end-to-end retraining for each type of conditioning variable. LDNS is also faster to train than AutoLFADS, which requires population-based training to optimize hyperparameters (appendix A3.1). Second, we show that LDNS autoencoders can be augmented with per-neuron autoregressive dynamics to capture single-neuron temporal dynamics (e.g., refractoriness), which otherwise cannot be captured with population-level shared dynamics. Third, as a result of the length-generalizable autoencoders and diffusion models using S4 layers, LDNS can generate variable-length trials in both the Lorenz example and human cortical recordings-a feature that will be particularly useful in modeling datasets recorded during naturalistic stimuli or behavior.

Altogether, these features enable LDNS to generate realistic neural activity, especially when conditioned on behavioral covariates. In our experiments, we demonstrate that unseen movement trajectories can be used to conditionally generate samples of neural activity, from which we can decode these hypothetical behaviors. These generated latent trajectories reflect behavioral information in an interpretable way. Our methodology is general and can be applied to recordings from any brain region, beyond the motor and speech cortex examples shown here. Thus, LDNS opens up further possibilities for hypothesis generation and testing in silico , potentially enabling stronger links between experimental and computational works.

Limitations In real neural data, the latent dimensionality of the system is not known, and as with all LVMs (which often assume that population dynamics are intrinsically low-dimensional), choosing an appropriate latent dimension can be challenging. Furthermore, any modeling errors at the encoding and decoding stage of the autoencoder will affect the overall performance of the latent diffusion approach. Nevertheless, in our experiments, we found that autoencoder training is fast, stable, and reasonably robust to hyperparameter configurations. While LDNS was still able to model the data well under relatively severe compression (e.g., 182-to-16 for the monkey recordings), optimizing latent dimensionality to balance expressiveness and interpretability remains a goal for future research.

Broader impact Realistic spike generation capabilities increase the risk of research manipulation by generating synthetic data that may be difficult to detect. On the other hand, LDNS could be useful for the dissemination of privatized clinical data, though we acknowledge the critical importance of protecting data privacy when working with sensitive human participant data. Finally, synthetically generated data (conditioned on unseen behavioral conditions) could be useful for augmenting the training of brain-computer interface decoding models.

References

Misha B. Ahrens, Michael B. Orger, Drew N. Robson, Jennifer M. Li, and Philipp J. Keller. Whole-brain functional imaging at cellular resolution using light-sheet microscopy. Nature Methods , May 2013.

Juan Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecasting with structured state space models. Transactions on Machine Learning Research , 2022.

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.

Evan W Archer, Urs Koster, Jonathan W Pillow, and Jakob H Macke. Low-dimensional models of neural population activity in sensory cortical circuits. Advances in neural information processing systems , 27, 2014.

Mehdi Azabou, Vinam Arora, Venkataramana Ganesh, Ximeng Mao, Santosh Nachimuthu, Michael Mendelson, Blake Richards, Matthew Perich, Guillaume Lajoie, and Eva Dyer. A Unified, Scalable Framework for Neural Population Decoding. Advances in Neural Information Processing Systems , 2023.

Mohammad Bashiri, Edgar Walker, Konstantin-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.

Manuel Brenner, Florian Hess, Georgia Koppe, and Daniel Durstewitz. Integrating multimodal data for joint generative modeling of complex dynamics, 2024.

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.

Mark Churchland and Matthew Kaufman. Mcmaze: macaque primary motor and dorsal premotor cortex spiking activity during delayed reaching (version 0.220113.0400) [data set], 2022. Version 0.220113.0400.

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.

John P Cunningham and Byron M Yu. Dimensionality reduction for large-scale neural recordings. Nature neuroscience , 2014.

Kevin Doran, Marvin Seifert, Carola A. M. Yovanovich, and Tom Baden. Distance function for spike prediction, 2023.

Lea Duncker and Maneesh Sahani. Temporal alignment and latent gaussian process factor inference in population spike trains. In Advances in Neural Information Processing Systems , 2018.

Partha Ghosh, Mehdi S. M. Sajjadi, Antonio Vergari, Michael Black, and Bernhard Scholkopf. From variational to deterministic autoencoders. In International Conference on Learning Representations , 2020.

Karan Goel, Albert Gu, Chris Donahue, and Christopher R'e. It's raw! audio generation with state-space models. International Conference on Machine Learning , 2022.

Rabia Gondur, Usama Bin Sikandar, Evan Schaffer, Mikio Christian Aoi, and Stephen L. Keeley. Multimodal Gaussian Process Variational Autoencoders for Neural and Behavioral Data. In International Conference on Learning Representations , 2024.

Ian Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio. Generative adversarial nets. In Advances in Neural Information Processing Systems , 2014.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. International Conference on Learning Representations , 2022.

Florian Hess, Zahra Monfared, Manuel Brenner, and Daniel Durstewitz. Generalized teacher forcing for learning chaotic dynamics. International Conference for Machine Learning , 2023.

Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems , 2020.

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.

Kristopher Jensen, Ta-Chu Kao, Jasmine Stone, and Guillaume Hennequin. Scalable Bayesian GPFA with automatic relevance determination and discrete noise models. In Advances in Neural Information Processing Systems , 2021.

James J. Jun, Nicholas A. Steinmetz, Joshua H. Siegle, et al. Fully integrated silicon probes for high-density recording of neural activity. Nature , 2017.

Mohammad Reza Keshtkaran and Chethan Pandarinath. Enabling hyperparameter optimization in sequential autoencoders for spiking neural data. Advances in neural information processing systems , 32, 2019.

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.

Diederik P. Kingma and Max Welling. Auto-encoding variational Bayes. In International Conference on Learning Representations , 2014. doi: 10.48550/arXiv.1312.6114.

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.

Scott Linderman, Matthew Johnson, Andrew Miller, Ryan Adams, David Blei, and Liam Paninski. Bayesian Learning and Inference in Recurrent Switching Linear Dynamical Systems. In Proceedings of the 20th International Conference on Artificial Intelligence and Statistics , 2017.

Edward N. Lorenz. Deterministic Nonperiodic Flow. Journal of the Atmospheric Sciences , 20(2):130-141, March 1963.

Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. International Conference on Learning Representations , 2019.

Ilya Loshchilov and Frank Hutter. Decoupled Weight Decay Regularization. In International Conference on Learning Representations , 2019. doi: arXiv:1711.05101.

Jakob H Macke, Lars Buesing, John P Cunningham, Byron M Yu, Krishna V Shenoy, and Maneesh Sahani. Empirical models of spiking in neural populations. In Advances in Neural Information Processing Systems , 2011.

Bartul Mimica, Tuçe Tombaz, Claudia Battistin, Jingyi Guo Fuglstad, Benjamin A. Dunn, and Jonathan R. Whitlock. Behavioral decomposition reveals rich encoding structure employed across neocortex in rats. Nature Communications , 2023.

Manuel Molano-Mazon, Arno Onken, Eugenio Piasini, and Stefano Panzeri. Synthesizing realistic neural population activity patterns using generative adversarial networks. International Conference on Learning Representations , 2018.

Joseph E. O'Doherty, Mariana M. B. Cardoso, Joseph G. Makin, and Philip N. Sabes. Nonhuman Primate Reaching with Multichannel Sensorimotor Cortex Electrophysiology, 2017.

Chethan Pandarinath, Daniel J. O'Shea, Jasmine Collins, Rafal Jozefowicz, Sergey D. Stavisky, Jonathan C. Kao, Eric M. Trautmann, Matthew T. Kaufman, Stephen I. Ryu, Leigh R. Hochberg, Jaimie M. Henderson, Krishna V. Shenoy, L. F. Abbott, and David Sussillo. Inferring single-trial neural population dynamics using sequential auto-encoders. Nature Methods , 15, 2018.

William S. Peebles and Saining Xie. Scalable diffusion models with transformers. IEEE International Conference on Computer Vision , 2022.

Felix C Pei, Joel Ye, David M. Zoltowski, Anqi Wu, Raeed Hasan Chowdhury, Hansem Sohn, Joseph E O'Doherty, Krishna V. Shenoy, Matthew Kaufman, Mark M Churchland, Mehrdad Jazayeri, Lee E. Miller, Jonathan W. Pillow, Il 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, Datasets and Benchmarks Track (Round 2) , 2021.

Biljana Petreska, Byron M Yu, John P Cunningham, Gopal Santhanam, Stephen Ryu, Krishna V Shenoy, and Maneesh Sahani. Dynamical segmentation of single trials from population neural data. In Advances in Neural Information Processing Systems , 2011.

David Pfau, Eftychios A Pnevmatikakis, and Liam Paninski. Robust learning of low-dimensional dynamics from large neural ensembles. In Advances in Neural Information Processing Systems , 2013.

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.

Poornima Ramesh, Mohamad Atayi, and Jakob H Macke. Adversarial training of neural encoding models on population spike trains. Real Neurons & Hidden Units: Future directions at the intersection of neuroscience and artificial intelligence @ NeurIPS 2019 , 2019.

Kashif Rasul, Calvin Seward, Ingmar Schuster, and Roland Vollgraf. Autoregressive denoising diffusion models for multivariate probabilistic time series forecasting. International Conference on Machine Learning , 2021.

Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In Proceedings of the 31st International Conference on Machine Learning , 2014.

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.

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.

Andrew R. Sedler and Chethan Pandarinath. lfads-torch: A modular and extensible implementation of latent factor analysis via dynamical systems. arXiv preprint arXiv:2309.01230 , 2023.

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.

Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning , 2015.

Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. International Conference on Learning Representations , 2021.

David Sussillo, Rafal Jozefowicz, L. F. Abbott, and Chethan Pandarinath. Lfads - latent factor analysis via dynamical systems, 2016.

Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. CSDI: Conditional score-based diffusion models for probabilistic time series imputation. Advances in Neural Information Processing Systems , 2021.

Julius Vetter, Jakob H Macke, and Richard Gao. Generating realistic neurophysiological time series with denoising diffusion probabilistic models. bioRxiv , 2023.

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 on the network activity of retinal ganglion cells. Journal of Computational Neuroscience , 2012.

Alison I. Weber and Jonathan W. Pillow. Capturing the Dynamical Repertoire of Single Neurons with Generalized Linear Models. Neural Computation , 2017.

Francis Willett et al. Data for: A high-performance speech neuroprosthesis [dataset], 2023. URL https://doi.org/10.5061/dryad.x69p8czpq .

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 V. Shenoy, and Jaimie M. Henderson. A high-performance speech neuroprosthesis. Nature , 2023. Publisher: Nature Publishing Group.

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.

Minkai Xu, Alexander Powers, R. Dror, Stefano Ermon, and J. Leskovec. Geometric latent diffusion models for 3d molecule generation. International Conference on Machine Learning , 2023.

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.

Byron M. Yu, John P. Cunningham, Gopal Santhanam, Stephen I. Ryu, Krishna V. Shenoy, and Maneesh Sahani. Gaussian-Process Factor Analysis for Low-Dimensional Single-Trial Analysis of Neural Population Activity. Journal of Neurophysiology , 2009.

Yuan Zhao and Il Memming Park. Variational latent gaussian process for recovering single-trial dynamics from population spike trains. Neural Computation , 2017.

Ding Zhou and Xue-Xin Wei. Learning identifiable and interpretable latent models of high-dimensional neural activity using pi-V AE. In Advances in Neural Information Processing Systems , 2020.

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)

A central component of our autoencoder architecture is the recently introduced structured state space models (S4)[18]. With an input sequence x = [ x 1 . . . x T ] ∈ R T and corresponding output y = [ y 1 . . . y T ] ∈ R T , an S4 layer applies the following operation for each timestep -

where the discretized state and input matrices A,B given continuous analogues A,B and step size ∆ are computed as

When the state s t is not required, this recurrent computation of the output y given input sequence y can be unrolled into a parallelizable convolution operation

We used the S4 implementation 1 provided by Gu et al. [18] that stably initializes the state transition matrix A using a diagonal plus low-rank approximation. For a multivariate input-output pair x , y ∈ R D × T , we apply D separate univariate S4 layers for each dimension and then mix them in the channel-mixing layer using an MLP (see next section). Each univariate input-output mapping consists of H separate S4 'head' that are expanded and contracted from and into a single dimension.

Due to its recurrent nature, S4 is a causal layer, enabling variable-length training and inference. To enable bidirectionality, we flip the input signal x 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

Weinclude temporal information only in the encoder and model the decoder as a lightweight pointwise MLP for the autoencoder (Supp. Fig. A1). 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 a fixed length and only backpropagate through the unpadded output rates.

Furthermore, to infer smooth rates and avoid spiking behavior, we use coordinated dropout [24]. For each time bin independently, we mask the input spikes to zero with random probability p and scale up the remaining spikes by 1 1 -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 a spiking prediction of the Poisson rates.

x

[B,C,L]

+

+

LayerNorm

Time mixing

LayerNorm

Channel mixing

(Linear

→

GELU

→

Linear)

x

[B,C,L]

AE block

AE block

MLPCin

→

C

h

Latents

MLPCh

→

C

L

Pointwise

MLPCL

→

C

in

Spikes

Rates

denoiser

block

MLPCL

→

C

Latents

t-1

MLPC

→

C L

Latents

t

t

conditioning

denoiser

block

MLP

x

[B,C,L]

c

[B,C']

+

+

Scale and Shift

Time mixing

(

bidirectional

Scale and Shift

Channel mixing

(Linear

→

GELU

→

Linear)

x

[B,C,L]

MLPC'

→

C

Scale

Scale

MLPC'

→

2C

MLPC'

→

2C

MLPC'

→

C

MLPC'

→

C

A1.3 Spike-history-augmented 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 a history length of T ′ = 20 , corresponding to 100 ms, using the AdamW optimizer [30] (learning rate 0.1, weight decay 0.01). In our implementation, we use the Softplus function given by Softplus( x ) = log(1 + exp( x )) as an approximation to the exponential function in Eq. (2), which is accurate for the low-count regime while increasing numerical stability. During autoregressive sampling, we limit the maximum possible spike count to 5 spikes, which corresponds to the biological maximum, limited by the refractory period for 5 ms time bins.

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. A1). For diffusion timestep and fixed-length conditioning vector, we shift and scale the inputs and outputs to the time mixing and channel mixing blocks using adaptive instance normalization, as done in Peebles and Xie [37].

Unconditional generation for synthetic spiking data with Lorenz Dynamics and cortical spiking data in monkeys - we use only time conditioning.

Angle-conditioned generation for cortical monkey spiking data - we add an embedding MLP ([cos θ, sin θ ]) of the reach angle θ to the timestep embedding output.

Trajectory conditioned generation for cortical monkey spiking data - we concatenate the hand velocities v x , v y of the monkey with the input as two additional channels.

Unconditional variable-length generation for cortical human spiking data - we concatenate the desired length (with a maximum sequence length of 512) as a centered binary mask channel in the input. We only backpropagate through the central section of the output corresponding to the binary mask.

We use a DDPM scheduler with 1000 timesteps and ϵ -parameterization. To stabilize and speed up training, we train all diffusion models using a smooth L 1 loss, written as

A2 Dataset access information

All real-world datasets used in this work are publicly available under open-access licenses. Our work does not involve the collection of new experimental data.

The human BCI dataset is available at https://datadryad.org/stash/downloads/file_ stream/2547369 under a CC0 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.org/dandiset/000128 , ID: 000128) under a CC-BY-4.0 license. This dataset contains sorted unit spiking times and behavioral data from primary motor and dorsal premotor cortex during a delayed reaching task.

A3 Hyperparameters and compute resources

A3.1 Computational Resources

We performed all training and evaluation of LDNS on the Lorenz and Monkey reach datasets on an NVIDIA RTX 3090 GPU with 24GB RAM. For the Human BCI data, we used an NVIDIA A100 40GB GPU.

The autoencoder for the Lorenz dataset is trained in ≈ 6 minutes, and the diffusion model in ≈ 20 minutes. For the evaluation, all sampling is performed on the GPU in 5 minutes. The effective GPU wallclock time (time when the GPU is utilized) for the entire training and evaluation run is within 30 minutes.

For the Monkey reach dataset, the autoencoder with the given hyperparameters is trained in ≈ 8 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 is approximately within one hour. Optimizing the autoregressive observation model took less than 1 minute.

AutoLFADS, the baseline used for unconditional sampling for the Monkey reach dataset, was trained on a cluster of 8 NVIDIA RTX 2080TI GPUs for one day. As it requires automated hyperparameter tuning to achieve the best accuracy using population-based training (PBT, [24]), AutoLFADS is significantly more compute-expensive to train than LDNS.

For the Human BCI dataset, due to larger trial lengths, more data points, and more heterogeneous temporal dynamics, we trained a slightly larger autoencoder and diffusion model than in Monkey reach. The autoencoder took 50 minutes to train, and the diffusion model took 10 hours to train. Sampling from the trained model took 9 minutes, resulting in a total of under 12 hours of effective GPU wallclock time.

We ran several preliminary experiments for LDNS to optimize the architecture and hyperparameters, as well as for designing appropriate evaluations. We estimate the total effective GPU wallclock time to be ≈ 10 × that of the final model runs. The AutoLFADS baseline was only trained once with PBT, as this framework automatically optimizes the model hyperparameters.

We implemented all training and evaluation code using the Pytorch framework 2 , and used Weights & Biases 3 to log metrics during training.

A4 Baseline comparison: Latent Factor Analysis via Dynamical Systems LFADS

Latent Factor Analysis via Dynamical Systems (LFADS) is a sequential variational autoencoder used to infer latent dynamical systems from neural population spiking activity [51, 36]. LFADS consists of an encoder, a generator, and optionally, a controller, all of which are RNNs. The generator RNN implements the learned latent dynamical system, given an initial condition and time-varying inputs. The internal states of the generator are mapped through affine transformations to lower-dimensional latent factors and single-neuron Poisson firing rates. The encoder RNN maps the neural population activity into an approximate posterior over the generator's initial condition.

At each timestep, the controller RNN receives both encoded neural activity and the latent factors from the previous timestep and outputs an approximate posterior over the input to the generator. The entire model is trained end-to-end to maximize the ELBO, as is done in VAEs. To address the difficulty of hyperparameter optimization for LFADS, Population-Based Training (PBT) has been proposed to automate hyperparameter selection, termed AutoLFADS [25].

In our experiments with the monkey reach dataset, we use the PyTorch implementation of AutoLFADS [47]. We use the hyperparameters and search ranges from Pei et al. [38], but omit the controller RNN to simplify generation from prior samples. Although this might limit the model's expressiveness, prior research indicates that the monkey reach data can be well-modeled as autonomous, without external inputs from the controller [10]. LFADS has previously performed well on this data without the controller [36].

We generate samples from LFADS by sampling initial conditions from the Gaussian prior, running the generator RNN forward, and Poisson-sampling spikes from the resulting firing rates. For inclusion of spike history in the observation model of LFADS, we used the same training method and hyperparameter settings as in LDNSsh (appendix A1.3).

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, [21]) and Poisson-identifiable VAE (pi-VAE, [63]). It is important to note that while both TNDM and pi-VAE have demonstrated success in analyzing neural and behavioral data, neither was specifically designed for realistic spike train generation. The architectural choices in our implementation of these methods reflect their original intended applications in neural data analysis rather than generation of neural spiking data. Nevertheless, our comparisons show that LDNS, especially with spike-history, is superior or on par with all other methods (Table 4).

A5.1 Targeted Neural Dynamical Modeling (TNDM)

TNDM [21] is a VAE-based model designed to jointly model neural activity and behavior. TNDM extends LFADS by using an RNN to generate latent dynamics that are mapped to both neural activity and behavioral variables. TNDM separates the latent space into behavior-specific and behaviorindependent subspaces to disentangle task-relevant and intrinsic neural dynamics.

For our comparison on the unconditional monkey reach task, we trained TNDM using the architecture and hyperparameters proposed in the original implementation of the paper 4 . We used 64-dimensional latent dynamics for each of the two sets. These project to a total of 10 latent factors z (5 behaviorspecific z r and 5 behavior-independent z i ), which is the maximum number demonstrated in the original work.

To generate unconditional samples, we sampled initial generator states from a standard normal prior N (0 , I ) , then generated the latent dynamics and projected into neuron rates the same way as in LFADS. TNDM, performs well in matching real data in spike statistics (Supp. Fig. A13 cyan, Supp. Fig. A16e) and temporal dynamics (Supp. Fig. A15). Overall, we observe that LDNS captures spike statistics better than TNDM, except for the the population spike history count. In all metrics, LDNS augmented with spike history outperforms TNDM on spike statistics.

A5.2 Poisson-identifiable VAE (pi-VAE)

Poisson-identifiable V AE (pi-V AE) [63] is a V AE-based model for count data that ensures identifiability in the latent space. pi-V AE does not model temporal dependencies, instead treating each time point as an independent sample.

We trained pi-VAE on the monkey dataset using the original architecture and hyperparameters 5 . We used a General Incompressible-flow Network as a decoder, with 2 behaviorally relevant dimensions and 2 independent dimensions in the latent space. However, our evaluation context differs significantly from the original paper's demonstrations: while pi-VAE was initially evaluated on 50ms time bins and straight reaches only, our comparison uses 5ms bins and conditions on angles across all reaches at both middle and end trajectory points. The 'label', or behavior, is presented as a 4-dimensional vector containing the cosine and sine of initial and final reach angles. Since this has a conditional latent space, sampling is performed by sampling angles randomly.

Importantly, sampling from pi-VAE does not introduce any temporal dependence between spike bins within a trial - pi-V AE was not intended to be a generative model of neural spiking data. The lack of temporal modeling in pi-VAE's is a fundamental limitation for generating realistic spike trains, as evident in our empirical comparisons (Supp. Fig. A13 in yellow, Supp. Fig. A15,A16f). Note that this failure cannot be diagnosed simply from looking at the sampled spiking data (Supp. Fig. A14).

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)-a task often ignored by other LVMs designed for neural data analysis such as LFADS, pi-VAE, and TNDM.

The S4 autoencoder and diffusion model in LDNS are trained in separate stages, offering modularity, while both components naturally account for temporal dependencies (unlike pi-VAE).

S4 is autoregressive, similar to other RNN-based models, but empirically we found it to perform better when extending past the training trial length (compared to LFADS, see Sec. A6.1).

One feature provided by some neural-behavioral analysis models (such as pi-VAE and TNDM) is an explicit disentangling of neural vs. behavior-relevant latents. While we found LDNS latents contain relevant behavioural information (Fig. 5e-g, Supp. Fig. A17), we did not explicitly supervise the latent space to induce this property.

Finally, the spike history-dependent observation model in LDNS is modular and can be optimized post-hoc using rate predictions of any model to improve spike generation quality. We observed this with LDNS as well as LFADS (Table 1).

A6 Supplementary Figures Lorenz

0

1

2

0

1

2

0

50

100

time (s)

0.0

0.2

0.4

0.6

rate (Hz)

0

50

100

0.0

0.2

0.4

pred

real

error

a

0

100

200

0

50

100

0

100

200

time (s)

neuron id

b

rate (Hz)

10

-2

10

-1

frequency [Hz]

10

-1

10

2

power [dB]

ae

diff

10

-2

10

-1

10

-1

10

2

10

-2

10

-1

10

-1

10

2

10

-2

10

-1

10

-1

10

2

10

-2

10

-1

10

-1

10

2

10

-2

10

-1

10

-2

10

0

10

-2

10

-1

10

-2

10

0

10

-2

10

-1

10

-1

10

2

-2

0

latent 0

latent 1

0

100

200

0.0

0.5

latent 2

0

100

200

latent 3

0

100

200

-2.5

0.0

latent 4

0

100

200

latent 5

0

100

200

time (s)

-2.5

0.0

latent 6

0

100

200

time (s)

latent 7

inferred latents

b

a

a

gt correlation

autoencoder

Idns

gt rates vs. model

1.0

neuron id

100

0190

100

口

100

model corr

0.5

·ae

0.0

50

50

50

-0.5

0

0

0

-1.0

0

100

0

100

0

100

-1

0

neuronid

b

C

temporal auto-corr

temporal cross-corr

0.15

I1

gt1

0.075

gt 2

auto-corr

0.10

0.050

1

gt 3

gt 4

0.05

0.025

0.000

0.00

-20

0

20

-20

0

20

lag

lag

A6.1 Length generalization of LFADS on Lorenz

To analyze whether LFADS [36] exhibits similar length generalization properties as LDNS, we trained an AutoLFADS model on the Lorenz dataset (256 bins). We used the same architecture as the Monkey dataset, with 40-dimensional latent dynamics. We sampled initial conditions from the LFADS prior, then generated dynamics for both the original 256 steps and for an extended duration of 16 × the training length.

Qualitatively, we observed that while LFADS produced trajectories resembling the attractor dynamics of the ground truth Lorenz system (Supp. Fig. A5a, across various dimension combinations), these trajectories often diverged when run for longer intervals (Supp. Fig. A5b). However, the system eventually returned to typical dynamics.

Furthermore, when generating for extended durations, we observed that the mean population firing rates sometimes reached extreme values in some samples (Supp. Fig. A6), though they eventually returned to typical ranges. This behavior was not observed in LDNS samples, suggesting that bidirectional generation in the diffusion model provides more stability in variable-length generation.

a

Ifads latents original length

2

31

b

Ifads latents exceeding original length

.

0

2000

4000

time steps [a.u.]

0.0

0.5

1.0

mean rate

sample 1

0

2000

4000

time steps [a.u.]

0

20

40

mean rate

sample 2

0

2000

4000

time steps [a.u.]

0

20

mean rate

sample 3

0

2000

4000

time steps [a.u.]

0.0

0.5

1.0

mean rate

sample 4

0

2000

4000

time steps [a.u.]

0.0

0.5

1.0

mean rate

sample 5

0

2000

4000

time steps [a.u.]

0.5

1.0

mean rate

sample 6

0

2000

4000

time steps [a.u.]

0

20

mean rate

sample 7

0

2000

4000

time steps [a.u.]

0

20

40

mean rate

sample 8

A7 Supplementary Figures Human BCI

real dataspikes

Idns sampled spikes

P!!!un

unit id

unit id

unit id

unit id

1s

1s

data

Idns

10

PC1

0

-10

10

PC2

0

-10

0.0

5.0

10.0 0.0

5.0

10.0

time (s)

time (s)

gtneuroncorrmatrix

aeneuroncorrmatrix

diffneuroncorrmatrix

neuronvsneuroncorrelation

0.5

0.5

Idns

0.0

0.0

ae

0.0

0.5

gt

A8 Supplementary Figures Monkey Reach task

gt correlation

autoencoder

Idns

gt vs.model

150

150

150

0.03

corr

neuron

100

100

100

0.00

moc

p

50

50

50

-0.03

ae

Idns

0

0

0

0

100

0

100

0

100

-0.03

0.00

0.03

neuronid

data

neuron

100

neuron

100

0

neuron

100

0

neuron

100

0

neuron

100

0.0

0.5

0.0

0.5

time (s)

time (s)

IfadsPoisson

Ifads with spike history

0.0005

auto-corr

0.0000

data 2

data

1

data 1

data 2

-0.0005

Ifads 1

Ifadssh 2

lfads 2

Ifadssh 2

-25

0

25

-25

0

25

lag

lag

tndm

lfads lfads

0.000

0.002

0.004

0.006

value (lower is better)

tndm

lfads lfads

shldns ldnssh

0.000

0.001

0.002

0.003

tndm

lfads lfads

shldns

ldnssh

0.00

0.02

0.04

0.06

tndm

lfads lfads

0.00

0.01

0.02

0.03

ours

ours

ours

ours

D KL

pop. spike count hist

RMSE pairwise corr

RMSE mean isi

RMSE std isi

data

Idns

Idnssh

Ifads

Ifadssh

tndm

neuron

time

0.0

0.5

-5

0

5

PC1

ldns

0.0

0.5

ldnssh

0.0

0.5

lfads

0.0

0.5

lfadssh

0.0

0.5

tndm

0.0

0.5

0.0

0.5

time [s]

-5

0

5

PC2

0.0

0.5

0.0

0.5

0.0

0.5

0.0

0.5

0.0

0.5

10

1

10

2

frequency [Hz]

10

-

5

PC 1 power [dB]

data

10

1

10

2

frequency [Hz]

PC 1 power [dB]

10

1

10

2

frequency [Hz]

PC 1 power [dB]

10

1

10

2

frequency [Hz]

PC 1 power [dB]

10

1

10

2

frequency [Hz]

PC 1 power [dB]

10

1

10

2

frequency [Hz]

PC 1 power [dB]

ldns

ldnssh lfads

10

-

4

10

-

3

10

-

2

MSE median psd

PC1

ldns

ldnssh lfads

10

-

4

10

-

2

MSE median psd

PC2

a

b

c

a

0.2

b

0.2

data

data

density

Idns

density

Idnssh

Idnssh

0.1

1

Idns kde

data kde

Idns

0.00

0.1

hist kde

datakde

0.00

0.0

-0.05

0.0

-0.05

0

10

-0.05

0.00

0

10

-0.05

0.00

spike count

16

spike count

16

Idnssh mean isi

Idnssh std isi

Idns mean isi

0.2

Idns std isi

0.2

0.2

0.2

0.1

0.1

0.1

0.1

0.1

0.2

0.1

0.2

0.1

0.2

0.1

0.2

gt mean isi

gt mean isi

C

d

0.2

data

0.2

data

density

Ifads

data kde

Ifads

0.00

density

Ifadssh

datakde

Ifadssh

0.00

0.1

Ifads kde

0.1

Ifadssh kde

0.0

-0.05

0.0

-0.05

0

10

-0.05

0.00

0

10

-0.05

0.00

spike count

6

spike count

16

Ifads mean isi

0.2

Ifads std isi

0.2

Ifadssh mean isi

0.2

Ifadssh std isi

0.2

0.1

0.1

0.1

0.1

0.1

0.2

0.1

0.2

0.1

0.2

0.1

0.2

gt mean isi

gt mean isi

e

f

0.2

data

0.2

data

density

11

tndm

data kde

tndm

0.00

density

11

data kde

0.00

0.1

0.1

pivaekde

0.0

-0.05

0.0

-0.05

0

10

-0.05

0.00

0

10

-0.05

0.00

T

spike count

16

spike count

16

0.6

0.2

pivae mean isi

0.3

pivae std isi

0.2

0.4

0.2

0.2

0.1

0.1

0.0

T

0.0

0.25

0.50

0.0

0.2

0.1

0.2

0.3

0.0

0.2

gt mean isi

gt mean isi

b

Idns latent 0

Idns latent 1

C

IdnsPC1:PC2

sampledlatent

PC of latents

2

a

straighttestconditions

0

0

0

20

TT

testreaches

-2

position

0.0

0.5

0.0

0.5

0.0

0.5

0.0

0.5

-2

0

2

0

0

time [s]

time [s]

time [s]

time [s]

PC1

>

d

aelatent0

aelatent1

e

aePC1

aePC1:PC2

-20

T

LL

autoencoderlatent

2

PC of latents

2

-25

0

25

xposition

0

0

P

0

0

-2

0.0

0.5

0.0

0.5

0.0

0.5

0.0

0.5

-2

0

T

2

time [s]

time [s]

time [s]

time [s]

PC1

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 we justify our claims by evaluating our method's performance on synthetic and real-world data.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of our work in a paragraph in the last section of the main paper.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [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 crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: In the Appendix, we provide the full architecture of the proposed LDNS model, along with hyperparameter configurations for the model, data, and training algorithm. We also make the training loss explicit for the autoencoder and diffusion model. Finally, we provide a Github repository https://github.com/mackelab/LDNS with our code implementation.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in

some way (e.g., to registered users), but it should be possible for other researchers to have some path to 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 instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We have open sourced the code and instructions for reproducing our experiments on our Github repository - https://github.com/mackelab/LDNS. The data is open access, and we provide links to downloading the data.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) 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 new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide a table with all relevant hyperparameters in the Appendix. We also provide yaml configuration files in the zipped code directory.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

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 1 in main paper, Table 4 in appendix), we report performance and standard deviation over 5 sets of samples.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 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 error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We list the hardware configurations used for training and evaluation, as well as estimates on wallclock time for this hardware in the Appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We have read and confirmed that we conform to all the points mentioned in the NeurIPS Code of Ethics, in every respect.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We discuss the broader impact of our work in the last section of the main paper. We do not foresee any negative impact of our work, apart from possibly synthesizing scientific or clinical data with dishonest intent.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: In our opinion, we do not foresee a high risk of misuse of our work.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We cite the real-world datasets that we use in the main paper (CC-BY-4.0 for the Monkey reach dataset, CC0 1.0 for the Human BCI dataset.) We have made sure to not use any code or data whose license prevents us from using it. (Apache-2.0 for code for the S4 layers, MIT license for Weights & Biases, BSD-3 Clause for Pytorch, Apache-2.0 license for the baseline implementation LFADSTorch.)

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 is used and, if 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, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We provide a zipped directory of the code we use for training and evaluating our models and baselines. As previously mentioned in questions 4,5 of the checklist, we provide instructions to understand and run the code for reproducibility.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. 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 as 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 of human BCI data in one experiment. Both ethics reviewers found no issues with our paper, as this data was CC0 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or 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

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: We did not interact with human subjects for any experiments or data collection in this work.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nAs 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.\n\nLIMITATIONS:\nI think the authors have adequately addressed limitations and possible negative societal impacts of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. There is only one comparison method in this paper.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThis 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nWhile 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.\n\nQUESTIONS:\n- 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.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nI 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.\n\nWEAKNESSES:\nI 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\"\n\nQUESTIONS:\nAre 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?\n\nLIMITATIONS:\nSee 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-13"}