Energy-based Autoregressive Generation for Neural Population Dynamics

TL;DR

The paper tackles the trade-off between computational efficiency and high-fidelity neural population modeling. It introduces Energy-based Autoregressive Generation (EAG), a latent-space, energy-based transformer framework that learns temporal neural dynamics via strictly proper scoring rules, enabling single-pass generation of realistic spike data. Empirical results on Lorenz simulations and Neural Latents Benchmark datasets show state-of-the-art generation quality with substantial speedups over diffusion-based methods, plus strong conditional generation capabilities that generalize to unseen behavioral contexts and improve motor BCI decoding. The work demonstrates the practicality of energy-based latent modeling for neuroscience research and neural engineering, with broad implications for efficient, high-fidelity synthetic neural data and BCI applications.

Abstract

Understanding brain function represents a fundamental goal in neuroscience, with critical implications for therapeutic interventions and neural engineering applications. Computational modeling provides a quantitative framework for accelerating this understanding, but faces a fundamental trade-off between computational efficiency and high-fidelity modeling. To address this limitation, we introduce a novel Energy-based Autoregressive Generation (EAG) framework that employs an energy-based transformer learning temporal dynamics in latent space through strictly proper scoring rules, enabling efficient generation with realistic population and single-neuron spiking statistics. Evaluation on synthetic Lorenz datasets and two Neural Latents Benchmark datasets (MC_Maze and Area2_bump) demonstrates that EAG achieves state-of-the-art generation quality with substantial computational efficiency improvements, particularly over diffusion-based methods. Beyond optimal performance, conditional generation applications show two capabilities: generalizing to unseen behavioral contexts and improving motor brain-computer interface decoding accuracy using synthetic neural data. These results demonstrate the effectiveness of energy-based modeling for neural population dynamics with applications in neuroscience research and neural engineering. Code is available at https://github.com/NinglingGe/Energy-based-Autoregressive-Generation-for-Neural-Population-Dynamics.

Figures (16)

• Figure 1: Energy-based Autoregressive Generation (EAG) framework. Known latent positions (blue) provide context for predicting masked positions (gray). The MLP generator incorporates noise $\boldsymbol{\epsilon}$ via adaptive layer normalization to enable stochastic generation, trained with energy loss for distributional prediction.
• Figure 2: Unconditional generation on Lorenz Dataset. The generated data closely matches the ground truth across four metrics: spike count distribution, pairwise correlation, mean-isi, and std-isi.
• Figure 3: The latency/quality trade-off for EAG and LDNS. We vary number of diffusion steps (200, 400, 600, 1000) of LDNS and number of autoregressive steps (16, 32) of EAG. EAG-32 achieves a 96.9% reduction in latency, and a 32.4% improvement on RMSE mean ISI compared to LDNS-1000.
• Figure 4: Generalization to unseen angle labels. (a) real firing rates and sampled firing rates. (b) Decoded trajectory from real rates and sampled rates conditioned on unseen angle labels. (c) Single-trial neural trajectories in latent space extracted from real and sampled activity. (d) The first 2 principal components averaged over eight reach directions of real and sampled firing rates. For all panels, top: real data; bottom: sampled data.
• Figure 5: Generalization to unseen velocity labels. (a) Real hand trajectory and decoded trajectory (left panel), decoded trajectory from real rates and sampled rates (right panel). (b) Sampled spike trains conditioned on two nearly identical trajectories show trial-to-trial variability.