Langevin Flows for Modeling Neural Latent Dynamics

Abstract: Neural populations exhibit latent dynamical structures that drive time-evolving spiking activities, motivating the search for models that capture both intrinsic network dynamics and external unobserved influences. In this work, we introduce LangevinFlow, a sequential Variational Auto-Encoder where the time evolution of latent variables is governed by the underdamped Langevin equation. Our approach incorporates physical priors -- such as inertia, damping, a learned potential function, and stochastic forces -- to represent both autonomous and non-autonomous processes in neural systems. Crucially, the potential function is parameterized as a network of locally coupled oscillators, biasing the model toward oscillatory and flow-like behaviors observed in biological neural populations. Our model features a recurrent encoder, a one-layer Transformer decoder, and Langevin dynamics in the latent space. Empirically, our method outperforms state-of-the-art baselines on synthetic neural populations generated by a Lorenz attractor, closely matching ground-truth firing rates. On the Neural Latents Benchmark (NLB), the model achieves superior held-out neuron likelihoods (bits per spike) and forward prediction accuracy across four challenging datasets. It also matches or surpasses alternative methods in decoding behavioral metrics such as hand velocity. Overall, this work introduces a flexible, physics-inspired, high-performing framework for modeling complex neural population dynamics and their unobserved influences.

• The paper presents a novel framework combining underdamped Langevin equations with a sequential VAE to improve neural dynamics prediction.
• The paper leverages oscillatory potential functions and a Transformer decoder to capture both local and long-range neural interactions.
• The paper achieves state-of-the-art performance on synthetic and benchmark datasets, demonstrating robust kinematic predictions.

Langevin Flows for Modeling Neural Latent Dynamics

Introduction

This paper introduces a model designed to capture the latent dynamics of neural populations by employing Langevin equations, integrating both deterministic and stochastic elements. This framework caters to the demand within neuroscience for models that can account for both internal neural dynamics and external influences, such as unobserved stochastic environmental inputs. The utilization of the Langevin equation in this context enables the modeling of complex neural systems with a focus on harnessing the inherent oscillatory and flow-like dynamics of neural activities.

Methodology

Underdamped Langevin Equation

The proposed model leverages the underdamped Langevin equation, integrating inertia, damping, potential functions, and stochastic forces to describe the evolution of latent neural variables. The potential function is parameterized using a network of locally coupled oscillators, fostering oscillatory behaviors akin to those observed in neural systems. The Langevin dynamics incorporate both autonomous dynamics and stochastic environmental effects, offering a principled approach to modeling the variability in neural systems.

Figure 1: Workflow of our method: the RNN encoder takes the spike data as input at every timestep and updates the hidden states $h_t$, and the latent variables $z_t, v_t$ evolve in time according to the Langevin equation. Finally, the Transformer decoder predicts the firing rates from the entire sequence.

Sequential Variational Auto-Encoder

The model is built upon a sequential Variational Auto-Encoder (VAE) framework, combining a recurrent encoder and a Transformer decoder. The recurrent encoder focuses on capturing local temporal dependencies, while the Transformer attends to the entire latent sequence, refining the predictions of firing rates by capturing long-range interactions. The model significantly improves upon conventional methods by predicting neural dynamics with enhanced accuracy.

Latent Langevin Posterior Flow

The time evolution of the posterior in this framework involves alternating deterministic and stochastic steps. The Hamiltonian flow preserves probability mass over time, whereas the Ornstein-Uhlenbeck process introduces a stochastic transition component. This sophisticated integration of deterministic and probabilistic flows ensures refined prediction accuracy for dynamic neural data modeling.

Experiments

Synthetic Lorenz Attractor

The model demonstrates its efficacy using a synthetic dataset generated from a Lorenz attractor system, achieving higher correlation scores in predicting the firing rates compared to established baselines like AutoLFADS and NDT.

Figure 2: Trial-average firing rates (top) and the corresponding spike trains (bottom) of some neurons of the Lorenz system.

Neural Latent Benchmark

Empirical tests on the Neural Latents Benchmark (NLB) reveal the model's consistent performance across multiple datasets. The LangevinFlow achieves state-of-the-art results in terms of likelihood of held-out neurons and forward prediction accuracy, validating its robustness and generalizability across different neural data types and sampling rates.

Figure 3: Spatiotemporal waves induced by our LangevinFlow in different views on MC_Maze. Here each group denotes an independent set of convolution channels.

Kinematic Predictions

The model accurately predicts kinematic variables such as hand velocities on neuro-physiological datasets, showcasing its ability to decode behaviorally relevant signals from neural activity accurately.

Figure 4: Kinematics (hand velocities and trajectories) of the ground truth and predicted by our method on Area2_Bump.

Conclusions

The introduction of the LangevinFlow framework represents a meaningful enhancement for modeling neural latent dynamics by embedding stochastic differential elements like the Langevin equation. Integrating these elements enables the model to handle complex neural population dynamics with stronger predictive capabilities. The results motivate further exploration into physics-informed inductive biases for neural modeling, potentially leading to better dynamical systems approaches for neuroscientific applications. Future work could explore more nuanced input-dependent potential functions to account for external influences, hinting towards more refined models of neural activity.