Latent Diffusion for Neural Spiking Data

TL;DR

LDNS addresses the dual need for revealing low-dimensional neural population structure and generating realistic, behaviorally conditioned spiking data. It achieves this with a two-stage approach: a regularized autoencoder employing structured state-space (S4) layers to produce time-aligned latent trajectories $\mathbf{z} \in \mathbb{R}^{d\times T}$, and a conditional diffusion model operating in latent space to sample $\mathbf{z}^*$ with variable length and covariate conditioning. An expressive, autoregressive spike-history observation model augments the Poisson likelihood to capture single-neuron dynamics without perturbing latent dynamics, improving realism. The method is validated on synthetic Lorenz dynamics and real datasets from human cortex during attempted speech and monkey reach tasks, showing faithful reproduction of spike-count distributions, inter-spike-interval statistics, and population correlations, and enabling conditional generation on reach direction and velocity profiles. LDNS thus provides a practical, modular framework for simultaneous latent inference and high-fidelity generative modeling of neural spiking data, with potential for closed-loop in silico experiments and hypothesis testing, while acknowledging latent-dimension selection and privacy considerations for synthetic 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.

Figures (22)

• Figure 1: Latent Diffusion for Neural Spiking data. LDNS allows for (un)conditional generation of neural spiking data through combining a regularized autoencoder with diffusion models that act on the low-dimensional latent time series underlying neural population activity.
• Figure 2: Realistic generation of spiking data with underlying chaotic dynamics.a) Synthetic spiking data from an underlying Lorenz system with a Poisson observation model. b) Accurate, smooth rate predictions of the autoencoder for held-out spiking data. c) Plotted trace of sampled latents (256 bins training length, left) and $16\times$ the original training length (middle). The sampled latent distribution matches the PSD of the autoencoder latents (right; median, 10%, and 90% percentiles). d) LDNS population spike count histogram (kde: kernel density estimate) and pairwise cross-correlations match the training distribution. e) LDNS single neuron statistics, i.e., mean inter-spike interval (isi) and std isi, match the training distribution.
• Figure 3: Unconditional generation of variable-length trials of human spiking data during attempted speech.a) Multi-unit activity is recorded from speech production-related regions of the brain (top) during attempted vocalization of variable-length sentences (bottom). b) Neural activity during sentences of different lengths. c) LDNS unconditionally sampled trials with different lengths, using the Poisson observation model. d) LDNS population spike count histogram, and mean and std of the isi match those of the data. e) Correlation matrices of the data (left) and LDNS samples (middle), and scatterplot of the pairwise correlations of data vs. LDNS samples (right).
• Figure 4: Realistic generation of spiking data in a monkey performing reach tasks.a) A monkey performs diverse reach movements in different mazes. b) Neural activity during a reach trial and a sampled trial from LDNS with a Poisson observation model. c) The LDNS population spike count histogram, and pairwise correlations match those of the data. d) LDNS mean- and std isi match the monkey data distribution. e) Auto-correlation of data, LDNS samples with Poisson observations (left), and LDNS samples with spike history, grouped according to correlation strength.
• Figure 5: Generation conditioned on monkey reach directions and velocity traces.a) Closed loop assessment: do conditionally generated latents translate to neural activity consistent with the desired direction or reach movement? b) Unseen reach movements (data) and corresponding movements decoded from the rates predicted by the autoencoder (ae). c) Decoded reach directions of LDNS samples conditioned on initial reach angles $\theta$. d) Decoded reach directions of LDNS samples conditioned on 3 unseen reach movements (velocities $v_x, v_y$). e) Straight reaches from the test set used for velocity conditioning. f) LDNS sampled latents conditioned on trajectories shown in e) vary smoothly over time and reflect information about reach angles. g) PCs of sampled LDNS latents shown in f) reveal meaningful and separable information about behavior.