Langevin Flows for Modeling Neural Latent Dynamics

TL;DR

LangevinFlow is a physics-inspired sequential VAE that models neural latent dynamics using Langevin equations, capturing oscillatory behaviors and outperforming existing methods on synthetic and real neural data.

Contribution

It introduces LangevinFlow, a novel neural latent dynamics model incorporating physical priors and Langevin equations, with a locally coupled oscillator potential for biological oscillations.

Findings

Outperforms state-of-the-art baselines on synthetic Lorenz attractor data.

Achieves superior likelihoods and prediction accuracy on Neural Latents Benchmark datasets.

Effectively decodes behavioral metrics such as hand velocity.

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.