BroGNet: Momentum-Conserving Graph Neural Stochastic Differential Equation for Learning Brownian Dynamics
Authors: Suresh Bishnoi
Jayadeva Jayadeva
Sayan Ranu
N M Anoop Krishnan
Venue: ICLR 2024 poster (2024)
Keywords: Brownian dynamics
stochastic differential equation
graph neural network
scientific machine learning
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
TL;DR: Here, we present a momentum-conserving graph neural network for learning brownian dynamics represented by a stochastic differential equation

Abstract: 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, such as Newtonian or Hamiltonian. 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 modify the architecture of BroGNet to enforce linear momentum conservation 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. Finally, we show that BroGNet conserves the momentum of the system resulting in superior performance and data efficiency. 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.

Published as a conference paper at ICLR 2024
BRoGNET: MOMENTUM CONSERVING GRAPH NEU- RAL STOCHASTIC DIFFERENTIAL EQUATION FOR LEARNING BROWNIAN DYNAMICS
Suresh Bishnoi, Jayadeva, Sayan Ranu, N. M. Anoop Krishnan Indian Institute of Technology Delhi, Hauz Khas, New Delhi, India 110016 (srz208500, jayadeva, sayeevaaa,yyyya sayanranu, trishhan] @iitt.ddc..... ac. in
ABSTRACT
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, such as Newtonian or Hamiltonian. Here, propose framework, namely Brownian graph neural networks (BroGNEt), combining stochastic differential equations (SDEs) and GNNS learn Brownian dynamics directly from the trajectory. We modify the architecture of BROGNET to enforce linear momentum conservation of the system, which, in turn, provides superior performance on learning dynamics 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 simulate unseen system sizes that are two orders of magnitude larger and to different temperatures than those used during training. Finally, show that BROGNET conserves the momentum the system resulting in superior performance and data efficiency. Altogether, our study contributes to advancing the understanding of the intricate dynamics of Brownian motion and demonstrates the effectiveness graph neural networks in modeling such complex systems.
INTRODUCTION AND RELATED WORKS
Learning the dynamics of physical systems directly from their trajectory is an active area of research due to their potential applications in materials modeling Park al. (2021), drug discovery Vamath- evaa et al.] (2019) ,, motion planning Ni & Qureshi (2022), robotics Sanchez- et al. (2019); Greydanus et al. 2019), and even astrophysics Sanchez- Gonzalez al. (2020). Recent works demon- strated that physics-bat laws such as energy and momentum while also simplifying the learning making the model data efficient Thangamuthu et al. (2022); Finzi et al. (2020). Among these, family of models, such as Lagrangian or Hamiltonian neural networks Bhattoo et al. (2023) 2022); et al. (2021) [Sanchez] Gonzalez al. (019);; Greydanus al. (2019) and Neural ODEs Bishnoi al. (2022) Zhong et al. (201);; et al. (202)); Chen et al. (2018), enforces the physics-based inductive biases a strong sense. Here. a governing ordinary differential equation (ODE) is used along with neural network to learn the abstract quantities, such as energy or force, directly from the trajectory of the system. These models have shown remarkable success in learning the dynamics of variety of systems in an efficient fashion, viz., particle-based systems Bishnoi et al. (2022); TThngammuthu et al. (002)).
inductive biases could enable the learned models to follow conservation
atomistic dynamics Park al. 2021); |Huan al. (2022), physical systems Bhattoo et al. 2023
Sanchez-Gonzalez et al. (2019) Greydanus et al. (2019), and articulated systems Bhattoo et al. (2022)
Despite their success, these works focus on purely deterministic systems where the dynamics governed by an ODE Kariiadaiis et al. (202T): Thangamuthu et al. (202)); Zhong et al. (202T). An
alternative approach to model physical systems is to formulate the governing equation as stochastic
'SB: School of Interdisciplinary Research, J: Department of Electrical Engincering, SR: Department of Computer Science, NMAK: Department of Civil Engineering, SR and NMAK: Yardi School of Artificial Intelligence (joint appointment).

---

Published as a conference paper at ICLR 2024
differential equation (SDE), for instance, Brownian dynamics (BD) widely used study the dynamics of particles in a fluid or solvent Van Gunsteren & Berendsen (182)). Brownian dynamics simulations have proven to be highly valuable in numerous scientific disciplines. In physics, they have been used to study the diffusion of particles Chung et al. (1999), the self-assembly of colloidal systems Noguchi systems/Noguchi & Takasu (200T), and the behavior of polymers in solution Helfand et al. (T980). In
chemistry, Brownian dynamics simulations have shed light on reaction kinetics Northrup & Erickson (1992), molecular diffusion Madura et al. (1995), and the behavior of macromolecules Cates ( (1985).
In biology, they have provided insights into the movement of cells Klank et al. (2018). the folding of proteins Karplus Weaver (1994), and the dynamies of biomolecular complexes Gabdoulline & Wade (2001). Despite their importance in such wide domains, while there have been approaches to inferring the nature of Brownian dynamics from observations in a statistical sense Gnesotto et al. (2020): Genkin et al. (202T), to the best of the authors' knowledge, no attempt has been made to learn the dynamics of Brownian systems in particular, or SDEs in general, from their trajectory.
An alternate approach involves learning the dynamics directly using the graph representation of the data Poli et al. (202I): Chamberlain et al. (202T); Eliasof et al. (2021) Gravina et al. (2023); Rusch et al. (2022)) Wang al. (2019); Eliasof et al. (2023); Wu al. (2023)- Here, a physical system is modeled using a graph neural network (GNN) where the nodes represent the particles and the edges represent the interactions. To these models, additional inductive bias in terms of physics-informed loss functions ("soft constraints") or architectures ("hard constraints") can be employed. Diverse approaches integrating soft and hard constraints have advanced the understanding and learning of dynamical systems. Previous works explored the utilization of coupled ODEs for real-time diagnosis of COVID-19 spread Dandekar et (2020) capturing system invariants Zhang et al. (022)), and incorporating conservation laws Hansen et al. 22044) to name a few. Furthermore, physics-informed informed inductive biases to enhance the performance of GNNS has also been studied Thangamuthu et al. (2022). Architectural modifications for hard constraints include graph neural ODEs, and Lagrangian and Hamiltonian graph neural networks Poli et al. (202T);
the significance of physics
Gondzlez- et al. (2019); Bhattoo al. (2023) 2022). HHaber & Ruthotto (2017) presents an architecture that mitigates gradient instability with ODEs and Hamiltonian- inspiieee forward propagation. E et al. (2017) successfully models complex dynamics through connection with backward stochastic differential equations. However, most of these works focus on deterministic systems, and there are limited efforts to extend them to stochastic systems Eliasof al. (2023).
Here, we propose BROGNET, which can learn the Brownian dynamics of system directly from the trajectory. Exploiting GNN-based framework, we show that the Brownian dynamics can be learned from small systems, which can then be generalized to arbitrarily large systems. The major contributions of the work are as follows.
Learning stochastic dynamics from trajectory: We present a framework, BROGNET, that learns the abstract quantities governing the dynamics of stochastic systems directly from their trajectory. While demonstrated for Brownian dynamics, the framework presented is generic and may be applied to other SDEs as well.
Graph-based modeling for zero-shot generalizability: We model the physical system as graph with nodes representing particles and edges representing their interactions, which can then be used to learn BD. The inductive graph-based modeling naturally imparts the desirable properties of zero-shot generalizability to unseen system sizes and permutation invariance. In addition, we also demonstrate generalizability to unseen temperatures in a zero-shot fashion.
Momentum conservation: We demonstrate that the proposed framework follows Newton's 's third law, thereby conserving the total linear momentum of the system. This feature, in turn, provides superior performance for the model.
2 BACKGROUND ON BROWNIAN DYNAMICS
The fundamental concept behind Brownian dynamics is the influence of thermal fluctuations on the motion of particles. In a fluid, particles are subject to random collisions with the surrounding molecules, which lead to erratic motion characterized by continuous changes in direction and speed (see the videos in the supplementary for a visual depiction). The behavior of particles undergoing Brownian motion can be described mathematically using SDEs, such over-damped Langevin equations Coffey & Kalmykov| (2020). These equations incorporate deterministic forces, such as external potentials or interactions between particles, as well as stochastic forces that represent the
2

---

Published as a conference paper at ICLR 2024
NLP
IISA rodd
rxpresentaion laseen
node Nuchar
MLF
Iraid orbeekdng
GNN
Model lupellat sware irternal frow
Figure BROGNET's architecture. MLP refers to multilayer perception, GNN to graph neural networks and ohe to one-hot encoding.
random effects of the surrounding fluid. By numerically integrating these equations, it is possible to simulate and analyze the dynamic behavior of the particles over time.
Consider system of N particles with masses M and interacting with each other through interaction potential U (X+), where X, represents the position vectors of all the particles at time I, constituting time-dependent random variable. The Langevin equation governing the dynamics of this system given by Coffey & Kalmykov (2020)
MX.- VU.(X) CMX,+ + V2MçkBT Qt
where (X. F(X,) is the force due to the inter-particle interactions, ( is the damping constant, T is the temperature, k is the Boltzmann constant, and Q is delta-correlated stationary Gaussian process with zero mean that follows (Πtz) 0 and (52, ô(t - t'), where ó is the standard Dirac-delta function. the over-damped limit Langevin dynamics, where no acceleration exists, the equation reduces to Brownian dynamics as Brańka & Heyes (1998)
dXx VU(X)) V27kgT Ωr
drift term
diffusion term
where 7 = .MC. Thus, the evolution of system can be obtained by numerically integrating the SDE given by Eq. The first term in the RHS of Eq. i the deterministic inter- soer-patticle interaction (drift term), and the stochastic nature in the dynamics comes from the second term (diffusion).
BROGNET: A GRAPH NEURAL SDE APPROACH
To learn the trajectories of multi-particle system governed by Brownian dynamics, we design a neural network called BROGNET. BROGNET is essentially Graph Neural SDE ire-inspired by Graph Neural ODE, which parametrizes an ODE using GNN. Thus, BROGNET parameterizes the stochastic dynamics ((x) using a neural network and learns the approximate function B(x,) by minimizing the loss between the distribution of ground truth trajectories with that of the predicted trajectories. However, in contrast to classical dynamics, the loss is applied over a distribution of positions rather than the exact deterministic trajectory due to the stochastic nature of the dynamics. We discuss the exact loss function in detail later.
The learning process is, therefore, decomposed into two factors: the standard deviation associated with the ground-truth stochastic process corresponding to each particle and their expected positions. In BROGNET, the standard deviation is learned through MLP. On the other hand, the expected positions are learned through a Graph Neural SDE. The choice of using a GNN to model interaction dynamics is motivated by two key observations. First, GNN allows number of model parameters is independent of the number of particles in the system. Hence, enables zero-shot generalizability systems of arbitrary unseen sizes. addition, a GNN allows modeling of the system. Below, we discuss the details.
be inductive, where the
permutation-invann
3.1 THE GNN ARCHITECTURE
BROGNET models the interaction dynamics among participating particles graph. The graph topology is used learn the approximate dynamics B by minimizing the distance between distribu- tions representing the true and predicted positions. Note that henceforth all the variables with a hat O correspond to the approximate variable leamned by the machine learning model.
Graph structure. We represent an N-particle system as an directed graph 9 = {V,E). V represents the set of nodes, where each node represents particle. Similarly, & represents interactions
3

---

Published as a conference paper at ICLR 2024
connections among particles in the form of edges. For example, in ball-spring system or colloidal gels, the balls or colloid particles correspond to nodes, and the springs or the interactions between the colloidal particles correspond to edges. Note that in the case of systems such as springs, where explicit connections are pre-defined, the graph structure remains static. However, in the case of systems like colloidal gels, the edges are defined based a neighborhood cut-off, and hence the graph structure can be dynamic, with the edges between particles changing as a function of the particle configuration. As will see in the next paragraph, directed graph is required only because the edge weights we use are directional.
Input features. Each node corresponds to a particle and is defined by its intrinsic features. such as the particle's type (7z). position (X,,)), and velocity (X,,)) at time t. The position and velocities are 3-dimensional tuples corresponding to the x, and dimensions. The particle type distinguishes between different characteristics, such as varying masses or friction coefficient for balls or particles. For brevity, X is referred to as X; henceforth unless specified otherwise. On the other hand, each edge represents the connection between nodes and is represented by the 3-dimensional edge feature Wij (X; Xj), which signifies the relative displacement of the connected nodes along each coordinate. Note that the edge weights are asymmetric due to the vectorial nature of the input. Summary of neural architecture. Fig. [ provides pictorial description of BROGNET. In the pre pre-processing layer, we utilize Multilayer Perceptrons (MLPs) to create dense vector representations for each node v; and edge eij. In systems where internal forces play significant role in governing dynamics, the structure's 's topology holds crucial importance. To capture this dependency, we employ deep, phessage-passing Graph Neural Network (GNN). The final representation z, h and Ziy h/j are generated by the GNN for each node and edge respectively, capturing relevant information. These representations are subsequently passed through another MLP to predict the force of each node. We next detail these individual steps.
Pre-processing. In the pre-processing layer, a dense vector representation is constructed for each node iú and edge eig using MLP denoted as MLP em The construction process can be described as
h squareplus (MLP (MLP (one-hor hot (T:)))
h squareplus (MLP (Wiz))
(4)
Here, squareplus represents an activation function. It is important to note that the node and edge embedding functions, parameterized by MLP cm, ML, utilize separate weights. For the sake of brevity, we refer to them simply as MLP cm- Note that the node representations do not explicitly encode their ground-truth positions. Only the distances between nodes are encoded in the edge representations. This allows to ensure translational invariance.
Force prediction. We utilize multiple layers of dessaae-passing to facilitate communication between nodes and edges. In the eth layer, the node embeddings updated follows:
h squareplus Wh.
htll
Σ JEN!
Σ E
hit
h
JEN
Here, Min {uz E | ega E E) and Nout (u, EVI eij E E) represents the incoming and outgoing neighbors of node i; respectively. || denotes the vector concatenation operation. The term W signifies layer-specific learnable weight matrix. More simply, we sum-pool the edge embeddings over both incoming and outgoing neighbors, concatenate them along with the target node's 's own embedding and then pass it through one linear transformation followed by non-linearity. The embedding of an edge eij on node o; in the 8th layer, denoted as hfz- is computed as:
h e+1 squareplus (We (hā][bf[[in)
(6)
g Similar to W W is layer-specific learnable weight matrix that applies to the edge set. The In the th layer, the final
cressage-passingg process spans L layers. with L being a node and edge representations are denoted as n, h and zij h respectively.
Finally, the pair-wise interaction force F.j from particle i to j is predicted as:
= squareplus(MLPy/ee
(7)
where, MLPy denotes a Multilayer Perceptron with a squareplus activation function Barron (202T). One may also use ReL.U as activation function, see empirical comparison at Fig. App. H
4

---

Published as a conference paper at ICLR 2024
The force node is defined the summation of all forces from incoming edges and the reactive force from its outgoing edges. Mathematically,
F =
Σ JEX
Σ
Pij
(8)
JENT
Here, the second term accounts for the reactive force from each outgoing edge, which is essentially the negated version of the force particle i imparts particle j. This inductive bias ensures Newton's ' third law that every action has an equal and opposite reaction. More importantly, this inductive bias implicitly ensures that the net force the system is zero, thereby strictly ensuring momentum conservation. Formally:
Theorem (Momentum Conservation) In the absence of an external field, BROGNET exactly conserves the linear momentum of the system.
Proof of the theorem is provided in Appendix A Further, as demonstrated later the empirical results comparing the performance of BROGNET with a version that does not employ the momentum conservation, namely, Brownian dynamics GNN (BDGNN), suggests that the momentum conservation indeed results in better performance of the model.
Learning the Brownian term. While we can learn the distribution associated with the diffusion term (see Eq. ) it is well-known that Brownian motion follows multivariate Gaussian distribution. The diffusion term independent of the system topology and depends only the particle type and its attributes such radius or friction coefficient. To learn this term, represent the node type one-hot encoding, which is passed through an MLP to obtain the 5% associated with each particle (see Fig. I]. This 94. along with force obtained from the GNN, is substituted in the governing equation (Eq. 2 and learned in an end-to-end fashion. It is worth noting that the standard deviation of the stochastic process is given by o, k2 Thus, all other terms other than % in the Eq. 2 are
known apriori or are constants. Trajectory prediction and training. Based the predicted forces and yxs the positions derived using the Euler Maruyama integrator Brańka & Heyes [1998). The loss on the prediction is computed using Gaussian Negative Log- Log-likelihood loss, which is then back- proppparpe Specifically, let X,, and X, denote the ground truth and predicted positions of particle at time f, respectively. The loss function is defined as follows.
train the model.
C 1
ΣΣ log max (o, (X.. X.,)² max(6?,e)
Here, c is a small constant, and À controls the relative weightage of the two terms. à; is the standard deviation associated with the position of each of the particles. The overall objective of the loss is to minimize the pariance-normalized deviation from the ground truth (second term), while also ensuring regularizing term to keep the variance low. We use Gaussian NLL loss since the stochasticity in Brownian dynamics is a Gaussian as well (recall Eq.
Note that the positions, both predicted and ground-truth, are computed directly from the predicted and ground-truth forces, respectively, using Euler Maruyama integrator as follows.
2kgT
X,(t + +t) - = x,()) + Fi/MAt+
AN
(10)
%
where AΩ, is a random number sampled from a standard Normal distribution. It should be noted that the training approach presented here may lead to learning the dynamics dictated by the Euler Maruyama integrator. Thus, the learned dynamics may not represent the "true" dynamics of the system, but one that is optimal for the Euler Maruyama integrato
EXPERIMENTS
In this section, we benchmark the ability of BROGNET to learn Brownian dynamics directly from the trajectory and establish:
Accuracy: BROGNET accurately models Brownian dynamics and outperforms baseline models.
See App. B for details Fuler Marayama integrator and other integrators for SDE.
5

---

Published as a conference paper at ICLR 2024
Zero-shot generalizability: The inductive architecture of BROGNET allows it accurately generalize to systems of unseen sizes and temperatures.
The codebase and all datasets used in this work can be accessed from the repository at https: /ggihub.com/MBBR Coom/M3RG-IITD/Bree
4.1 EXPERIMENTAL SETUP
We provide the details of the hardware and software environment in App. C.1
To the best of the authors' knowledge, there no prior works on learning Brownian dynamics directly from the trajectory. Thus, in order to compare the performance BROGNET, we propose five different baselines.
Baselines:
(I) Neural Network (NN): Here, we use feedforward MLP that takes X; vector at time i as the input and predicts the positions ++ A/ the output.
(2) Brownian NN (BNN): NN in (1) does not have any physics based inductive biases. To address this, we propose BNN, a Neural SDE that parameterizes the Brownian dynamies using feed-forward MLP. Specificcl, BNN takes X, of all the particles concatenated into vector as input and predicts the F and 9 as output employing fully connected MLP. These outputs are then used in Eq. 10 to predict the updated positions. Thus, BNN employs the inductive-bias in terms the Brownian equation of motion.
(3) Brownian full graph network (BFGN): The models in (1) and (2) do not exploit the topology of the system and are not permutation invariant. To this extent, we employ BFGN graph-based baseline that takes the position as input and outputs the forces and ry for each node. This output is then sub- stituted in Eq. 10 to obtain the updated position. Thus, BFGN is graph-based version of BNN. Note that BFGN employs the full graph network architecture Cranmer et al. (2020); Sanchez-Gonzale et al. 2220)) a versatile architecture that has been demonstrated for variety of particle-based systems. It worth noting that BFGN is not translational invariant due to the explicit use of positions as node inputs. (4) Brownian Dynamics GNN (BDGNN): We also present physics-informed baseline, BDGNN. BDGNN has all the inductive biases of BROGNET except Newton's 's third law. Specifically, the force on each particle computed as:
F. = (MLPV
(11)
Note that BDGNN has the same architecture as BROGNET and hence allows us to evaluate the role of momentum conservation as inductive bias in the BROGNET.
(5) Brownian Neural Equivariant Interatomic Potentials (BNequlP): BNequIP employs an equiv- ariant GNNS proposed by (Batzner et al. (022) and directly predicts the force vector as the output. (6) Momentum-informec BDGNN (MIB In order to study the effect of adding momentum conservation as "soft-constraint,, we add an additional term to the loss function of BDGNN for momentum conservation. Specifically, we add third term to Eq. 9 namely, IIΣ Fill = O.
Datasets and systems: To compare the performance BROGNET, we selected three different systems, namely, linear spring. linear spring with binary particle types, and non-linear spring systems, all following Brownian dynamics at finite temperatures. The first is a model system with particles interacting with each other based a linear spring and subjected to Gaussian noise. The second system comprises a non-linear spring where the force on a spring is proportional to the fourth power of displacement. The third system corresponds of particles having different 7. For linear and non- non-linear spring systems, models are trained on 5-spring system only, which are then evaluated on other system sizes. In case of linear spring with binary particle types, models are trained on 10-spring system and evaluated on other system sizes. The details of all the systems to generate the ground truth data are provided in App. C Further, the detailed data-generation procedure is given in App. D
linear spring system with two different types
To quantify performance, we use the following metrics. The position error is the normalized Euclidean distance between the ground truth and predicted positions. Mathematically, let (ar, y, and ă. ÿ, å be the ground truth and predicted positions, respectively, of a given particle. Furthermore, let 0z/u/= be the ground-truth standard deviation the positions of the particle. The error is, therefore:
Evaluation Metric:
Position error:
1/2
2
PE =
Σ
(12)
Tc
OE
(E()))))
Brownian error: The Brownian error quantifies the difference between the standard deviation

---

Published as a conference paper at ICLR 2024
estimated by the neural network and that of the ground truth using the RMSE metric. Trajectory roll-out error: The above two metrics individually look at the errors in positions and the standard deviation. In this metric, we holistically compare the distance between the distribution of trajectories in the ground truth against the ones predicted by BROGNET. Specifically, given a Brownian system, we simulate 100 ground-truth trajectories across 100 time steps over each particle in the system. Similarly, we predict the trajectories using BROGNET. We next compute the average distance between the ground-truth distribution and the predicted one using KL-divergence across all particles. KLL-divergence is defined as follows:
(x) X(x)
DKL (X 12)
Σ aEx
xxx log
(13)
Here, X' and X' represent the distributions over the ground truth and predicted trajectories, respec- tively. Both these distributions are univariate normal distributions since the stochastic component in Brownian dynamics is normally distributed (Eq. 2 Hence, Eq. reducee the following:
σ1 σ + (Mo 21) 1
DEL (x = log +
(14)
du
209
2
Here, do and (40 are the standard-deviation and mean of distribution & and similarly on and μr denote the standard deviation and mean of X.
Model architecture and training setup: We use 10000 data points generated from 100 trajectories with random initial conditions to train all the models. Detailed training procedures and hyper- parameters are provided App. E For all versions of the graph architectures, MI.s are two layers deep. All models were trained till the decrease in loss saturates to less than 0.001 over 100 epochs. The model performance is evaluated on a forward trajectory, a task it was not explicitly trained for, of 0.1s. To compute the evaluation metrics, all the results are obtained over trajectories generated from 100 different initial conditions. Further, to evaluate the zero-shot generalizability, the models trained on 5 spring systems evaluated on 50, and 500 spring systems. Further, the models trained at 1 unit temperature are evaluated at 10 and 100 unit temperatures.
4.2 COMPARISON WITH BASELINES
10
Fig. 2 shows the trajectory error, Brownian error, and the position er- ror and the geometric mean of these errors over the trajectory for all the baselines, namely, NN, BNN, BFGN, BDGNN, MINNGNN, and BNequlP in comparison to BROGNET for lin- ear 5 5-spring system. We observe that BROGNET and MIBDGNN sig- nificantly outperform all other base- lines in terms of overall evaluation metrics. Specifically, we observe that the BROGNET outperforms BDGNN, which differs from the BROGNET only in terms of momentum con- servation. This suggests that the momentum conservation bias added BROGNET sigmificantly enhances the performance of the model. Interest- ingly, we observe that the MIBDGNN significantly outperforms BROGNET. This shows that the additional regular- izer in the loss function enables im- proved learning of the dynamics.
IPON BOCIa
10"
RECN
Shera Ee00ie
0.00
0.50
0.06
(d)
(e)
(
10
2 10
10"
10
6x
10
NN
BW
BFGN
BOONY
MIDGWN
BaN BFSN
BDGMN MEDSH
BroGH
I0 co
Figure 2: Comparison of BROGNET on a linear 5 5-spring system against baselines, NN, BNN, BFGN, BDGNN, MIBDGNN and BNequIP using (a) trajectory roll-out error, (b) Brownian error, and (c) position error. All models are trained using 10000 data points. The plots are obtained based on 1000 forward simulations by the trained model from 100 initial conditions, each evaluated with 10 random seeds random. Additionally, the geometric mean of (d) trajectory roll-out error, (e) Brownian error, and (f) position error is shown in the bottom subplot.
Figure shows the performance of BROGNET on non-linear 5 5-spring system, and Figure shows the performance on binary linear 10- 10-spring systems, respectively. As in the case of the linear spring system, we observe that BROGNET outperforms all other baselines in learning the Brownian dynamics for non-linear spring systems and binary linear spring systems. However, we note that
7

---

Published as a conference paper at ICLR 2024
the difference between the performance of BROGNET with other baselines is lower in binary spring systems. Nevertheless, these results suggest the ability of BROGNET to learn the dynamics of complex
systems interacting with each other based non-linear interaction forces.
4.3 ZERO-SHOT
[b]
GENERALIZABIIITY
DHM BFon
2004
Now, evaluate the zero-shot gen- eralizability of the learned models to unseen system sizes and temperatures. To this extent, we use the learned mod- els. BFGN, BDGNN, MIBDGNN, BNequlP and BROGNET trained on N =5 and T= units. First, these mod- els are evaluated on system sizes of 50 and 500 springs, up to two orders of magnitude larger than those used in the training data. Figure 5 (a-c) shows the trajectory rollout error of BFGN, BDGNN, MIBDGNN, BNequlP and BROGNET on 5 50-, and 500- spring systems, respectively. Inter- estingly, we observe all four models exhibit inductivity to larger system sizes, thanks to the GNN. Further, we observe that the trajectory roll- out error for all the models is com- parable with the increasing system size, even when evaluated on a sys- tem size of 500. All the error met-
RE.04
Bridoor
(d)
(6)
10'
B0 1806
10"
10
10
MN
BNN
SPON
BDGAN
MEDGON
marii
SFON 33
SDGIN
MEDOON
L CN
Figure 3: Comparison of BROGNET on non-linear
O-spring system against baselines NN, BNN, BFGN, BDGNN, MIBDGNN and BNequIP using (a) trajectory roll- out error, (b) Brownian error, and (c) position error. All models are trained using 10000 data points. The plots are obtained based on 1000 forward simulations by the trained model from 100 initial conditions, each evaluated with 10 random seeds. Additionally, the geometric mean of (d) tra- jectory roll-out error, (e) Brownian error, and (f) position error are shown in the bottom subplot.
rics on all three types of systems linear spring, binary linear spring, and non-linear spring included in the App. [F] These results suggest that the BROGNET datasets and generalized to arbitrarily large system sizes for performing inference on the Brow- nian dynamics of the system. It should be noted that both NN and BNN, due to their architec- ture (feed-forward MLP), are not inductive and do not allow inference on larger system sizes.
be trained small
Another important aspect of Brown- ian dynamics is the ability to simu- late the learned systems correspond- ing to different thermodynamic condi- tions, namely, different temperatures. Figure e). shows the trajectory roll- out error linear spring systems with five particles at temperatures 10 and 100 units, respectively. We again observe that BROGNET and MIB- DGNN outperforms BFGN, BDGNN, and BNequlP with very low errors on the roll-out. This suggests that both the stochastic Brownian terms and the deterministic force term learned by the BROGNET are accurate and scal- able to large system sizes and different thermodynamic conditions.
[e)
[b)
[e]
ENN Macica
10"
50CNN
SDONN
Vaccii
10
0.00
0:25
0.10
0.0
6156
0.10
10
0.00
3025
0.50
10)
4
3+ 10
30
2x 10
M
BNN
BPON
6DSNN
MEDONN
SPON
BDGAN MEDONN
Blegg BroDiie
6x 10
oN
Figure 4: Comparison of BROGNET
binary linear
10 (o-ssring system against baselines using (a) trajectory roll- out error, (b) Brownian error, and (c) position error. All models are trained using 10000 data points. The plots are obtained based on 1000 forward simulations by the trained model from 100 initial conditions, each evaluated with 10 random seeds random. Additionally, the geometric mean of (d) trajectory roll-out error, (c) Brownian error, and (f) position error are shown in the bottom subplot.
4.4 MOMENTUM CONSERVATION
Note that the Brownian motion has net- zero acceleration, so the momen- tum remains conserved. To evaluate this, we compute the total force ex- perienced by each particle during the
8

---

Published as a conference paper at ICLR 2024
N= 5
N 50
N 500
N=5, T=
N=5, T= 100
Dz (2(00
10
3.00
0.06
0.50
0.00
0.05
1110
0.30
0.06
0.10
0.00
0.0b
0.10
0.00
0.06
0.10
Time
BFGN
BOGNN
MIBDGNN
BNoquIP
BroGNet
Figure 5: Trajectory roll-out error of BFGN, BDGNN, MIBDGNN, BNequlP and BROGNET trained on N 5 and T =1 unseen system sizes = 50, and 500) and unseen temperatures 10 and 100) on the linear spring system. The plots are obtained based on 100 forward simulations by the trained model with random initial conditions.
dynamics predicted by the models. Figure 6 shows the geometric mean the total force expe- rienced by the system. Interestingly, we see that all the models except BROGNET predict finite force, while BROGNET predicts force close to zero. Although MIBDGNN exhibits lower er- ror than other models, BROGNET outperforms MIBDGNN in this metric; as expected from hard constraint vs. soft constraint. This shows that the additional inductive bias of momentum conservation enforced in the architecture of BROGNET contributes to its improved performance.
4.5 DATA EFFICIENCY
We also evaluated the data efficiency of BROGNET compared to the basclines. Specifically, we evaluated the performance of BROGNET trained with dataset size varying 100, 500, 1000, 5000, and 10000 and compared with the base- lines BFGN, BDGNN, MIBDGNN and BNequIP, trained with the respective dataset. The results shown in Fig. Suggest that BROGNET can Iearn efficiently from small dataset size compared to the baselines, which could be attributed to the momentum conservation bias in the system.
10"
11:00
0.10
010
0.101
0.00
0700
650M
Figure 6: Momentum Error of linear, non-linear, and binary-lineat spring based on 100 forward simulations by the trained model with 10 random initial conditions.
CONCLUSION
We demonstrate a graph-based neural SDE framework that can learn Brownian dynamics directly from the data. We demonstrate the approach to linear and non-linear springs and binary linear spring systems subjected to a stochastic dynamic term and following Brownian dynamics. In order to compare the performance of BROGNET, we proposed several baselines, namely, NN, BNN, BFGN, BDGNN, MIBDGNN, and BNequIP. Note that these baselines were proposed due the paucity of baselines in the literature. We show that the BROGNET and MIBDGNY architectures, with their unique momentum conservation, significantly outperforms all other baselines. Specifically, we show that MIBDGNN exhibits superior performance in learning the dynamics due to the additional term in the loss function, while BROGNET exhibits superior momentum conservation due the hard architectural constraint. We also demonstrate the zero-shot generalizability of the learned model to unseen system sizes and unseen temperatures. These results suggest that BROGNET presents robust framework to learn Brownian dynamics from small amounts of data.
14.10
3-10
12'
200W
ME00W
Figure 7: Geometric mean of (a) trajectory error, (b) Brownian error, and (c) position eгror of BFGN, BDGNN, MIBDGNN, BNEQUIP, and BROGNET trained on varying dataset size. All the results are eval- uated 1000 forward simulations; 100 initial condi- tions, each evaluated with 10 random seeds.
Limitations and future works: This work focuses on Brownian dynamics, which an over-damped limit of the Langevin equation. How does the model perform in the case of a Langevin equation where the acceleration is non-zero? How can the model be extended to other SDEs? We aim to explore these questions in future work. Moreover, the present architecture cannot learn the dynamics in the presence of external fields. Extending the model to the dynamics of real-world systems, including experimental datasets, also an outstanding problem.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
NMAK and SR acknowledge the financial support for this research provided by Intel and Google Research Scholar Award. The authors thank the IIT Delhi HPC facility for providing computational and storage resources. SB acknowledges the financial support from the Prime Minister's 's Research Fellowship (PMRF), Ministry of Education, Government of India.
REFERENCES
Jonathan T. Barron. Squareplus: A softplus- like algebraic rectifier, 2021.
Simon Batzner, Albert Musaelian, Lixin Sun, Mario Geiger, Jonathan P Mailoa, Mordechai Korblluth, Nicola Molinari, Tess E Smidt, and Boris Kozinsky. E (3)-equivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature communications, 13(1):2453, 2022.
Ravinder Bhattoo, Sayan Ranu, and NM Krishnan. Learning articulated rigid body dynamics with lagrangian graph neural network. Advances in Neural Information Processing Systems, 35: 29789-29800, 2022.
Ravinder Bhattoo, Sayan Ranu, and NM Anoop Krishnan. Learning the dynamics of particle-based systems with lagrangian graph neural networks. Machine Learning: Science and Technology, 2023.
Suresh Bishnoi, Ravinder Bhattoo, Sayan Ranu, and NM Krishnan. Enhancing the inductive biases of graph neural ode for modeling dynamical systems. International Conference Learning Representations, 2022.
AC Brańka and DM Heyes. Algorithms for brownian dynamics simulation. Physical Review E, 58 (2):2611, 1998.
ME Cates. Brownian dynamics of self-similar macromolecules. Journal de Physique, 46(7):1059 1077, 1985.
Benjamin Paul Chamberlain, James Rowbottom, Maria I. Gorinova, Stefan D Webb, Emanuele Rossi, and Michael M. Bronstein. GRAND: Graph neural diffusion. In The Symbiosis of Deep Learning and Differential Equations, 2021. URL httpss/////pee https:/pppaneiieeeee//oonmmmmm Ifu__Ifu__jjaRR
Ricky TQ Chen, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. Neural ordinary differential equations. In Proceedings of the 32nd International Conference on Neural Information Processing Systems, pp. 6572 -223,, 2018.
Shin-Ho Chung, Toby W Allen, Matthew Hoyles, and Serdar Kuyucak. Permeation of ions across the potassium channel: Brownian dynamics studies. Biophysical Journal, 77(5):2117-2533, 1999.
W. T. Coffey and Yu P. Kalmykov. The Langevin Equation: With Applications to Stochastic Problems in Physics, Chemistry and Electrical Engineering (Third edition), volume 27 of World Scientific Series Contemporary Chemical Physics. World Scientific, Trinity College, Dublin, Ireland and Université de Perpignan, France, 2020.
Miles Cranmer, Alvaro Sanchez Gonzalez, Peter Battaglia, Rui Xu, Kyle Cranmer, David Spergel, and Shirley Ho. Discovering symbolic models from deep learning with inductive biases. Advances in Neural Information Processing Systems, 33, 2020.
Raj Dandekar, Chris Rackauckas, and George Barbastathis. A machine learning- global diagnostic and comparative tool to assess effect of quarantine control in covid-19 19 spread. Patterns, 109): 2020. ISSN 2666-3899, doi: https://doi http:////oooor///// 1016// 10166/ patter.2020. 100145. URL
Weinan E, Jiequn Han, and Arnulf Jentzen. Deep learning based numerical methods for high- dimensional parabolic partial differential equations and backward stochastic differential equations. Communications in Mathematics and Statistics, 5(4): -8000 November 2017. ISSN 2194-6711 IX. doi: 10. 10.1007/440304-01--01111
10

---

Published as a conference paper at ICLR 2024
Moshe Eliasof, Eldad Haber, and Eran Treister. PDE-GCN: Novel architectures for graph neural networks motivated by partial differential equations. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https://ppenneview.neeffee net/forum?iddw WikkGixJB2x
Moshe Eliasof, Eldad Haber, and Eran Treister. Adr-gnn: Advection-diffusio--rrccii graph neural networks, 2023.
Marc Finzi, Ke Alexander Wang, and Andrew G Wilson. Simplifying hamiltonian and lagrangian neural networks via explicit constraints. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Bal- can, and H. Lin (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 13880-13888 Curran Associates, Inc., 2020. URL http://pooeddnnn...uuuii.cc/pppp Tii55c688844d4aaaaaaaahhhchCCCC
Razif R Gabdoulline and Rebecca C Wade. Protein-protee association: investigation of factors influencing association rates by brownian dynamics simulations. Journal of molecular biology, (0665))11339113)) 2001.
Mikhail Genkin, Owen Hughes, and Tatiana A Engel. Learning non-stationary langevin dynamics from stochastic observations of latent trajectories. Nature Communications, 12(1):5986, 2021.
Federico S Gnesotto, Grzegorz Gradziuk, Pierre Ronceray, and Chase P Broedersz. Learning the non-equilibrium dynamics of brownian movies. Nature communications, 11 (1):53337, 2020.
Jonathan Godwin Thomas Keck*, Peter Battaglia, Victor Bapst, Thomas Kipf, Yujia Li, Kimberly Stachenfeld, Petar Veličković, and Alvaro Gonzaez- Jraph: A library for graph neural networks in jax., 2020. URL http://gituub.comddeepmmiti
Alessio Gravina, Davide Bacciu, and Claudio Gallicchio. Anti-symmetric DGN: a stable architecture for deep graph networks. In The Eleventh International Conference on Learning Representations, 2023. URL http:///wwweeweee.eee////o/////a
Samuel Greydanus, Misko Dzamba, and Jason Yosinski. Hamiltonian neural networks. Advances in Neural Information Processing Systems, 32:15379- 15389,, 2019.
Nate Gruver, Mare Anton Finzi, Samuel Don Stanton, and Andrew Gordon Wilson. Deconstructing the inductive biases of hamiltonian neural networks. In International Conference on Learning Representations, 2021.
Eldad Haber and Lars Ruthotto. Stable architectures for deep neural networks. Inverse Problems, 34(1):014004, December 2017. ISSN 1361-6420. doi: 10. 10.1088//3311-6422/aaaaaaaa URL http:
Derek Hansen, Danielle C. Maddix, Shima Alizadeh, Gaurav Gupta, and Michael W. Mahoney, Learning physical models that respect conservation laws. Physica D: Nonlinear Phenomena, 457:133952, January 2024. ISSN 0167-2789. doi: 10.1016/j 10.1016)..hyysd.2222 fdx.doi.org/1001000 p0yd 20t..
URL http
Eugene Helfand, ZR Wasserman, and Thomas A Weber. Brownian dynamics study of polymer conformational transitions=.. Macromolecules, 13(3):526-533, 1980.
Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equiv- ariant graph mechanics networks with constraints. In International Conference on Learning Representations, 2022.
George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3(6))422-440, 2021.
Martin Karplus and David L Weaver. Protein folding dynamics: The diffussion-collisionn model and experimental data. Protein science, 3(4):650-668. 1994.
Rebecca L Klank, Steven S Rosenfeld, and David Odde. A brownian dynamics tumor progression simulator with application to glioblastoma. Convergent science physical oncology, 4(1):015001, 2018.
11

---

Published as a conference paper at ICLR 2024
Seungjun Lee, Haesang Yang, and Woojae Seong. Identifying physical law of hamiltonian systems via meta-learning. In International Conference on Learning Representations, 2021. URL https:
Jeffry D Madura, James M Briggs, Rebecca Wade, Malcolm Davis, Brock A Luty, Andrew Ilin, Jan Antosiewicz, Michael K Gilson, Babak Bagheri, L Ridgway Scott, et al. Electrostatics and diffusion of molecules in solution: simulations with the university of houston brownian dynamics program. Computer Physics Communications, 91(1-3):57-95, 1995.
Ruiqi Ni and Ahmed H Qureshi. Ntfields: Neural time fields for physics-informed robot motion planning. arXiv preprint arXiv:2210.00/200 2022.
Hiroshi Noguchi and Masako Takasu. Self-assembly of amphiphiles into vesicles: a brownian dynamics simulation. Physical Review E, 64(4).041913, 2001.
Scott H Northrup and Harold P Erickson. Kinetics of protein-protein association explained by brownian dynamics computer simulation. Proceedings of the National Academy of Sciences, 89(8): 3338-3342, 1992.
Cheol Woo Park, Mordechai Kornbluth, Jonathan Vandermause, Chris Wolverton, Boris Kozinsky, and Jonathan P Mailoa. Accurate and scalable graph neural network force field and molecular dynamics with direct force architecture. npj Computational Materials, 7(1):1-9, 2021.
Michael Poli, Stefano Massaroli, Junyoung Park, Atsushi Yamashita, Hajime Asama, and Jinkyoo Park. Graph neural ordinary differential equations, 2021.
T. Konstantin Rusch, Ben Chamberlain, James Rowbottom, Siddhartha Mishra, and Michael Bron- stein. Graph-coupled oscillator networks. In Kamalika Chaudhuri. Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning. volume 162 of Proceedings of Machine Learning Research, pp. 18888-18909. PMLR, 17-23 Jul 2022. URL httph//spppocceediggg.ooo www.s..ss....eie.e. wwt.treeee///vssessesss 122/r//c222aahama
Alvaro Sanchez-Gonzalee Victor Bapst, Kyle Cranmer, and Peter Battaglia. Hamiltonian graph networks with ode integratons. arXiv e-prints, pp. arXiv-1909, 2019.
Alvaro Banchez-Gonzaler Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovee, and Peter Battaglia. Learning to simulate complex physics with graph networks. In International Conference on Machine Learning, pp. 8459-8468. PMLR, 2020.
Samuel Schoenholz and Ekin Dogus Cubuk. Jax md: a framework for differentiable physics, Advances in Neural Information Processing Systems, 33, 2020.
Abishek Thangamuthu, Gunjan Kumar, Suresh Bishnoi, Ravinder Bhattoo, NM Krishnan, and Sayan Ranu. Unravelling the performance of physics-informed graph neural networks for dynamical systems. Advances in Neural Information Processing Systems, 2022.
Jessica Vamathevan, Dominic Clark, Paul Czodrowski, Ian Dunham, Edgardo Ferran, George Lee, Bin Li, Anant Madabhushi, Parantu Shah, Michaela Spitzer, et al. Applications of machine learning in drug discovery and development. Nature reviews Drug discovery. 18(6):433-44 1766-477,400 2019.
WF Van Gunsteren and HJC Berendsen. Algorithms for brownian dynamics. Molecular Physics, 45 (3):637-647, 1982.
Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, and Justin M. Solomon. Dynamic graph enn for learning on point clouds, 2019.
Qitian Wu, Chenxiao Yang, Kaipeng Zeng, Fan Nie, Michael Bronstein, and Junchi Yan. Advective diffusion transformers for topological generalization in graph learning. 2023.
Wang Zhang, Tsui- Wei Weng, Subhro Das, Alexandre Megretski, Luca Daniel, and Lam M. Nguyen. Concerne:: A contrastive learning based framework for automated conservation law discovery and trustworthy dynamical system prediction, 2023.
12

---

Published as a conference paper at ICLR 2024
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Symplectic ode-net: Learning hamiltonian dynamics with control. International Conference on Learning Representations. 2019.
Yaofeng Desmond Zhong, Biswadip Dey, and Amit Chakraborty. Benchmarking energy-conserving neural networks for learning dynamics from data. In Learning for Dynamics and Control, pp. 1218-1229. PMLR, 2021.
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A PROOF OF THM.
PROOF. Consider system without any external field where the dynamics is governed only by the internal deterministic forces, for instance, n-spring system and stochastic force terms representing the interaction with the environment, for example, spring-particle systems in fluid system with the stochastic term representing the collision of the particles with the atoms of the fluid environment. Note that the deterministic term consisting of the forces is also termed as the drift term as it contributed to the overall drift of the system, if any. Accordingly, the stochastic term is referred to as the diffusion term. Since, the acceleration of each the particles zero, the net force on the system should be zero, that is, the total linear momentum of the system due to the internal deterministic forces remain conserved. The conservation of linear momentum for this system implies that >i F; O. Assume that the total error between the predicted and actual forces of all the particles at time t to be é. To prove that BROGNET conserve the linear momentum exactly, it is sufficient to prove that é 0. since the stochastic term has zero mean. Consider,
22
n
Σ E -
P.)
F
(15)
Note that Brownian dynamics, which is the over damped limit of Langevin equation (Eg. T the acceleration of each particle is assumed to be zero. Thus, without loss of generality, Σ P. F; = 0, in the absence of an external field. Hence,
12
Σ
F
Fig
jEANS
jENou
Thus, the sum the predicted force of the system remains zero leading to the conservation of linear momentum of the system.
B EULER MARUYAMA INTEGRATOR
Consider a system of particles having masses M and interacting with each other through an interaction potential (()(,), where ((t) represents the position vectors of all the partieles at time t, that constitute a time-dependent random variable. The Taylor series expansion the position of a particle at time t it given by
(16) In the over-damped limit of Langevin dynamics, namely, Brownian dynamics, acceleration and higher order terms are assumed to be zero. Thus, substituting for velocity given by Eq. in Eq. T6
X1(1 + At) = x.()) + ALX,()) + +1)
get
2k
X,(t At) = x,(t) Fi/7
ANG
(17)
%
where AΩ, is a random number sampled from a standard Normal distribution. Note that the Euler- Maruyama is the simplest form of stochastic numerical integrator. Further, higher order integrals can also be employed to integrate the SDEs.
EXPERIMENTAL SYSTEMS
C. SIMULATION ENVIRONMENT
Simulation environment. the training and forward simulations are carried out the JAX environment Sccoenholz & Cubuk (2020). The graph architecture is implemented using the jraph package Godwin et al. (2020). All the codes related to dataset generation and training are available in Baggghhubc.oooMMMR-IIIDD/Boo
Software packages: numpy-I 1. 1,20.3, jax-0.2.24, jax-md-0.1.20, jaxlib-0.1.73. jraph-O.0. 1.0.0.0.............. Hardware: Memory: 16GiB System memory, Processor: Intel(R) Core(TM) 17-10750H CPU @
2.60GHz.
14

---

Published as a conference paper at ICLR 2024
C.2 LINEAR S-SPIING SYSTEM
In this system, n-point masses are connected by clastic springs that deform linearly with extension compression. Note that similar to pendulum setup, each mass m, connected only to two masses m,-1 and m,+1 through springs so that all the masses form a closed connection. Thus, the deterministic force experienced by each mass i due spring connecting Fij = = - Xjll - R,j), where Rıj is the equilibrium length of the spring and represent the un-deformed length and the stiffness, respectively, of the spring.
its neighbor is
C.3 NON-LINEAR S-SPRING SYSTEM
This system is similar to a linear n-spring system with the difference that the spring force non-linear. Specifically, the deterministic force experienced by each mass due to spring connecting to its neighbor j is Fiy [X - Xy Ray) )5 where Ry is the equilibrium length of the spring and k represent the un-deformed length and the stiffness, respectively, of the spring.
C.4 BINARY LINEAR S-SPRING SYSTEM
This system similar a linear r-spring system with the difference that the system consists of two different types of particles with different masses, friction coefficient, and Thus, the deterministic force experienced by each mass due to spring connecting to neighbor j remains the same as Fi₃ = - X; Xill Rıj). where R,j is the equilibrium length of the spring and k represent the un-deformed length and the stiffness, respectively, of the spring. However, the stochastic part varies depending on the type of the particle.
D DATASET GENERATION
All the datasets generated using the known deterministic forces of the systems, along with the stochastic, as described in Section C and 2 For each system, we create the training data by performing forward simulations with 100 random initial conditions. For both linear and nonlinear n-spring system, n = 5 is used for generating the training data. For the binary system, 10 particle system is used for generating the data with the ratio particle type to type 2 as 3:7, A timestep of 10-$3 is used to integrate the equations of motion for all the systems. The Euler-Maruyama algorithm used integrate equations of motion due to its ability handle stochastic differential equations. The details of the parameters used for each of the systems are given below.
Linear n-spring system
Parameter
Value I unit
Mass (M)
Stiffness (k)
1 unit
Damping constant ())
1 unit
TgT Equilibrium length the spring (Ri;)
1 unit
unit
Nonlinear n-spring system
Parameter
Value
Mass (M)
I unit
Stiffnesss (k)
1 unit
Damping constant (C)
1 unit
kgT
1 unit
Equilibrium length of the spring (Ryy)
unit
Binary linear n-spring system
15

---

Published as a conference paper at ICLR 2024
Parameter Mass (M)
Value
1 unit
Stiffness (k)
1 unit
(7))
1 unit
(72)
2 unit I unit I unit
kgT
Equilibrium length of the spring (Riy)
Ratio of particle typel type2
30:70
From the 100 simulations for the systems obtained from the rollout starting from 100 random initial conditions, 100 data points are extracted per simulation, resulting in total of 10000 data points. For training, we create several sets of trajectories with positions from two consecutive timesteps t and + At, where the input for the model is the positions of all the particles at / and the output against which the model trained the position of all the particles at + At employing the loss only trajectories of length one
function defined in Eq. 9 Note that in the present work, timestep only. The trained models are evaluated on 1000 trajectories generated from 100 different initial conditions with each initial condition simulated for 10 random seeds, all of which unseen during the training. The trajectories considered are of 100 timesteps, that is, 0.1 s. Note that this is a task the models are trained for; the training was done only on one timestep trajectories. All the error metrics presented in the work are averaged over the 1000 trajectories for each of the models and casees unless specified otherwise.
For zero-shot generalizability. for each case of different system sizes and temperatures, we directly evaluate the trained models on 1000 trajectories. No fine tuning or additional training is performed for any of the different system sizes or temperatures.
E TRAINING DETAILS AND HYPER- PARERR-AAAMMRESS
The detailed procedures followed for the training of the models and the hyperparameters employed for each of the models, identified based on good practices, are provided this section.
E. TRAINING DETAILS
The training dataset is divided in 80:20 ratio randomly, where the 80% is used for training and 20% used the validation set for hyperparametric optimization. Further, the trained models tested on its ability to predict statistically equivalent trajectory, a task it was not trained on. Specifically, the systems are evaluated on a 0.1s long trajectory, that is 100x larger than the training trajectory. The error metrics are computed based on 1000 different trajectories from 100 initial conditions, each run with 10 random seeds. All models are trained till the losses saturated. A learning rate of 10-3 was used with the Adam optimizer for the training. Detailed hyperparameters employed for all the models are provided next.
E.2 HYPER-PARE PARAMETERS
The hyper-paramette used for different models, namely, NN, BNN, BFGN, BDGNN, and BROGNET are given in the tables below.
NN, BNN
Parameter
Value
Hidden layer neurons (MLP) Number of hidden layers (MLP)
16
2
Activation function Optimizer
squareplus
ADAM
Learning rate
1.0e
Batch size
20
BFGN
16

---

Published as a conference paper at ICLR 2024
Parameter
Value
Node embedding dimension Edge embedding dimension
Hidden layer neurons (MLP)
16
Number of hidden layers (MLP)
Activation function
squareplus
Number of layers of message passing Optimizer
ADAM
Learning rate
1.0e
Batch size
20
BDGNNN BROGNET
Parameter
Value
Node embedding dimension Edge embedding dimension
5
Hidden layer neurons (MLP)
Number of hidden layers (MLP)
Activation function
squareplus
Number of layers of message passing Optimizer
ADAM
Learning rate
1.0e
Batch size
20
BNNequI
Parameters
Value
graph_nee Saaph_net_stte _nee__steppsssss________________
use_sc nonlinearities
True
I'e': "rw_swwish' 'o': 'tanh')
n_element hidden_irreph
1
'128x0e + 64xle 4x2e'
sh sh_irreps
'Ix3e + 1x0e' 8
num_basis
r_max
2.5
madial_net_nonlinee
'raw swish" 64
radial_net_n_hidenn
radial_net_n_layers
1
shift
0
scale n_neighbors
1
10.
scalar_mlp_atd
5.
Optimizer
ADAM
Learning rate Batch size
1.0e
100
F ZERO-SHOT GENERALIZABILITY
Here, we evaluate the zero-shot generalizability BROGNET comparison to BFGN, BDGNN, MIBDGNN, and BNequlP for linear, non-linear, and binary linear n spring systems. Figures H| and]]show the trajectory, Brownian, and position error, respectively, of BFGN, BDGNN, MIBDGNN, BNequIP, and BROGNET evaluated on 50, 500, and 5000 particle systems. Note for linear and non-linear springs, the models are trained on 5-particle systems, while for binary linear spring system, the models trained on 10 10-particle systems. We observe that BROGNET significantly outperforms other models. Further, Fig. K]shows the Inference time w.r.t number of particle for BFGN, BDGNN, MIBDGNN, BROGNET and BNEQUIP.
17

---

Published as a conference paper at ICLR 2024
N=5000
10
linear spring
10
10
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0000
0.10
- non-linear Sew.immaarra..... spriing
Dxc (XII))
10
100
10
5000
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.10
10
binary- limear spring
10
MV
N=10
10
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
Time
BFGN
BDGNN
MIBDGNN
BNequIP
BroGNet
Figure H: Rollout error of BFGN, BDGNN, MIBDGNN, BNequIP, and BROGNET tested on N 50, 500, 5000 for linear (row 1), non-linear spring (row 2), and binary linear spring (row 3) systems evaluated on 1000 forward trajectories. Note that for linear and non-linear systems, the models are trained on 5 5-particle systems whereas for binary system, the models are trained on particle system.
18

---

Published as a conference paper at ICLR 2024
N 5
N 50
500
N=5000
linear spriing
10
0.05
0.10
000
0.05
0.10
0.00
0.10
10
nom spring
RMSE(O,6)
0.06
0.10
0.00
0.06
0.10
0.00
0.10
10
binary-limee spring
N=10
0.10
0.10
0.00
Time MIBDGNN
BFGN
BDGNN
BNequIP
BroGNet
Figure Brownian error of BFGN, BDGNN, MIBDGNN, BNequIP, and BROGNET trained on N system and tested N 50, 500, 5000 for linear (row 1), non-linear spring (row 2), and binary linear spring (row 3) systems evaluated on 1000 forward trajectories.
19

---

Published as a conference paper at ICLR 2024
N= 5
N = 50
N 500
N=5000
linear spring
10
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
0.50
0.05
0.10
non lin-annnnr spring
| -1l, "l
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
0.50
0.05
0.10
binary linear spring
N=10
10
0.00
0.05
0.10
0.00
1.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
Time
BFGN
BDGNN
MIBDGNN
BNcquIP
BroGNet
Figure J: Position error of BFGN, BDGNN, MIBDGNN, BNequIP, and BROGNET trained on N 5 system and tested on N = 50, 500, 5000 for linear (row 1), non-linear spring (row 2), and binary linear spring (row 3) systems evaluated 1000 forward trajectories.
10
BFGN
BDGNN
MIBDGNK
Inference time
BroGNet
BNequIP
10
10 -2
10
10
10
Number of particles
Figure K: Inference time w.r.t number of particle for BFGN, BDGNN, MIBDGNN, BROGNET and BNEQUIP. Note that in this the forward simulation time computed as an average 30 trajectories.
20

---

Published as a conference paper at ICLR 2024
Models
Training time (in sec)
Forward Simulation time (in sec)
NN
391
0.012
BNN
635
0.222
BEGN
1755
0.642
BDGNN
1934
0.411
BNequIP
27535
21.319
BROGNET
2126
0.506
Table A: Training time for 10,000 epochs and forward simulation time for 100 timesteps of NN, BNN, BFGN, BDGNN, NequlP, and BROGNET on linear spring system. Note that the forward simulation time is computed as the average of 100 initial condition with each initial condition simulated for 10 different random seeds (altogether, an average of 1000 trajectories).
(a)
(b)
( 0
10
10
10
Dx,081x;
RMSE(C.8)
[|x 2l.d
-3
10
BroGNet
BroGNet
BroGNet
BroGNet-Rell
BroGNei-ReLU
BroGNet.RelL
10
BroGNet LeakyReLu
10
BroGNet LeskyReLu
BroGNet LeakyRelu
BroCNett Sigmoid
BrogiNe--Simmod
10
BroGNett BroCNNet-Sigmedd
0.00
0,05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
time
time
tmm
Figure L: Comparison of activation function SquarePlus, ReLU, LeakyReL.U and Sigmoid linear 5- spring system with BROGNET. All the results are evaluated based on 1000 forward simulations based on 100 initial conditions, each evaluated with 10 random seeds. G TRAJECTORY VISUALIZATION
Visualizations of the trajectories of linear n-spring systems using ground truth and BROGNET is included as a video in the link: http:///hh....oommA.acccDDDttont BROGNET is trained on 5 5-particle system.
Note that
H SQUAREPLUS, RELU, LEAKYRELU vs SIGMOID
Here, we evaluate BROGNET performance with respect to different activation functions: SquarePlus, ReLU, Lea and Sigmoid. From Fig. L we observe that all the activation functions give comparable performance. We used SquarePlus in order to potentially extend BrogNet to learn energy instead of force in future work. In this case, the activation function needs to double differentiation.
I GENERALIZATION TO UNSEEN Y
Here, we evaluate BROGNET performance with respect different gamma values: 0.1, 0.6, 0.8, 1.0. 1.2, 1.4, and 10. As y increases performance improves. This is due to the fact that for a higher value of y, there will be a smaller standard deviation for the stochastic term.
J IMPORTANCE OF GRAPH TOPOLOGY
Here, evaluate BROGNET when there is correct graph topology vs BROGNETRND when there random connection between senders and receivers. From Fig. N is clear that performance is significantly high with correct graph topology
21

---

Published as a conference paper at ICLR 2024
(a)
(b)
(c)
10
Da
10"2
RMSE(G.
10
10
10
10
10
0.00
0,05
0.10
0.00
0.05
0.10
0.00
0.05
0.
tmo
tme
time
gamma: 0.1
gamma: 0.8
gamma: 1.2
gamma: 10
gamma: 0.6
gamma: 1.0
gamma: 1.4
(a)
(b)
(c)
10
Gy (2)
RMSE(a.d)
10
|x 11-16
10
BroGNet
10
BroGNet
BroGNe:
BroGNeIRND
BroGNetRND
10"
BroONe:RND
0.05
0.00
0.10
0.00
0.05
time
time
fme
Figure N: BROGNET vs BROGNETRND architecture
K PERFORMANCE W.R.T. HYPERPARAMETERS
Here, we evaluate the performance of BROGNET w.r.t. number of layers of message passing [1, 2. 3] and the number of hidden layer neurons [5, 16, 32]. the paper for BROGNET we choose the number of layers of message passing 1 and the number of hidden layer neurons 5. For other architectural details, check Appendix E
22

---

Published as a conference paper at ICLR 2024
(a)
(b)
(c)
10
10
Dx. (X)
10
RMSE d)
10
10
15
[16]
10
(22
10
10
[32]
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
time
6mc
fmc
(d)
(e)
1700 0
10
10
10
De. (.(0)))
RMSE(O.6)
10
[5. 5]
[5, 5]
[5 5]
18. B]
10
18 8]
[8. 8]
(16,
[16. 16]
[10. 16)
[32, 32]
[32, 32]
10
[32,32] 32)
10
0.00
0.05
0.10
0.00
0.05
0.10
0.00
0.05
0.10
(g)
(h)
1%00 0
10"
(0000
RMSE(a.8)
10"
10
55
15. 5,
|
5
8,
8,8,8)
8.
[16, 18, 16]
[16, 16, 16)
116, 16, 16)
[32, 32, 321
10
[32, 32, 32)
(32, 32, 321
0.00
0.05 time
0.10
0.00
0.05 time
0.10
0.00
0.05
0.10
tim
Figure O: Performance of BROGNET architecture w.r.t. hyperparameters
23


Reviews:
[{"review_id": "xQK324d9Ji", "reviewer": "Reviewer_KKvb", "summary": "In this paper, the authors introduce an innovative framework, Brownian graph neural networks (BROGNET), that integrates stochastic differential equations and GNNs to directly learn Brownian dynamics from trajectories. Their method enforces the conservation of linear momentum within the system, leading to empirically observed improved performance in learning dynamics. The authors showcase the effectiveness of BROGNET by applying it to various benchmarked Brownian systems. They also demonstrate its ability to generalize to simulate previously unseen system sizes and temperatures", "questions": "Please see above!", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The main idea of the paper is interesting. It is well-written and the proposed method seems to be novel.", "weaknesses": "Some parts are unclear and require further explanations. There are questions and vague points that need addressing:\n\n1. How does the suggested framework manage noisy or incomplete trajectory data, and is it capable of accurately learning the underlying  dynamics in such cases? \n\n2. How does the choice of activation function affect the performance of the MLPs? Were other activation functions considered, and if so, how did they compare to the chosen function?\n\n3. Can the MLP be replaced with other types of neural networks, such as convolutional neural networks or recurrent neural networks? How would this affect the performance of the proposed framework?\n\n4. Can you provide more details on the scalability of the proposed framework? How does the computational complexity scale with the number of particles, and how does this affect its applicability to large-scale systems?\n\n5. How were the hyperparameters chosen, and how does the choice of hyperparameters affect the performance of the proposed framework?\n\n6. How the choice of graph topology affects the performance of the proposed framework? Were other graph topologies considered, and if so, how did they compare to the chosen topology?\n\n7. How does the proposed method handle systems with external fields or other sources of non-deterministic forces?\n\n8. Can you provide more details on the benchmarked Brownian systems used to evaluate BROGNET's performance? How do these systems compare to real-world applications?\n\nI am happy to increase my score if the authors could address my concerns.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PAxJWMXgxP", "reviewer": "Reviewer_VaBU", "summary": "This paper introduces BroGNet, a GNN that is momentum conservative designed based on SDEs for Brownian dynamics learning.\nThe authors provide a very thorough introduction and related work section to motivate the paper and to provide the reader with sufficient background. Then, the method is presented, followed by several experiments in different scenarios. The proposed method seems to significantly outperforms existing methods.", "questions": "Regarding the dynamic graph used here, how different is the proposed procedure than [8] ?\n\nRegarding equation (8), this seems a bit like a discretized version of an advection operator (see [9,10]) for example. Can the authors expand on this point and clarify the differences?\n\n[8] Dynamic Graph CNN for Learning on Point Clouds\n\n[9] ADR-GNN: Advection-Diffusion-Reaction Graph Neural Networks\n\n[10] Advective Diffusion Transformers for Topological Generalization in Graph Learning", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is mostly easy to follow and read.\n\nThe authors provide a very good background section to explain different terms, such that even non expert readers can understand the paper.\n\nThe experimental section looks promising under various settings.\n\nThe authors provide good explanations about the baselines and the experiment details.", "weaknesses": "Missing literature about GNNs: while this paper is concerned with learning brownian dynamics from data, there is a complementary topic in GNNs and that is the design of GNN architectures inspired by ODEs. I believe that the authors should add a discussion to the related work section to clarify the difference between the two. Some references are provided in [1-5].\n\nMissing literature about Neuro ODEs: please see [6,7].\n\nIt is not clear why the authors propose to use the square plus activation. Is there a specific reason? (besides the experimental result provided in the appendix)\n\nReading the appendix, I understand that the authors used only one message passing layer in their implementation. Can you please elaborate on this point? What would the performance be like when adding more layers?\n\nRefereces:\n\n[1] Graph Neural Ordinary Differential Equations  \n\n[2] GRAND: Graph Neural Diffusion \n\n[3] PDE-GCN: Novel Architectures for Graph Neural Networks Motivated by Partial Differential Equations\n\n[4] Anti-Symmetric DGN: a stable architecture for Deep Graph Networks  \n\n[5] Graph-Coupled Oscillator Networks\n\n[6] Stable Architectures for Deep Neural Networks\n\n[7] Deep learning-based numerical methods for high-dimensional parabolic partial differential equations and backward stochastic differential equations", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "iFh4HSB9q0", "reviewer": "Reviewer_cdjq", "summary": "The paper proposed Brownian graph neural networks (BROGNET) which is a new framework combining stochastic differential equations and graph neural networks to learn Brownian motion dynamics directly from trajectories. The architecture ensures linear momentum conservation, leading to improved learning of dynamics. Several baselines were proposed for comparison due to the limited existing benchmarks. The BROGNET's distinctive momentum conservation feature made it significantly superior to all other baselines. It also demonstrated the ability to generalize to much larger system sizes and different temperatures than those seen during training.", "questions": "1.\tTypos: 2nd line under equation 7: ativation \n\n2.\tIs there any specific reason to use squareplus as activation function? I understand the comparison experiments with ReLU in appendix H which shows similar performance, but why squareplus is chosen at firsthand as it’s less popular?\n\n3.\tEquation 10 is questionable. \\Delta \\Omega should have the magnitude of \\Delta t, but the following sentence mentioned “is a random number sampled from a standard Normal distribution”. Do you miss a magnitude of \\Delta t?\n\n4.\tWhat does || || mean in equation 5,6?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. I do like the design of predicting interacting forces rather than total forces on each node which naturally conserves total momentum conservation. Such “hard constraint” (or physics-based inductive biases in the paper) not only rigorously adheres to physical principles but also significantly enhances the model's performance, compared with more straight forward methods such as adding regularizers to penalize the physics violation. I suggest adding a few sentences in the first paragraph to explicit distinguish the “hard constraints” vs “soft constraints”.\n\nThere are a few additional papers in this trajectory worth mentioning. The natural constraint design in this paper is more less like ref [A], as the constraint is pure summation. While it can be generalized as special cases of Ref[B],[C],[D] as well.\n\n[A] A machine learning-aided global diagnostic and comparative tool to assess effect of quarantine control in COVID-19 spread, Patterns, 2020.\n\n[B] ConCerNet: A Contrastive Learning Based Framework for Automated Conservation Law Discovery and Trustworthy Dynamical System Prediction. ICML 2023.\n\n[C] Learning Physical Models that Can Respect Conservation Laws. ICML 2023.\n\n[D] Unravelling the performance of physics-informed graph neural networks for dynamical systems. NeurIPS, 2022. (this is already cited.)\n\n2. The model can generalize to unfamiliar system sizes and temperatures with zero-shot learning.\n\n3. The outperforms existing baselines in various tasks.", "weaknesses": "1.\tMy major concern lies on the comparison baselines. Firstly, the prior work baselines (BFGN, BNequIP) do not seem very strong to me. However, I’m not an expert in partical-based systems and I’m not expecting each accepted paper will include comparisons with all the popular models. More importantly, to show the benefits of the physics-informed inductive bias, it is worth comparing the other model with the same backbone structure + training with regularization on the physics violation. I appreciate the ablation study with BDGNN which essentially shares the same backbone model. But adding a comparison experiment by training BDGNN with regularizing on momentum conservation will better prove the power of inductive bias. \n\n2.\tRegarding the sde integrator, I’m a bit concerned about the gradient stability due to the stochastic integral. Did you meet any issues when the random distribution of noise term leading to large error during back-propagation? Or do you need some sampling method to get the approximation of the drift term along with the variance of the stochastic term?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ok0DP6NnD6", "reviewer": "Reviewer_Va4f", "summary": "## Summary\n\nThe paper proposes a Graph Neural Network (GNN)-based method for simulating the over-damped limit of Langevin dynamics, which reduces to Brownian dynamics when no acceleration is present. Contributions include the use of additional MLPs for predicting random diffusion and decoding edge latent as interacting forces to enforce linear momentum conservation. However, the paper has several significant shortcomings, including a lack of motivation, limited benchmarking, and insufficient definition and references. Due to these issues, I recommend rejecting the paper.\n\n## Detailed Comments\n\n### 1. Lack of Motivation and Benefits\n\nThe paper does not provide adequate motivation for replacing traditional simulation techniques for the over-damped Langevin system with GNNs. Given that the dataset is generated from simulators that can easily model such systems, the authors need to justify the utility of their approach. This could be in the form of improved simulation speed or to-reality accuracy, although the latter would require validation beyond simulated data.\n\n### 2. Lack of Definitions and References\n\n- In Figure 1, the term \"ohe(type)\" is used without definition or context, making it unclear to the reader.\n- The \"squareplus\" activation function is mentioned but not cited.\n\n### 3. Limited Benchmarking\n\nThe paper restricts its experiments to simple spring systems, providing only a narrow validation of its methodology. Prior work in this domain typically includes experiments on 3-4 different datasets to establish the method's applicability.\n\n## Conclusion\n\nWhile the paper introduces a GNN-based method for simulating over-damped Langevin dynamics, it suffers from multiple critical flaws, including a lack of clear motivation, insufficient references, and limited benchmarking. These drawbacks severely compromise the paper's value and applicability. Therefore, I recommend rejecting this submission.\n\nThe authors should consider submitting to a workshop on this specific topic or, including more datasets to show the method's capability.", "questions": "How hard is it to simulate Brownian motions? If not hard, w.r.t. simulation time or numerical or modeling challenges, there is no point in replacing it with NN?\n\nIf otherwise, please show the challenges in your draft.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Small contributions like: decoding edge as force so linear momentum is conserved.", "weaknesses": "See above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)