Thermo-Symplectic Neural Langevin Flows

Molecular dynamics (MD) underlies biomolecular simulation by modeling atomistic motion over physical time. In thermostatted settings, this evolution is commonly formulated as underdamped Langevin dynamics, whose sequential force-driven integration makes long horizon simulation computationally expensive. While learning based approaches use MD data to model ensembles or state transitions, their computational primitive is defined over empirical state transitions rather than the Langevin update itself. We propose a Thermo-Symplectic Neural Langevin Flow (TSL-Flow), a learnable framework that models underdamped Langevin dynamics through phase-space updates designed to preserve thermo-symplectic structure. Our approach lifts local stochastic updates into Lie algebraic summaries, allowing the learned dynamics to preserve the geometric structure of phase-space while exposing an associative composition across time. To enable accurate and efficient long horizon simulation, we use this compositional structure to define a prefix scan compatible product that propagates trajectory segments in parallel. Empirically, our method improves coordinate and physical fidelity across biomolecular benchmarks and achieves the lowest inference time. These results indicate that temporal efficiency in learned molecular simulation can be improved by representing the Langevin update as a composable stochastic physical primitive.