Langevin Flows for Modeling Neural Latent Dynamics

TL;DR

This work introduces LangevinFlow, a sequential Variational Auto-Encoder where the time evolution of latent variables is governed by the underdamped Langevin equation, a flexible, physics-inspired, high-performing framework for modeling complex neural population dynamics and their unobserved influences.

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.

Figures (4)

• Figure 1: Workflow of our method: the RNN encoder takes the spike data as input at every timestep and updates the hidden states ${\bm{h}}_t$, and the latent variables ${\bm{z}}_t,{\bm{v}}_t$ evolve in time according to the Langevin equation. Finally, the Transformer decoder predicts the firing rates from the entire sequence.
• Figure 2: Trial-average firing rates (top) and the corresponding spike trains (bottom) of some neurons of Lorenz system.
• Figure 3: Spatiotemporal waves induced by our LangevinFlow in different views on MC_Maze. Here each group denotes an independent set of convolution channels.
• Figure 4: Kinematics (hand velocities and trajectories) of the ground truth and predicted by our method on Area2_Bump.