  Neural networks (NNs) that exploit strong inductive biases based on physical
laws and symmetries have shown remarkable success in learning the dynamics of
physical systems directly from their trajectory. However, these works focus
only on the systems that follow deterministic dynamics, for instance, Newtonian
or Hamiltonian dynamics. Here, we propose a framework, namely Brownian graph
neural networks (BROGNET), combining stochastic differential equations (SDEs)
and GNNs to learn Brownian dynamics directly from the trajectory. We
theoretically show that BROGNET conserves the linear momentum of the system,
which in turn, provides superior performance on learning dynamics as revealed
empirically. We demonstrate this approach on several systems, namely, linear
spring, linear spring with binary particle types, and non-linear spring
systems, all following Brownian dynamics at finite temperatures. We show that
BROGNET significantly outperforms proposed baselines across all the benchmarked
Brownian systems. In addition, we demonstrate zero-shot generalizability of
BROGNET to simulate unseen system sizes that are two orders of magnitude larger
and to different temperatures than those used during training. Altogether, our
study contributes to advancing the understanding of the intricate dynamics of
Brownian motion and demonstrates the effectiveness of graph neural networks in
modeling such complex systems.
