Physics-Inspired All-Pair Interaction Learning for 3D Dynamics Modeling
Authors: Kai Yang
Yuqi Huang
Junheng Tao
Wanyu Wang
Qitian Wu
Venue: ICLR 2026 Poster (2026)
Keywords: 3D Dynamics Prediction
Attention Mechanism
Primary area: learning on graphs and other geometries & topologies

Abstract: Modeling 3D dynamics is a fundamental problem in multi-body systems across scientific and engineering domains and has important practical implications in trajectory prediction and simulation. While recent GNN-based approaches have achieved strong performance by enforcing geometric symmetries, encoding high-order features or incorporating neural-ODE mechanics, they typically depend on explicitly observed structures and inherently fail to capture the unobserved interactions that are crucial to complex physical behaviors and dynamics mechanism. In this paper, we propose PAINET, a principled SE(3)-equivariant neural architecture for learning all-pair interactions in multi-body systems. The model comprises: (1) a novel physics-inspired attention network derived from the minimization trajectory of an energy function, and (2) a parallel decoder that preserves equivariance while enabling efficient inference. 
Empirical results on diverse real-world benchmarks, including human motion capture, molecular dynamics, and large-scale protein simulations, show that PAINET consistently outperforms recently proposed models, yielding 4.7% to 41.5% error reductions in 3D dynamics prediction with comparable computation costs in terms of time and memory. Our codes, baseline models and datasets are available at https://github.com/Icarus1411/PAINET.

Published as a conference paper at ICLR 2026
PHYSICS-INSPIRES ALL-PAIR INTERACTION LEARN- ING FOR 3D DYNAMICS MODELING
Kai Yang!.? Yuqi Huang", Huanne,Junheng Junheng Tao Wanyu Wang", Qitian Wu?* "Department of Computer Science and Engineering. Shanghai Jiao Tong University "Shanghai Artificial Intelligence Laboratory School of Artificial Intelligence, Shanghai Jiao Tong University
*SJTU Paris Elite Institute of Technology, Shanghai Jiao Tong University DDepartment of Physics, Harvard University
-arvard-MITT Center for Ultracold Atoms
FEiic and Wendy Schmidt Center, Broad Institute of MIT and Harvard ficarus1411, hhhuangq, wwyspeit it e jtu edu. en jtaoßfas. harvard. wugitiannmit edu
ABSTRACT
Modeling 3D dynamics is a fundamental problem in multi-body systems across scientific and engineering domains and has important practical implications in trajectory prediction and simulation. While recent GNN-based approaches have achieved strong performance by enforcing geometric symmetries, encoding high- order features or incorporating neural-ODE mechanics, they typically depend on explicitly observed structures and inherently fail to capture the unobserved in- teractions that are crucial to complex physical behaviors and dynamics mecha- nism. In this paper, propose PAINET, principled SE(33)-quivvaiiat neu- ral architecture for learning all-pair interactions in multi-body systems. The model comprises: (1) novel physics-inspines attention network derived from the minimization trajectory of an energy function, and (2) parallel decoder that preserves equivariance while enabling efficient inference. Empirical results on diverse real-world benchmarks, including human motion capture, molecu- lar dynamics, and large-scale protein simulations, show that PAINET consis- tently outperforms recently proposed models, yielding 4.7% to 41.5% error re- ductions in 3D dynamics prediction with comparable computation costs terms of time and memory. Our codes, baseline models and datasets are available at https: htt//gittuuuut wttpa//c.....aaaa1111111.1..
INTRODUCTION
Modeling the 3D dynamics is a fundamental challenge across wide spectrum of scientific and engineering disciplines, encompassing molecular dynamics (Schütt et al., 2017; Unke et al., 2021), celestial mechanics (Sanclez-Gonzalez Gonzalez et al., 2019), physical simulation (Battaglia al., 2016), etc. Traditional approaches are grounded in classical mechanics, leveraging physical laws such as New- tonian dynamics or force-field models. Despite decent accuracy and interpretability, these methods require substantial computational costs for solving partial differentiable equations in systems with complex interactions or large numbers of particles (Kipf et al., 2018).
Recent years have witnessed the promise of data-driven approaches, particularly those based on deep learning methods. Graph neural networks (GNNs), which treat particles as nodes and their interactions edges, have emerged powerful paradigm for learning physical dynamics directly from observed data (Battaglia et al., 2016; Sanchez-Gonzalez et al., 2019; Pfaff et al., 2021). A typical GNN model is Equivariant Graph Neural Network (EGNN) (Satorras al., 2021) which has shown strong capability for modeling N-body dynamics by preserving SE(()-equivariance, a property recognized as essential geometric prior in physical systems (Bronstein al., 2017). It inspires a series of follow-up extensions with higher-order features (Liu et al., 2022; Cen et al.,
"Corresoonding author.

---

Published as a conference paper at ICLR 2026
2024), space transformation (Du et al., 2022), equivariant local frames (Yin et al., 2025), local attention (Fuchs et al., 2020), and temporal convolution (Xu et al., 2024), which pushes the frontier of Mfate-of-the-al performance diverse benchmarks.
However, current models predominantly hinge on explicitly observed structures, e.g., an adjacency graph extracted from distances of particles, for computing representations. This design neglects unobserved all-pair interactions among particles, which play an indispensable role in the dynam- ies mechanism. For example, in atomic and molecular systems, the long-range tails of the Van der Waals-type potentials (London, 1937) are often overshadowed by the stronger short-range in- teractions (French et al., 2010), yet neglecting them degrades manghorizon trajectory accuracy and obscures long-range correlations (Sagui & Darden, 1999; Defenu et al.. 2023). Another example is the spontaneous formation of unobserved structures, such as in erystallization (Schreiber Fersht, 1996) and protein folding (Richards, 1977; Cao, 2020). In these cases, the observed structures rep- resent only a transient temporal snapshot; stipulating the model on fixed observed structures may lead to systematic bias and significant error accumulation in time, as the true interaction landscape dynamically evolves (Taudt et al., 2015).
Despite the significance, modeling unobserved interactions is non-trivial since the searching space for latent structures (which involves the potential interactions between arbitrary pairs particles) is exponential to the number of particles and the optimal structures can dynamically evolve across time. Without a principled formulation for unobserved interactions and how they couple with the dynamics mechanism, uncovering the latent structures is underconstrained, computationally prohibitive, and also prone to spurious correlations. Furthermore, it is not straightforward to preserve the SE(3)- equivariance property for models accommodating the latent interactions beyond observed structures.
In this paper, we propose PAINET (Physics-inspired All-pair Interactions Network) for 3D dy- namics modeling. Starting from an energy function that quantifies the latent interactions among particles through smoothness criteria of particle embeddings, we derive a principled attention mech- anism with adaptive pairwise mappings that captures long-range, particle-type typ--secciii dependen- cies. Moreover, the model is implemented with a parallel decoder that preserves -qquivrraae while maintaining efficiency in inference.
We evaluate PAINET on diverse real-works datasets, including human motion sequences, molecular dynamics (MD17), and large-scale protein simulations. Empirical results show that PAINET con- sistently outperforms recently proposed methods with up to 41.5% / 18.0% / 4.7% decrease of mean square errors (MSEs) in motion capture/ / MD17 protein dynamics simulation, while consuming comparable computation costs those competing models. Furthermore, ablation studies that corroborate the efficacy of the key proposed components and demonstrate that the computation costs of PAINET grow nearly linearly w.r.t. partiele numbers and time steps.
conduct comprehensive
The contributions of our work are summarized below:
An energy-based formulation for latent structure learning in 3D dynamics. We formulate the problem of uncovering latent structures in 3D dynamics as minimization of an energy function that characterizes the smoothness of particle embeddings in latent space, which paves the way principled approach for accommodating the unobserved interactions into 3D dynamics modeling.
A physics-isspired attention network with (33)iuauuvaa and efficient decoding. On top of the formulation, we propose a new attention network with adaptive pairwise mappings to capture long-range, particle-ypee-peecfii dependencies beyond observed structures at each step. The model is compatible for parallel decoding that preserves -qquivarria with efficient inference.
Comprehensive empirical evaluation and comparison in multi-faceted aspects. As demon- strated by experiments on eleven datasets ranging from human motions to molecules and proteins, the proposed model achieves 4.7%-41.5% improvements in terms of 3D dynamics prediction as measured by MSE and shows desired scalability to large-scale multi-body systems with increasing numbers of particles and simulation time steps.
2 PRELIMINARIES
Notations and Problem Formulation. The 3D dynamics modeling can be formulated as learning over a sequence of geometric graphs g(t) = (V, F, v(9), where V denotes the set of particles (nodes) and & represents the set of observed interactions (edges). F NNxdy denotes the
2

---

Published as a conference paper at ICLR 2026
node features and A E WNNxx.. is edge-attribute tensor, where aij R denotes the feature vector associated with the edge between nodes and j. x = [x]
E R Nx3 and V [v(") :..... : R2x3 denote the position and velocity matrices at time step t. respectively, where (t) and V denote the position and velocity of node i in the 3D coordinate system. The problem can be described as: given the initial observation g (V. E, F, A, ( V(0)), the model aims at predicting the future dynamics trajectory (
xi
[ (1) X
The notion of equivariance formalizes a type of symmetric property in geometric space regardless of the chosen coordinate system. Formally, let g E G denote an abstract group, then a function μ X + y is defined as equivariant with respect to g if there exists S,, y - y for the set of transformations T X - X' on g such that: (x]z three types of equivariance in special Euclidean group SE(3), including rotations, translations, and permutations. For 3D dynamics, rotation matrix Q E (3)) respectively, where SO(3) is special Orthogonal group in 3D space.
Equivariance.
(z((x)). In this paper, we focus on
can instantiate Tg and Sg as translation g R and orthogonal
Then, an -333ivaauvva predictor f for next-step positions satisfies:
((X(() +1g]) Q fx (1) +1g
1 1 [1,111
(I)
inductive bias for plausible prediction. Further details on translation, rotation and permutation properties are presented in Appendix B.
We enforce this property throughout the architecture
Equivariant Graph Neural Networks. Standard GNNs are permutation-equuvarias with respect to node relabeling. Beyond this, Equivariant Graph Neural Network (EGNN) (Satorras et al., 2021) attains SE(3) eequivrraae via the message passing layers that operate on invariant quantities (e.g., pairwise distances and relative positions). A typical EGNN layer updates node embeddings h, E R
and positions x; E I from layer to + 1 via:
m
2m h
(1)
h.
(1)
||x
x
1012 agg),
h on
(1+1)
h
Σ
mij
(1)
JEN(
(1+1)
(1) x
(x( (1)
) 2z (m!)) (?)
1
JEN())
where om. Φn and di are commonly instantiated as MLPs, and (i)) denotes the set of neighbored nodes of node i. Built upon the main design of EGNN, recent model EGNO (Xu et al., 2024) additionally incorporates temporal convolution to capture temporal correlations, while a follow- up work HEGNN (Cen et al., 2024) augments EGNN with higher-order interactions to enhance the model's s expressivity. Furthermore, GF-NODE (Sun et al., 2025) further adopts a neural-ODE cittinuuuu-ttmmme evolution on graphs. Despite their effectiveness for
formulation for accommodating patterns hinging on observed structures, these models restrict message passing at each layer to local neighborhoods stipulated by & and overlook the latent interactions that are unobserved yet make large differences in dynamics mechanism.
modeling the spatial-tempon
3 LEARNING ALL-PAIR INTERACTIONS FOR 3D DYNAMICS
In this section, we propose PAINET (Physics-inspired A All-pair Interactions Network) for 3D dy- namics modeling. We first derive a physics-inspired feed-forward layer from the minimization of a regularized energy (Sec. 3.1.. We further introduce the model architecture including a principled attentive encoder as well as parallel equivariant decoder that allows efficient inference (Sec. 3.2). Finally, we present the training and inference details for 3D dynamics prediction (Sec. 3.3).
3.1 MODEL FORMULATION
Energy with Latent Structures. Learning unobserved interactions among particles is non-trivial, as the searching space goes to the exponential order w.r.t. the number of particles. To cast the problem into a solvable formulation, we start from an energy function that characterizes the plausible latent structures through regularizing the smoothness of partiele embeddings in latent space:
E t; (px)) = Σ [h.
h 1 + Σ Pig (lbb h h)1l))
(3)
3

---

Published as a conference paper at ICLR 2026
Input t-0
Ourpats
Equivariant
x
Fopasaiion ONN
Encoder
Decoder
Lagen
Priucipled Atteation Layer
Primution
Mapping
Energy
Miln
Figure I: Illustration of PAINET framework. The model takes the initial state (including positions, velocities and observed features such as edge attributes of particles) as input and encode observed information into particle embeddings in latent space. The particle embeddings a stack of principled attention layers, where each layer corresponds to a descent step on the energy. The attention network includes adaptive pairwise mappings to capture long-range, particle-type- specific dependencies. For decoding, the model harnesses equivariant GNNs that incorporates the observed structural information without breaking SE(3)-qquivariance and generates predicted tra- jectory of particles at multiple time steps in parallel.
updated through
where h. denotes the current embedding of particle at layer t, h, represents its updated (to-be- optimized) embedding for the next layer, and H [h :...; hyy is the stack of the embeddings of N = [ particles. Without loss of generality, Pty IR - R. is non-linear non-decreasing function specific to particle pair (i, and À 0 is trading weight. Eqn. 3 which extends the quadratic energy introduced by the prior work (Zhou et al., 2004) integrates the local and global smoothness criteria: the first term regularizes the distance between the updated embeddings and the current embeddings, preventing the abrupt change of embeddings between adjacent time steps; the second term penalizes the distance among updated embeddings of all particles, encouraging the global smoothness stipulated by latent structures. The latter essentially accommodates the unobserved interactions in an implicit manner: any two particles whose embeddings are close (resp. distant) in latent space add small (resp. large) quantities to the energy. In this sense, Eqn. 3 characterizes the internal consistency of particle embeddings of any given layer. To improve its robustness, it natural assume the concavity of Pij that avoids over-regularization on large differences (Yang et al., 2) that is beneficial for filtering noisy interactions.
Ford-Forward Layers from Energy Minimization. In common physical systems, the evolution of particle states often transitions from high-energy to low-energy regimes and ultimately reaches cer- tain equilibrium. Projected to our context, where our goal is to learn desired particle representations that accommodate the unobserved interactions, descending the energy defined by Eqn. 3 corresponds to pursuing final state that optimizes the internal consistency among particle embeddings in latent space. From this physics-inspiedd standpoint, we next derive feed-forward layer that updates par- ticle embeddings whose formed trajectory in latent space minimizes Eqn. 3. To this end, we provide the following result that suggests a closed-form solution for the optimal embedding trajectory where each step (i.e. ferdwrd layer) contributes to a rigorous descent step on the non-convex energy.
Theorem For any energy function defined by Eqn. with given À > there exists 0 < such that the iterative updating rule (from the initial state h (0)
(e)
(h2)
h (t-1) (1- 17 +
Σ 5 Σ
h
where wiy (0)
(4)
w (
Oh2
AA= 11
yields descent step on the energy, i.e., ((H 1111; (pa))) S (HH() 4; (p5))) for any / 1.2 1.
The proof follows the principles of convex analysis and Fenchel duality (Rockafellar, 1970). This result implies principled attentive feew-forward layer where the attention of particle pair (i, is
4

---

Published as a conference paper at ICLR 2026
given by the gradient of the pairwise penalty function Pij evaluated at the current step. Assuming that fay denotes the first-order derivative of Pij. we have the updating rule of one worwwrd layer:
[ag - h Jim hm
h (++1)_ = (1 - n) + Σ
(5)
where fiy is non-negative, decreasing function w.r.t. h² ||,; h such that Pig satisfies the concavity and non-negativity.
Physics-Inspired Attention Mechanism. The instantiations of fú have much flexibility. Here, we adopt concise yet generic polynomial potential form that respects the parity symmetry of h. In modern physies, the Landau-Ginzburg potential energy form (Laandau et al., 1937) plays a central role in the phenomenological explanation of the formation of matter states ranging from super- conducting (Ginzburg et al., 2009) and ferromagnetic materials (Landau et al., 1937) to complex biochemical media (Hohenberg & Krekhov, 2015). Moreover, it has inspired the Anderson- Hidgrson-HHigss mechanism for spontaneous symmetry breaking in quantum field theory (Anderson, 1963; Higgs, 1964). As per this physics motivation, instantiate Psj as quadratic potential form:
Pig (h²) aiyh² bigh", fig (h²)
ai 26,gh²
where Gij > 8b,y and bij > 0 such that the non-negativity and decreasing property of fú are guaranteed. Furthermore, we assume particle embeddings are normalized into unit sphere, i.e. ||hil22 1, so that llh; h = 2 2h, hy. This property can be easily satisfied in practice by using layer normalization at each layer. Therefore, we have the updating rule of an attention layer:
Φij + Whi (
h
h (++2) = (1- + E
hy' where in
(7)
Em Psm + e'm TTE
| 0
Here we replace agg. bis in Eqn. with variables dyg by satisfying dig > viy > O. They give rise to adaptive linear mappings for specific particle pairs which is beneficial for capturing informative unobserved interactions specific to particle types and states.
3.2 MODEL ARCHITECTURE
Based on the above model formulation, we next present the detailed architecture of the proposed model PAINET for 3D dynamics modeling. Fig. provides an overview of our model framework.
Principled Attention Network. We extend the attention layer derived in Sec. 3.1 to parameter- ized neural networks and present the matrix-form updating rule as implemented with modern deep learning library. For particle embeddings at layer t denoted by ( R first compute its key, query and value matrices: Q(+) ! (()), ( ! V W where (((() ": + ((()) and (0 R the 1.2-normalized versions of Q() and K respectively. Then following Eqn. 7, the equivalent matrix-form attention-based updates at layer i can be written as:
K ((( e IRdxd are trainable weights and d the hidden size. We assume Q(+) E Nxd
((++1) diag '(11 +
()))))))
H(++)) = (1- H ( m D (t+1) ( ( + (T O (Q (KO)))
where the operator diag(-) converts the input vector into diagonal matrix, 1 E IR Nx is an all-one vector, & denotes the Hadamard product, 7 E (0. is a hyper-parametes and & (Ptyk. and T [Pillijv serve for adaptive pairwise mappings. The model described so far treats 'layers' as equivalent to 'time steps', which case each time step the embeddings are updated by one attention layer defined by Eqn. 8. One can extend the model with multiple attention layers per time step, yet as validated in our experiments this simpler architecture using one attention layer per time step yields the optimal performance (Sec. 4.5).
Adaptive Pairwise Mappings. Motivated by the fact that Van der Waals Waals-type potentials employ dif- harbon-hydrogen) (London, 1937), ferent coefficients across particle types (c.g., carbon-carbon interactions. Assuming the num- we specify o and to parameterize ber of particle types (e.g., atomic types for molecules) is F,, we let 2, (0, 1] be one-hot
5

---

Published as a conference paper at ICLR 2026
vector indicating the type of particle (with 21,; 1 iff the type index is k) and stack them as Z = | lookup embeddings E E R ExE and By E REXE, The pairwise weight matrices are then calcu- lated by:
..... ; y {0, 1]NNEE On top of these. we introduce two learnable particle-type-speeeee
S1 a (ZE,
v = =2 o (ZE ZZ
where (()) is specified as s gmo id and 81, 82 > are two learnable scalars, which ensures the non-negarivity of Øi and thg.
Jt=1 are obtained through re- currently operating 7' attention layers across time steps, we integrate them with the observed graph structures through a parallel decoding network. For each time step t, the decoder generates the pre- dicted positions (!) from the initial positions x and velocities V(") via message passing of the current particle embedding H()) over the observed graph structure A. The predictions for 7 time steps can be achieved in parallel:
Parallel Equivariant Decoder. After the partiele embeddings (H())
X (0) = decoder (H (0) V(), ISIST,
(10)
where we instantiate decoder as EGNN defined by Eqn. 2. The full feed-forward flow of our model is described Alg. 1. Prior works have shown the importance of -Euuivaquieee for physical plausibility in 3D dynamics modeling (Satorras et al., 2021; Xu et al., 2024). We can show that our model maintains -33ui--quiva throughout the encoding and decoding process. As proof sketch, the embedding h; remains invariant to rigid transformations of the input coordinates, while the position x, remains equivariant after EGNN. As a result, the model preserves the desired transformation properties under any rotation Q E SO(3) and translation g 1R3 The complete proof is deferred to Appendix E.
3.3 TRAINING AND INFERENCE
Training Objective. With initial positions ( and (optionally) velocities V(O), the initial particle embeddings are given by an MLP: ())) MXP (0) V(").. The model trained in a supervised manner to predict future trajectories [x (0) Jr=1- Let (() denote the predicted positions at time step
t and (t) is the predicted position of partiele i. The training loss is the averaged mean squared error:
1
Luaj
(1I)
Inference. At inference time, PAINET first computes sequence of particle embeddings (H(0)7 via recurrently applying the attention network for T steps. These embeddings are then used to gener- ate the predicted positions {X(0) in parallel by the decoder, as shown in Fig. and specifically Alg. 1. This recurrent-paree parallel framework accommodates all-pair interactions in encoding phase and enables efficient generation that preserves (3(()))-uvvarianee in the decoding phase.
Ji=1
4 EXPERIMENTS
In this section, we apply PAINET to various real-world datasets that involve 3D dynamics modeling in different domains validate the practical efficacy of the proposed model.
Evaluation Protocols. Following Xu et al. (2024), we conduct comparisons on two tasks, referred to as state to state (S25) and state to trajectory (S2T). The task of s2s aims to directly predict the final state (T - 1) and assesses the Mean Square Error (MSE) from the ground truth, referred to Final Mean Squared Error (F-MSE). The task of S2T aims to predict the trajectories including multiple designated future time steps (t = 1,... .T where T > 1), as measured by MSE averaged across time steps, referred to as Average Mean Squared Error (A-MSE) (see Appendix
for details).
Competitors. We compare PAINET with representative classic models including Linear Dynamics (Linear) (Satorras et al., 2021), 1). Radial Field Networks (RF) (Köhler et al., 2019), Message Passing Neural Networks (MPNN) (Gilmer et al., 2017) and EGNN (Satorras et al., 2021) the-art models for 3D dynamics prediction including EGNO (Xu et al., 2024), HEGNN (Cen et al., 2024) and GE-NODE (Sun et al., 2025). For $25 task, compare with all of these models; for
1021) as well of

---

Published as a conference paper at ICLR 2026
Table 1: Test MSE (L), training GPU memory cost and training time per epoch on Motion Capture. F-MSE is reported for S2S (T = 1) and A-MSE for S2T (T = 5). We mark the test performance of our model and emphasize the best and the rumner-up in bold and underlined. All improvements with significance level p < 0.05 by Wilcoxon signed-rane raak test are marked with *.
Walk (325)
Run (323)
Model
F-MSE 2)
GPU (GB)
Time (8)
E- (x10 1)
GPU (GB)
Time (S)
I Lincar (Satorras et al., 2021)
971 + 0.01
0.01
0.02
625 + 0.00
0.01
0.03
RF (Kühler et al., 2019)
93.1 5.28
0.06
0.12
9.06 0.21
0.06
0.11
MPNN (Gilmer al., 2017)
27.2 2.15
0.12
7.42 1422:000
EGNN (Satorras
2021)
26.0
0.21
0.19
4.21 0.78
0.21
0.22
ClofNet (Du et ul., 2022)
12.6 = 1.21
0.29
0.18
8.42 +: 0.69
0.26
0.24
EGNO (Xu et ul., 2024)
14.2 2.62
0.23
0.33
4.15 0.42
0.23
0,30
HEGNN (Cen al., 2024)
16.8 5.55
0.25
5.48 4 0.87
0.39
GF-NODE (Sun al., 2025)
15.7 1.34
0.29
0.62
3.87 0.43
0.28
0.67
PAINET
8,45 0.23
0.26
0.20
3,50 1: 0.27
0.24
0.21
Walk (#2T)
Run (321)
Model
A-MSE (x10 !)
GPU (GB)
Time (8)
A-MSE (x10 1)
GPU (GBB)
Time (s)
EGNO (Xu et al., 2024)
1.48 0.55
0.58
0.18
5.70 + 1.43
0.58
0.19
HRGNN (Cen et al., 2024)
2.47
0.89
0.68
0.20
9.49 + 0.02
0.68
0.22
GF-NODE (Sun al., 2025)
1,25 0.17
0.77
1.10
7.49 +: 1.69
0.77
0.76
PAINET
0.86
1.04
3.331 0.12
0.31
S2T, since the classic models are not originally designed for this task, we only compare with the lat- est models EGNO, HEGNN and GF-NODE. For fair comparison, we use the same hyperparameter searching space on validation data for all models (see Appendix F.6 for details).
4.1 MOTION CAPTURE
Dataset and Implementation. CMU Motion Capture dataset (CMU Graphics Lab, 2003) contains 3D joint trajectories collected from various human motion sequences. Following prior works (Huang et al., 2022: Xu et al., 2024), we uniformly discretize the trajectories and focus on two actions: Sub- ject #35 (Walk) and Subject #9 (Run). More implementation details are provided in Appendix F.2.
Results. The results are reported in Table 1. PAINET consistently achieves the lowest prediction error in all the tasks. For S2s, PAINET yields 40.5% and 9.6% relative improvements over the most competitive baseline in Walk and Run, respectively. For S2T, F PAINET exceeds all the competitors with 1.2 and 41.5% decrease on A-MSE in the task Walk and Run, respectively, demonstrating its superior capability for predicting 3D dynamics. Additionally, the model consumes comparable time and GPU memory cost compared to the competitors, which corroborates the advantage of PAINET using the same level of computational costs achieve superior performance.
4.2 MOLECULAR DYNAMICS ON SMALL MOLECULES
Dataset and Implementation. MD17 (Chmiela et al., 2017) is widely used benchmark for simu- lating molecular trajectories governed by quantum- forces. Following the common pro- tocol, we evaluate the models with 52s (T = and s2T (T 8), using 500/2000/2000 trajectories for training/valiaationnnsss. More implementation details are deferred to Appendix F.3.
Results. The results are presented in Table 2 and 3 for 32S and S2T, respectively. Our model consistently outperforms the competitors across all cases with substantial improvements in quite a few cases. Specifically, the performance improvements (which are all statistically significant with least p < 0.05) range from 1.6% to 18.0% and notably, the improvements on bengene for s2s and naphthalene for S2T 13.4% and 18.0%, respectively. As demonstrated in Fig. 2, the predicted 3D trajectories of PAINET preserve the essential structural characteristics and contribute to steady results in long time horizons in contrast with the competitor GF-NODE.
4.3 MOLECULAR DYNAMICS ON LARGE MOLECULES
Dataset and implementation. While MD17 contains mainly small organic molecules, con- ducted experiments on MD22 (Chmiela et al., 2023) for evaluating the modeling of complex non- local interactions in larger systems adequately. We evaluate the models with $2$ (T 1) and S2T (T' = 5), using 500/2000/2000 trajectories for training/validationnees More implementation details deferred to Appendix F.3.
7

---

Published as a conference paper at ICLR 2026
Groud Truth
PAINET
F-M
9.10 9.50
GF-NODE
F-MISTIX 1000)
8-92
3.95 9:36
0.24
9.52
9.44
Figure 2: Representative snapshots of aspirin molecular dynamies: the top row shows the ground- truth trajectories, the middle row shows the predictions from PAINET and the bottom row shows the predictions from GF-NODE, with corresponding F-MSEs reported across time steps (where in the parenthesis we show the actual time stamps). More results are presented in Appendix F.10.
Table 2: Test F-MSE (1) for S2S (T = 1) on MD17.
828 (T' 1)
Model
aspirin (x10
benzene (x10 1)
ethanol (x10 3)
maloualdehyde (x10
Linear (Satorrus al., 2021)
12.0 0.000
+ 0,000
5.62 0.000
2.14 0,000
MPNN (Gilmer et al., 2017)
9.63 + 01053
5.73 + 0,569
4,92 + 0.084
1,31 + 0.1044
RF (Kobler et al., 2019)
10.9 1:9::00000 0.003
12.6 + 0.028
4.64 + 0.001
1.29 : 0.006
EGNN (Satorras
2021)
0.158
6.41 + 0.246
4,65 0.005
128 : 0.001
ClofNet (Du et al., 2022)
9.46 1:6::0.000 0.029
4.81 + 0.010
4.62 = 0.006
1.28 = 0.001
EGNO et al., 2024)
9.44 0.122
1 0.248
= 0.004
128 = 0.001
HEGNN (Cen al., 2024)
9.38 0.056
5.37 + 0.620
4,64 0.001
1.29 0.002
GF-NODE (Sun et al., 2025)
9.15 + 0.030
6.23 + 0.602
4.63 + 0.001
1.28 + 0.000
PAINET
8.98 0.040
4.65 + 0.022
4.31 0.004
1.26 = 0.001
Model
1)
naphthalene (x10->)
salicylic (x10-))
toluene (x10-1)
uracil (x10~>)
Linear (Satorras al., 2021)
6.26 0.000
13.8 + 0.000
1,24 + 0.000
9.68 : 0.000
MPNN (Gilmer et al., 2017)
4.67 +: 0.036
9.41 + 0.023
1.06 + 0.042
6.35 = 0.028
RF (Kohler et al.,, 2019)
3.98 0.009
12.4 : 0.105
1.09 = 0.001
6.21 : 0.050
EGNN (Satorras
2021)
+
10.4 0.154
0.015
5.84 = 0.082
ClofNet (Du et al., 2022)
3,90 + 01066
10.3 + 0,581
1.03 + 0.002
6.21 + 0.415
EGNO et al.,, 2024)
3,80 + 0.077
8.99 + 0.641
1,06 = 0.005
5.52 + 0.149
HEGNN (Cen al., 2024)
3,72 +
8.37 + 0.021
1.05 0.004
5,30 0.047
GF-NODE (Sun et al., 2025)
3,73 1 0.166
8.61 1 0.254
1.04 - 0.002
5.77 - 0.017
PAINET
3,36 0.008
7.80 : 0.007
1.00 0.000
0.004
The results are reported in Table for S2S and S2T, respectively. Across the two molecular systems, PAINET achieves the best performance under all settings, outperforming strong baselines. These improvements demonstrate the effectiveness of our all-pait interaction modeling larger and structurally diverse molecules.
Results.
4.4 PROTEINS DYNAMICS
Dataset and implementation. The Adenylate Kinase (Adk) Equilibrium dataset (Seyler & Beck- stein, 2017) is time--i--ec protein dynamics simulation of apo Adk as common benchmark for testing long-range structured modeling and scalability to large graphs. We randomly split the dataset into
with the ratio of 3:1:1. To assess the generalization capability of the model, we train all the models with one-step prediction (T 1) and test them with multi-step prediction (T = 5). More details are illustrated in Appendix F.5.
reported in Table 5 with test MSE of each time step and the overall A- MSE. Compared to all competitors, PAINET achieves the best performance across all intermediate prediction steps (from 1 = 5). We also find that the performance improvement by PAINET enlarges as time step increases, which shows that PAINET possesses better capability for modeling the dynamics mechanism that generalizes to simulations in long time horizons.
Results. The results
4.5 ABLATION STUDIES
We next conduct a series of ablation studies to analyze the effectiveness of the main components, the impact of key hyperparameters well as evaluations on the model's scalability.
8

---

Published as a conference paper at ICLR 2026
Table 3: Test A-MSE (1) for S2T (T' on MD17.
S2T S)
Model
aspirin (x10 2)
benzene (x
ethanol (x10 2)
malonaldehyde (x
EGNO (Xu al., 2024)
9.34 - 0.049
5.87 = 0.182
4,63 4: 0.002
1.28 4: 0.001
HEGGNN (Cen et al,, 2024)
9.54 + 0.024
7.02 + 0.233
4.67 + 0.001
1.29 + 0.001
GF-NODE (Sun et al., 2025)
9,25 0.017
5.48 + 0.228
+ 0.006
1.28 + 0.000
PAINET
8.84 0.045
4.87 0.295
4.55 : 0.001
1.26 : 0.001
82T (?
Model
naphthalene (x10
salicylic (x10-9)
tolueme (x10-1)
uradi (x10-9)
EGNO (Xu al., 2024)
3.95 + 0.055
8.51 = 0.070
1.06 + 0.009
5.68 = 0.167
HEGNN (Cen et al., 2024)
4,27 0.001
8.99 + 0.024
1.06 + 0.007
6.08 + 0.062
GF-NODE (Sun al., 2025)
0.065
8.61 0.009
0.002
5.82 : 0.008
PAINET
3.24 = 0.119
7,88 0.142
1.00 : 0.002
5.06 = 0.011
Table Test MSE (1) on MD22. F-MSE is reported for s2s (T 1) and A-MSE for S2T (T 5). We mark the test performance our model and emphasize the best and the runner-up in bold and underlined.
stackyose
Ac-Ala3-NHMe
Model
F.MSF (x10- (X(2S)
A-MSE (x 10110000XX00 (x2m)
F-MSE (x 10-))88)))
(x (0-
EGNN (Satorras al., 2021)
3.40 0.07
9.62 0.14
EGNO (Xu et al,, 2024)
3.33 0.05
4.43 0.11
9.11 = 0.17
0.12
HEGNN (Cen et al., 2024)
3,34 0.14
388 + 0.20
12.41 + 0.53
11.63 = 1.07
GF-NODE (Sun al., 2025)
3.11 0.27
2.54
8.52
PAINET
3,00 = 0.06
2,40 - 0.11
8.83 = 0.18
8.35 : 0.27
Table 5: Test MSE (4), inference GPU memory cost and inference time cost under multi-step pre- diction on Adk protein dynamics. All models are trained for single-step prediction and tested for five-step prediction. We report the MSE at time step t 1, ..... 5 and the overall A-MSE.
Model
52I (T' 5)
MSE MSE
MSE((
MSE(1 4)
MSE
A-MSE
GPU (GB)
Time
HUND (Xu al., 2024)
1.119
1.620
1.905
2.139
2.240
1.805
5.87
HEGNN (Con et al., 2021)
1.101
1.61
1.875
1.998
2.087
1.735
8.24
18.22
GF-NODE (Sun al., 2025)
1.055
1.642
1,913
2.047
2.113
1.762
9.91
27.71
Learnable pairwise mappings outperform fixed counterparts. We compare with multiple vari- ants of PAINET in Fig. 3a where replace learnable V and + with fixed counterparts (e.g., V using an all-one matrix). The results show that using learnable pairwise mappings yields the optimal A-MSE. In Fig. 3a, we also compare with simplified version of PAINET that entirely removes the attention module. This simplified model leads to clear performance degradation. These results altogether validate the necessity and superiority of our proposed attention network.
Parallel equivariant decoder balances accuracy and efficiency. To assess the efficacy of the cquivariant decoder, we substitute it (i.e., the decoder module in Eqn. 2) with other alternatives and compare with PAINET in Fig. 3b. 1) MLP-add: we add (() and (*) and then use a MLP to predict (1), 2) MLP-concat: we concatenate X (0) and (()) and then use a MLP to predict (1), 3) EGNN-recurrent: we change FGNN in our model to a recurrent counterpart that generates X iteratively from = 1
T. The results show that while using MIP more efficient than EGNN, it suffers from degraded performance. This is intuitive since MLP fails to accommodate the observed structural information. Also, compared with EGNN-recurrent, our model yields lower A-MSE with much fewer time costs, which shows the advantage of the parallel decoder in terms of both effectiveness and efficiency.
Moderate decoder depth balances accuracy and efficiency. We vary the number of decoding layers (i.e., the number of EGNN layers) from to 10 while keeping all the other factors the same and report the results in Fig. We find that as the number of decoder layers increases, the test F. MSE first decreases and then remains constant level, while the training time cost linearly increases. This suggests that using a moderate number of decoder layers (e.g., 3) yields the optimal prediction results and in the meanwhile brings up desired efficiency.

---

Published as a conference paper at ICLR 2026
1.8
9.0
A-MSE
-..
A-MSE
0.6
7.5
Time
1.2
A-MSE
6.0
0.4
4.5
0.9
original
0=00 q m
y
wi attention!
local attontion
scaled quadratic
linear
3.0
1.5
onginal
MLP-
MLP-
EGNN-
add
concat
recurent
(a) Comparison with attention variants.
(b) Comparison with decoder variants.
Figure Ablation studies learnable pairwise mappings in the attention network and the parallel equivariant decoder on Motion Capture Run. In (a), we group the model variants for ablation studies into three categories: ablation on and v (in blue), 2) ablation on all all-pair attention (in green), and 3) ablation on potential function p (in red).
Table 6: Test MSE (1) for different time steps Motion Capture. F-MSE is reported for T and A-MSE for = 5, 10, 15, 20.
Walk (x10-1)
Model
T=1
T=5
T=10
T=15
T=20
EGNO (Xu et al., 2024)
1.24 + 0.12
1,28 + 0.06
1,70 + 0.09
7.07
+ 0.32
7.19 1 0.22
HEGNN (Cen et al., 2024)
0.83 + 0.04
1.42 1 0.13
4.15 1 0.21
8.41 - 0.26
14.3 - 0.47
GE-NODE (Sun et al., 2025)
1.49 + 0.10
1.62 + 0.08
1.47 + 0.09
2.22
+ 0.26
1.32 + 0.17
PAINET
0.85 + 0.02
0.86 + 0.03
1.20 + 0.05
1.33 + 0.07
1.36 + 0.04
Run(x10-1)
Model
T=1
T=5
T=10
T=15
T=20
EGNO (Xu et al., 2024)
4.16
0.33
5.89 + 0.52
5.31 + 0.41
5.88
0.48
6.79 0.62
HEGNN (Cen al., 2024)
5.33 + 0.74
9.49 + 0.89
6.92
118 9.21
9.72
GE-NODE (Sun et al., 2025)
3.72 + 0.17
7.49 + 0.22
4.44 + 0.30
6.34 0.41
6.62 0.21
PAINET
3.50 1 0.27
3.33 1 0.12
4.34 1 0.18
5.42 - 0.27
5.96 0.17
One attention layer per time step suffices. We next study the impact of attention layers per time step and change its number from to 6 in Fig. 6. We find that the optimal performance is achieved by using one attention layer per time step, in which case the time cost is minimal in contrast with the model using multiple attention layers per time step. The indicates that one-layer attention is expressive enough for capturing all-pair interactions among arbitrary particles and compared the multi-layer counterpart, this simpler model is easier for optimization and less prone to over-fitting.
Robustness persists across time steps. To assess robustness, we further evaluate prediction across various time steps. As shown Table PAINET consistently performs best on the Motion Capture dataset, demonstrating strong stability in long trajectory prediction.
Computation costs scale nearly linearly w.r.t. particle numbers and time steps. To test the scalability of PAINET, vary the number of time steps and particles Proteins (Adk) and report the inference GPU memory cost and time cost in Fig. 7. We find that the GPU memory usage exhibits well particle numbers. Moreover,
linear increasing trend w.r.t. time steps for prediction Table 9 compares training and inference time across models, where PAINET achieves competitive efficiency. This demonstrates the desired scalability of PAINET large-scale multi-body systems.
5 CONCLUSION
We have proposed PAINET, physics-inspired equivariant architecture, for modeling 3D dynamics. PAINET consists of a principled attention layer grounded on the principle energy minimization and a parallel equivariant decoder to maintain SE(3))equivariance and inference efficiency. This architecture addresses the key limitation in existing graph-base based dynamics models, which often rely solely on local observed structures. Various experiments across diverse domains, including motion capture, molecular dynamics, and protein dynamics, demonstrate that our model consistently outperforms prior models across various benchmarks with desired efficiency.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
The research was supported by Shanghai Artificial Intelligence Laboratory. Additionally, We thank the open-source community for releasing code and datasets that made this research possible.
REFERENCES
Philip W Anderson. Plasmons, gauge invariance, and mass. Physical Review, 130(1):439) 1963.
Peter W Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. In Advances in Neural Information Processing Systems, volume 29, 2016.
Michael M Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geomet- ric deep learning: going beyond euclidean data. IEEE Signal Processing Magazine, 34(4):18 2017.
Aoneng Cao. The last secret protein folding: the real relationship between long-range interactions and local structures. The Protein Journal, 39(5):422 433, 2020.
Jiacheng Cen, Anyi Li, Ning Lin, Yuxiang Ren, Zihe Wang, and Wenbing Huang. Are high-degree representations really unnecessary in equivariant graph neural networks? Advances in Neural Information Processing Systems, 37:26238-26266 2024.
Stefan Chmiela, Alexandre Tkatchenko, Huziel E Sauceda, Igor Poltavsky, Kristof Schütt, and Klaus-Robert Müller. Machine learning of accurate energy-conserving molecular force fields.
Science advances, 3(5)x1660315,, 2017.
Stefan Chmiela, Valentin Vassilev-Galindo, Oliver T Unke, Adil Kabylda, Huziel Sauceda, Alexandre Tkatchenko, and Klaus-Robert Müller. Accurate global machine learning force fields for molecules with hundreds of atoms. Science Advances, 9())ccad(0873, 2023.
CMU Graphics Lab. Cmu motion capture database. https://mocas Accessed: 2025-05-09.
cmu, edu, 2003.
Nicolò Defenu, Tobias Donner, Tommaso Macri, Guido Pagano, Stefano Ruffo, and Andrea Trom- bettoni. Long-range interacting quantum systems. Reviews of Modern Physics, 95(3):035002, 2023.
Weitao Du, He Zhang, Yuanqi Du, Qi Meng, Wei Chen, Nanning Zheng, Bin Shao, and Tie- Yan Liu. Se (3) equivariant graph neural networks with complete local frames. In International Conference Machine Learning, pp. 5583-5608. PMLR, 2022.
PA Fedders and DA Drabold. Molecular-dynamics investigations of conformational fluctuations and low-energy vibrational excitations a-si: H. Physical Review B, 53(7):3841, 1996.
Roger H French, V Adrian Parsegian, Rudolf Podgornik, Rick F Rajter, Anand Jagota, Jian Luo, Dilip Asthagiri, Manoj K Chaudhury, Yet-ming Chiang, Steve Granick, et al. Long range interac- tions in nanoscale science. Reviews of Modern Physics, 82(2):1887-1944, 2010.
Fabian Fuchs, Daniel Worrall, Volker Fischer, and Max Welling. Se (3)-transformers: 3d roto- translation equivariant attention networks. Advances in neural information processing systems, 18:1970-19811 2020.
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning. pp- 1263-1272. PMLR, 2017.
Vitaly L Ginzburg, Vitaly Lazarevich Ginzburg, and LD Landau. On the theory of superconductivitt. Springer, 2009.
11

---

Published as a conference paper at ICLR 2026
Richard Gowers, Max Linke, Jonathan Barnoud, Tyler John Edward Reddy, Manuel N Melo, Sean L Seyler. Jan Domanski, David L Dotson, Sébastien Buchoux, lan M Kenney, et al. Md- analysis: python package for the rapid analysis of molecular dynamics simulations. Technical report, Los Alamos National I Laboratory (LANL., Los Alamos, NM (United States), 2019.
Jiaqi Han, Wenbing Huang, Hengbo Ma, Jiachen Li, Josh Tenenbaum, and Chuang Gan. Learning physical dynamics with subequivariant graph neural networks. Advances in Neural Information Processing Systems, 35:26256-26266 2022.
Jiaqi Han, Minkai Xu, Aaron Lou, Haotian Ye, and Stefano Ermon. Geometric trajectory diffu- sion models. Advances in Neural Information Processing Systems, 2024. URL https: ///p//o///////// /ooenzeeieww net/pdff id=00mmm 5Mv 9H.
Jiaqi Han, Jiacheng Cen, Liming Wu, Zongzhao Li, Xiangzhe Kong, Rui Jiao, Ziyang Yu, Tingyang Xu, Fandi Wu, Zihe Wang, al. A survey of geometric graph neural networks: Data struc- tures, models and applications. Frontiers of Computer Science, 19:1911375, 2025. doi: 10.1007/
$11004-255-4446. URL https:://// / /oii wrw.ol........oo...ooooo 10100 020---025--0000 41426-14466
Peter W Higgs. Broken symmetries and the masses of gauge bosons. Physical review letters, 13 (16):508, 1964.
Pierre C Hohenberg and Alexei P Krekhov. An introduction the ginzburg landau theory of phase transitions and nonequilibrium patterns. Physics Reports, 572:1-42, 2015.
Wenbing Huang, Jiaqi Han, Yu Rong, Tingyang Xu, Fuchun Sun, and Junzhou Huang. Equivariant graph mechanies networks with constraints. In International Conference Learning Represen- tations, 2022. URL https::tttt/t...t. http::////peeeveeeeeeeeeeeeeeeee
n So-rummii--SSnnneeeeeeeee
Zijie Huang, Wanjia Zhao, Jingdong Gao, Ziniu Hu, Xiao Luo, Yadi Cao, Yuanzhou Chen, Yizhou Sun, and Wei Wang. Physics-infe inyfore-informed regularization for Nomain-agnostee dynamical system modeling. In Advances in Neural Information Processing Systems, 2024. URL https:// openrevies ew. wet//pfffdd// qoNE8P7.
Bo Jiang, Ziyan Zhang. Doudou Lin, Jin Tang, and Bin Luo. Semi-supervised learning with graph networks. In Proceedings of the IEEE/CVF conference computer vision and patlern recognition, pp. 11313-11320, 2019.
George Em Karniadakis, Ioannis Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informat machine learning. Nature Reviews Physics, 3(6)422 440, 2021.
Thomas Kipf, Ethan Fetaya, Kuan-Chich Wang, Max Welling, and Richard Zemel. Neural relational inference for interacting systems. In International Conference on Machine Learning, pp. 2688- 2697. PMLR, 2018.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. In International Conference on Learning Representations (ICLR), 2017. URI. https: / //oppneeview..eoe
arXiv:1609.02907.
Jonas Köhler, Leon Klein, and Frank Noé. Equivariant flows: sampling configurations for multi- body systems with symmetric energies. arXiv preprint 191100191. 2019.
Lev Davidovich Landau et al. On the theory of phase transitions. Zh. eksp. teor. Fiz, (11-32):926, 1937.
Yi Liu, Limei Wang, Meng Liu, Yuchao Lin, Xuan Zhang, Bora Oztekin, and Shuiwang Ji. Spherical message passing for 3d molecular graphs. In International Conference Learning Representa- tions (ICLR), 2022.
Fritz London. The general theory of molecular forces. Transactions of the Faraday Society, 33: 8b-26, 1937.
Yao Ma, Xiaorui Liu, Tong Zhao, Yozen Liu, Jiliang Tang. and Neil Shah. A unified view graph neural networks as graph signal denoising. In Proceedings of the 30th ACM international conference on information & knowledge management, pp. 10p.1200 1202-1211, 2021.
12

---

Published as a conference paper at ICLR 2026
Alexander D MacKerell Jr, Nilesh Banavali, and Nicolas Foloppe. Development and current status of the charmm force field for nucleic acids. Biopolymers: original Research on biomolecules, 56 (4):257-265, 2000.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li F Fei-Fei, Josh Tenenbaum, and Daniel Yamins. Flexible neural representation for physics prediction. Advances neural information processing systems, 31, 2018.
John P Perdew, Kieron Burke, and Matthias Ernzerhof. Generalized gradient approximation made simple. Physical review letters, 77(18)3886, 1996.
Tobias Pfaff, Meire Fortunato, Alvaro Sanchez-
Sanchez-Gonzalee Gonzalez, and Peter W Battaglia. Learning mesh- based simulation with graph networks. In International Conference on Learning Representations (ICIR), 2021.
Frederic Richards. Areas, volumes, packing, and protein structure. Annual review of biophysies and bioengineering, 6(1):151-176, 1977.
R. T Rockafellar. Convex analysis. Princeton University Press, 1970.
Celeste Sagui and Thomas A Darden. Molecular dynamics simulations of biomolecules: long-range
electrostatic effects. Annual review of biophysics and biomolecular structure, 28(1):155-179, 1999.
Alvaro Sanzaez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter Battaglia. Hamiltonian graph networks with ode integrators. In Advances in Neural Information Processing Systems, volume 32, 2019.
Alvaro Sanehez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W Battaglia. Learning simulate complex physics with graph networks. In International Confer- ence on Machine Learning, 2020.
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) equivariant graph neural net- works. In International Conference on Machine Learning, pp. 9123-93322 PMLR, 2021.
Gideon Schreiber and Alan R Fersht. Rapid, electrostatically assisted association of proteins. Nature structural biology, 365):427-431, 1996.
Kristof Schütt, Pieter-Jan Kindermans, Huziel Enoc Sauceda Felix, Stefan Chmiela, Alexandre Tkatchenko, and Klaus Robert Müller. Schnet: A continuous-flle fillee convolutional neural network for modeling quantum interactions. Advances in neural information processing systems, 30, 2017.
Sean Seyler and Oliver Beckstein. Molecular dynamics trajectory for benchmarking mdanalysis. URL: https://igshare. com//rticles/Molecuaar dynamics trajectory for benchmarking MDAnaly- sis/5 70, doi, 10(m9):7, 2017.
Chence Shi, Shitong Luo, Minkai Xu, and Jian Tang. Learning gradient fields for molecular con- formation generation. In International conference on machine learning, pp. 9558-95666 PMLR, 2021.
Dai Shi, Lequan Lin, Andi Han, Zhiyong Wang, Yi Guo, and Junbin Gao. When graph neural net- works meet dynamic mode decomposition, 2025. URL id=duGygkA30R.
Fang Sun, Zijie Huang. Haixin Wang, Huacong Tang, Xiao Luo, Wei Wang, and Yizhou Sun. Graph fourier neural ODEs: Modeling spatial-temporal multi-scales in molecular dynamics. Transac- tions on Machine Learning Research, 2025. ISSN 2835-8856. URL https: tutttss://ppereeee //tp////ee/eee.eee//e/eeee net/ pet/foormm wecrrummii-xeee xd= 2F
Aaron Taudt, Axel Arnold, and Jürgen Pleiss. Simulation of protein association: Kinetic pathways towards crystal contacts. Physical Review E, 91(3):033311, 2015.
Fabian L. Thiemann, Thiago Reschützegger, Massimiliano Esposito. Tseden Taddese, Juan D. Olarte-Plata Plata, and Fausto Martelli. Force-free molecular dynamics through autoregressive equiv- ariant networks, 2025. URL httpt:ttt/// /xwxxvivvvv www.g//a///aaaee 23794.
13

---

Published as a conference paper at ICLR 2026
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation-and -equaviiinnnivviaiatt neural networks for 3d point clouds. arXiv preprint arXiv:/802.082/15 2018.
Benjamin Ummenhofer, Lukas Prantl, Nils Thuerey, and Vladien Koltun. Lagrangian fluid sim- ulation with continuous convolutions. In International conference on learning representations, 2019.
Oliver Unke, Stefan Chmiela, Huziel Sauceda, Michael Gastegger, Igor Poltavsky, Kristof T Schutt, Alexandre Tkatchenko, and Klaus-Robert Mu"'ler. Machine learning force fields. Chem- ical Reviews, 121((66:0044--00186 2021.
Hrishikesh Viswanath, Yue Chang, Julius Berner, Peter Yichen Chen, and Aniket Bera. Reduced- order neural operators: Learning lagrangian dynamics on highly sparse graphs, 2024. URL
htttss://arxin org/ wtt//as/////ooooooooooooo/..a.aa 2407 03925.
Felix Wu, Amauri Souza, Tianyi Zhang, Christopher Fifty, Tao Yu, and Kilian Weinberger. Sim- plifying graph convolutional networks. In International conference on machine learning. pp- 6861-6871. Pmlr, 2019.
Hao Wu, Fan Xu, Yifan Duan, Ziwei Niu, Weiyan Wang, Gaofeng Lu, Kun Wang, Yuxuan Liang, and Yang Wang. Opatio-temporal fluid dynamics modeling via physical-wwareness and parameter diffusion guidance, 2024. URL https: /ttpivaavvvvvvvv///i//e../aaaavv org/ abs 2403.4. 13850.
Kai Xu, Jianmin Wang, Mingquan Liu, Kewei Zhou, Shaolong Lin, Weihong Li, Lin Shi, Peng Zhou, Huanxiang Liu, and Xiaojun Yao. Efficient generation of protein and protein-protedd com- plex dynamies via (3)-parameterized diffusion models. Journal of Chemical Information and Modeling, 2025.
Minkai Xu, Lantao Yu, Yang Song, Chence Shi, Stefano Ermon, and Jian Tang. Geodiff: A geometric diffusion model for molecular conformation generation. In International Confer- ence Learning Representations, 2022. URL https: httts::///opeeeeee http::///ppereeieeeeee net nwt/forumiidi PZcvxEMZVQC.
Minkai Xu, Jiaqi Han, Aaron Lou, Jean Kossaifi, Arvind Ramanathan, Kamyar Azizzadenesheli, Jure Leskovee, Stefano Ermon, and Anima Anandkumar. Equivariant graph neural operator for modeling 3d dynamics. In Proceedings of the 41st International Conference on Machine Learn- ing, ICML'24, JMLR.org. 2024.
Yongyi Yang, Tang Liu, Yangkun Wang, Jinjing Zhou, Quan Gan, Zhewei Wei, Zheng Zhang, Zengfeng Huang, and David Wipf. Graph neural networks inspired by classical iterative algo- rithms. In International Conference on Machine Learning, pp. 11773-11783, 2021.
Bangchen Yin, Jiaao Wang, Weitao Du, Pengbo Wang, Penghua Ying. Haojun Jia, Zisheng Zhang, Yuanqi Du, Carla P Gomes, Chenru Duan, al. Alphanet: Scaling up local frame-based atomistic foundation model. arXiv e-prints, pp. arXiv-25011 2025.
Yuelin Zhang, Jiacheng Cen, Jiaqi Han, Zhiqiang Zhang, Jun Zhou, and Wenbing Huang. Improving equivariant graph neural networks on large geometric graphs via virtual nodes learning. In Forty- first International Conference on Machine Learning, 2024.
Dengyong Zhou, Olivier Bousquet, Thomas Navin Lal, Jason Weston, and Bernhard Schölkopf. Learning with local and global consistency. In Advances in Neural Information Processing Sys- tems, pp. 321-328, 2004.
14

---

Published as a conference paper at ICLR 2026
A RELATED WORKS
We briefly discuss the relevant literature to include more background.
3D Dynamics Prediction. Predicting 3D dynamics-such as particle motion, robotic trajectories, or molecular interactions-is a fundational challenge in physics simulation and roboties. Traditional physics- models offer interpretability but often fall short when modeling complex or unknown interactions. More recently, data-driven approaches, especially deep learning, have shown strong potential (Karniadakis et al., 2021; Wu et al., 2024; Huang et al., 2024). While RNNs, LSTMs, and Transformers can model temporal dependencies, they often fail to capture the underlying relational structure in multi-agent or multi-body systems (Kipf et al., 2018). Graph Neural Networks (GNNs) address this by modeling entities nodes and their interactions as edges (Battaglia et al., 2016; Han et al., 2025). Recent advances in Equivariant GNNs (EGNNs) (Satorras et al., 2021; Huang al., 2022; Han al., 2022) further improve physical fidelity by preserving Euclidean symmetries, enabling more accurate and generalizable 3D dynamics predictions. Some subsequent works focus on the efficiency of EGNNs through architectural optimizations (Zhang al., 2024).
Message Passing Neural Networks. Graph neural networks (GNNs) (Gilmer et al., 2017; Kipf & Welling, 2017) provide a general framework for learning over relational structures and have rec- ognized as powerful method for simulating physical systems. One line of related works har- ness physical priors to design more expressive message-passing operators to encode system interac- tions (Mrowca et al., 2018; Shi et al., 2025; Viswanath et al., 2024) incorporate classical physical mechanics into the architecture (Sanchez -aabbeez- et al., 2019). A parallel line of works focus on encoding the symmetries of Euclidean space, i.e., equivariance, as inductive bias into the architec- tures, including translation equivariance (Ummenhofer et al., 2019; Sanchez-Gonzalez et al., 2020; Pfaff et al., 2021), full rotation and reflection equivariance through spherical harmonics (Thomas et al., 2018; Fuchs et al., 2020) or equivariant message passing (Satorras et al., 2021; Huang et al., 2022: Thiemann al., 2025). Further refinements exploit local coordinate frames to process higher- order geometric features (Liu et al., 2022; Du et al., 2022; Han et al., 2024; Cen et al., 2024) or model the temporal information by Fourier transform (Xu al., 2024) while preserving equivari- ance. Recently, a neural-ODE (Sun et al., 2025) formulation to enhance expressivity. Despite these advances, these EGNNs are constrained on the observed structure and neglect to incorporate the latent all-pair interactions, which is crucial in physics.
B EQUIVARIANCE
In this work, we consider equivariance properties with respect to the special Euclidean group SE(3), which includes translations. rotations and permutations. For 3D dynamics, the three types of equiv- ariance considered are summarized below:
Table 7: Three types of equivariance and corresponding conditions.
Type
Transformation
Condition
Translation
3
#(X + g) = =X) +
Rotation
Q e R3x3
((QX) = Qu((X
Permutation
P Rexn
((PX) = PPμ(X)
Here X = : :x]] E Nx3 denotes a set of position vectors for N nodes in 3D space, and N3 + N3 is function on X. As illustrated in Fig. 4, translation equivariance implies that a uniform shift of all input positions by a vector g results in the output being shifted by the same amount. Rotation equivariance indicates that applying an orthogonal transformation Q to the input causes the output be rotated in the same way. Permutation equivariance means that reordering the input points leads to an identically permuted output.
Note that for velocity vectors :vf::..; v E Nx3 and & : Nx3 - R Nx3 translations do not alter their values. However, they transform equivariantly under rotations and permutations, i.e.,
QE(V) EQQV), Pε(V) ((PV).
15

---

Published as a conference paper at ICLR 2026
Original
Rotation
Translation
Permutation
QX
X +
PX
p(-)
Q=(X)
J(X) +
Pμ(()
Figure Illustration for three types of equivariance, including rotation, translation and permutation.
These equivariance properties essential in modeling systems that comply with spatial symme- tries, such as molecular dynamics or physical simulations.
C ENERGY MINIMIZATION ON GRAPHS
Graph representation learning can be rigorously formulated an optimization problem governed by an energy minimization framework, where node embeddings evolve to satisfy smoothness and con- sistency criteria. Modern Graph Neural Networks (GNNs) can be interpreted as iterative solvers for this energy function. Foundational approaches establish this through the minimization of quadratic energy functionals to enforce local and global consistency on the graph manifold. has been shown that the propagation rules in standard architectures, such GCN (Jiang et al., 2019) and SGC (Wu et al., 2019), effectively perform gradient descent steps on an objective energy functionals with a fitting term (Ma et al., 2021). Extending this theoretical framework modeling 3D dynamics with latent structures, we characterize the uncovering of unobserved interactions as the minimization of regularized energy function. Formally, this energy E layer t balances feature fidelity with structural smoothness:
E(H, t: ((t)))) -16 +
12
> Pis( h,113), 1J
(12)
where H represents the stack of particle embeddings, Pij denotes non-linear potential function, and is regularization weight. The first term serves as a fidelity constraint, regularizing the deviation of updated embeddings from their previous states, while the second term penalizes pairwise distances to enforce internal consistency among particles. This formulation draws inspiration from physi- cal systems where state evolution naturally transitions from high-energy to low-energy regimes to reach equilibrium. By instantiating Pij as a handau-Ginzbur potential, this framework connects the optimization process phenomenological explanations of matter formation and symmetry break- ing. Consequently, the minimization of this energy yields a principled attention mechanism where propagation steps correspond rigorous gradient descent updates.
D PROOF FOR THEOREM
We first convert the minimization of Eqn. 3 into the minimization of its variational upper bound. For any concave, non-decreasing function p + - IR, one can express it as the variational decompo- sition
p(z
minlwz² - x(w)) 2 wz² -(ww) 220
(13)
where w is variational parameter and is the concave conjugate of p. Eqn. 13 defines p(=²) the minimal envelope of a series of quadratic bounds wz p(w) defined by a different values of 2 0. The upper bound is given for fixed w when removing the minimization operator. We note that for any optimal w we have
w" 22 plw")) p(=2),
(14)
16

---

Published as a conference paper at ICLR 2026
which tangent to p at =² and w" the equality of Eqn. 13.
06(=") 8g
This gives the sufficient and necessary condition for
Based on the above result, we can derive the variational upper bound of Eqn.
E(H, 4 (wig). (1t))())))))))))))))))))) = Elh, - 1
+
Σ
waj h: h ilwsss)
(15)
f
where f is the concave conjugate of and the equality holds if and only if the variational parameters satisfy
Op.y (h?)
Wij
(16)
Oh
/= hy
In light of this relationship, we can minimize the upper bound surrogate Eqn. 15 which is equivalent to the minimization of Eqn. 3 on condition that the variational parameters way 's given by Eqn. 16. Then, consider a gradient decent step on Eqn. 15, the particle embeddings are updated through (assuming as the step size)
H = H(O OE(H, H, t; (
0H
H-HO
(17)
= HH (X(D()) 2
H
H
H())
H TDD 2 H
where (+) (0) JNNN and ((t) denotes the diagonal degree matrix associated with (
Common practice to accelerate convergence adopts positive definite preconditioner term, e.g.
(D())-1 to re-seale the updating gradient and the final updating form becomes
H (1 H +
(18)
The above iteration converges for step size 0 < < 2 We thus prove that
(*) f) (Pig)) 2 (++ ,t; (pi)))-
(19)
Besides, we notice that for a fixed H, EHH, 4; {p:r = H - + Pir
hy h)1
becomes function of t and its optimum is achieved if and only if H H. Such a fact yields that
(H t; (po))) 2 EH 10; (p:r H)
(20)
The result of the main theorem follows by noting that E(H('), ,t (pist) 2 (((++)) t; (p5s)) 2 ((H(+1), + 1; (pi))).
E PROOF FOR EQUIVARIANCE OVER PAINET
We provide a formal proof that our model architecture preserves SE(3)-cqui throughout trajectory prediction process. Let Q E SO(3) be rotation matrix and R3 translation vector. function f(-) is said to be equivariant to SE(3) transformations if, for all positions X and all transformations (Q, g). it holds that:
SQQ + g) = Qx +g.
(21)
Similarly, a function said to be invariant
f(QX - g) f(X..
(22)
Our model consists of two major components: (1) physics-inspired attention network that oper- ates on particle embeddings h." (3)s,, 's and (2) an equivariant decoder that updates positions x(" using
(t) and relative positions Xj. We can prove that, the atten- embeddings hn and the decoder equivvaiat with
messages aggregated both h tion network outputs respect its positional inputs.
17

---

Published as a conference paper at ICLR 2026
Invariance of embeddings. At each time step, the attention network operates merely latent embeddings h(")s,, which are independent of the input coordinates x;. The attention mechanism computes dot-product similarity scores (e.g., h; h2), and produces updates with adaptive pairwise mappings. Since h, 8s are defined in latent space and do not transform under SE(3), their updates are invariant. For instance, as for rotation operation, we have:
vQ SO(3),
h (() Artention( QX Attention (F) Attention (F; x
(23)
where is an orthogonal rotation matrix and F is the initial node features.
Equivariance of decoder. The decoder is composed of equivariant GNN layers. Each EGNN layer updates positions using the form:
11) m 2m (h. .b.
(1) (1)
||x
10,
any),
h
(++1)
Φb
h
(1) Σ
mij 1j
(1)
(24)
JEN(S)
x x;
(1+1)
(1)
+
(x
!1)
x
11
(m))) ())
JEN()
where Oz Pm, Uh are MLPs. For any orthogonal rotation transformation Q SO(3) and g R3,
we can prove that ||x (1) 102 is invariant, since:
[1(xx] (1) g) -Qxx)) ( +g]l1² |[Qx)" (1) Qx 1
(xx" x, Q( xx
(25)
=xx/ 1p.
Thus, my is invariant. Moreover, we have:
Qx (++)
Q(x!!)
+ Σ (x(1)
x x (m,") +g
+
JEN())
(Qx)
g + Q Σ (x!" -x . (m
JEN(1)
(26)
(Qx (1)
+ g + Σ
(Qx) - -Qx) (0) oz (m
JEN(i)
(Qx)
+g) + Σ ([Qx)) (1) +g) -(xx)) g)) dx (m
JEN(()
Since my is invariant and the position updates are equivariant with respect to SE(3) transformations, the overall model preserves equivariance by design:
X A PA (0) (QX g. H () + g: H)).
(H) H()),
(27)
F EXPERIMENT DETAILS
F.1 EVALUATION METRICS
All the experiments involve two tasks: state-to-state (S23) and torafe-t-ry (S2T). The $2s task directly predicts the final state, while the S2T task predicts the trajectories a specified number of future time steps. Correspondingly, 52S uses Final Mean Squared Error (F-MSE) to measure the
18

---

Published as a conference paper at ICLR 2026
Algorithm PAINET for 3D dynamics modeling and prediction
Require: Initial positions x initial features F; trajectory length decoder layer number hyperparam-
eter vj. adaptive pairwise mappings b, V.
Ensure: Position predictions for each time step (X(),
1: IT
W
2: for t=rt=0 0 T do
Q V -
HO
4: Compute the n.2-normaizeed query and key matrices @ (()))
5: D(++1) diag" (1 + (v (Ô ()1
6:
H(--)) (1 - H + ((++1)
(V()) + (V @ (Q (KO))))))))
7:
H(01:00 x (1+100) <- H (++) X
8:
for = t L do
((++1++1) &- EGNN ( ((t1) H(++))
10: end for
11: ((++) - (++11t)
12:
end for
13:
return (X(1) {(7)
MSE between the predicted final state and the ground truth, and S2T calculates Average MSE (A- MSE) to assess the MSE averaged across all discretized time steps along the decoded trajectory. Specifically, the definitions of two evaluation metrics are:
1
F-MSE
Σ (T) (T)11?
(28)
1 ([x(()) (x)1 TN
A-MSE
(29)
where i E V
N = [V (() denotes the predicted position for particle at time step 4, and x;())
denotes the corresponding ground truth.
F.2 MOTION CAPTURE
We two subsets from the CMU Motion Capture dataset (CMU Graphics Lab, 2003), which pro- vides high-fidelity 3D trajectories of human joint movements. Specifically, we follow the protocols adopted by Huang al. (2022); Han al. (2022), focusing Subject #35 (Walk) and Subject #9 (Run). The dataset is partitioned into training, validation, and test sets with 200/600/600 samples for Subject #35 and 200/240/240 for Subject #9, respectively. The trajectories are uniformly discretized by AT 30. Each snapshot in every trajectory consists of 3D coordinates of 31 human joints. Following prior work. joints are modeled as graph nodes, with connections representing physical or kinematic constraints.
The model is conditioned on the initial positions and velocities of joints, and trained to predict future joint positions over a 100-step horizon. Training is performed using the Adam optimizer with a learning rate of 10-3, minimizing the mean squared error over predicted trajectories. For both training and evaluation, we set T 5 in S2T tasks. Namely, for each initial position, the model needs to predict the positions in the next 5 time step.
F.3 MOLECULAR DYNAMICS ON SMALL MOLECULES (MD17)
We evaluate our model on the MD17 dataset (Chmiela et al., 2017), which provides initio molec- ular dynamics trajectories for several small organic molecules. For each molecule, we randomly partition the dataset into 500 training, 2000 validation, and 2000 test trajectories with designated temporal interval of
random seed for reproducibility. Following the common practice, we AT 3000 for uniform discretization and dispose of the first and last 10000 snapshots in case of instability. Each training sample consists of sequence of atomic positions and velocities at the
19

---

Published as a conference paper at ICLR 2026
initial time step, with the task of predicting atomic positions at future time steps. The initial velocity is computed from consecutive positional differences within the trajectory. Following common pre- processing conventions, we focus heavy atoms by excluding hydrogens from the graph. For the molecular graph structure, we follow prior works (Shi et al., 2021; Xu et al., 2022; 2024) and aug- ment the connectivity by adding 2-hop neighbors. Edge features include combination of atomic types, bond types, and hop distance between connected atoms.
For both training and evaluation, we set T in S2T tasks, Namely, for each initial position, the model needs to predict the positions in the next 8 time step.
F.4 MOLECULAR DYNAMICS ON LARGE MOLECULES (MD22)
MD22 (Chmiela et al., 2023) contains trajectories for significantly larger and more flexible sys- tems (ranging from 42 to 370 atoms), including supramolecules and major biomolecules like Ac- Ala3-NHMe (peptide), DHA (lipid), and Stachyose (carbohydrate), which further validates model effectiveness for predicting the dynamics of larger, more complex molecules. For each molecule, we randomly partition the dataset into 500 training, 2000 validation, and 2000 test trajectories with designated random seed for reproducibility. We use temporal interval of AT = 1000 for uniform discretization and dispose of the first and last 10000 snapshots in case of instability. Each training sample consists of sequence of atomic positions and velocities at the initial time step, with the task of predicting atomic positions at future time steps. The trajectories were sampled at temperatures between 400 and 500 K at a resolution of fs, and energies and forces were computed at the PBE+MBD (Perdew et al., 1996) level of theory.
For both training and evaluation, we set T model needs to predict the positions in the next 5 time step.
S2T tasks. Namely, for each initial position, the
F,5 PROTEIN DYNAMICS (ADK)
We evaluate our model on the Adk protein equilibrium trajectory dataset, originally introduced by Seyler & Beckstein (2017) and later preprocessed by Han al. (2022). The simulation captures the molecular dynamics of apo adenylate kinase (Adk) using the CHARMM27 force field (MacKerell Jr al., 2000), under NPT conditions at 300 K and bar. Water molecules and ions are modeled explicitly, and snapshots are recorded every 240 picoseconds, yielding a total trajectory length of approximately 1.004 microseconds. This dataset is integrated into the MDAnalysis toolkit (Gowers et al.. 2019), which facilitates parsing and processing of molecular simulation outputs. We adopt the same data split as MacKerell Jr al. (2000), with 2481 sub-trajectories for training, 827 for validation, and 878 for testing. The molecular graph is constructed from backbone atoms using 10 Å cutoff for edge definition. Each sample provides initial atomic positions and velocities as input, and the task is to forecast future atomic configurations.
For generalization, we set the model to train with one-step prediction (T= 1), evaluate with multi- step prediction (7 = 5). Namely, the model is trained to directly predict the final position, but evaluated on prediction of positions in 5 time steps.
F.6 HYPERPARAMETERS
Table 8 summarizes the key hyperparameters used for training PAINET across all datasets, includ- ing learning rate r),,, weight decay (weight decay), batch size (batch), number of decoder layers (layer), hidden size (hidden) and number of attention heads (num_heads). We use Adam optimizer and models are trained until convergence with early stopping based on validation loss, using a patience of 50 epochs. All runs use the same random seed to ensure reproducibility.
For Motion Capture and Protein datasets, we use relatively smaller batch sizes due to memory con- straints, while MD17 allows larger batches given its lower per-sample dimensionality. For other hy- perparameters, we search them the validation set with the searching space: 1r (1x 10 -004,,0 4, 5 x 10 10-4, 4 1 x 10 10-3, 3 2 x 10 33. weight decay E {1 x 10
10-10 10 x 10 10-12 2,, x 10 5, x 10 10-18) 18)
layer E (3,4, 4, 5, 6, 7). hidden E (32, 64, 128), num heads (2, 4)
20

---

Published as a conference paper at ICLR 2026
Table Summary of hyperparameter settings PAINET.
datasct
time stop
batch
1x
weight decay
Layez
hidden
nusLheada
Motion Capture (Walk)
12
10-1
10
128
Motion Capture (Walk)
12
5 x 10
10
128
3
Motion Capture (Rws)
12
5 x
X 10 15
128
2
Motion Capture (Rwn)
12
5 10
10-15
128
MD17 (aspirin)
100
x 10
1x 10-15
7
64
4
MD17 (bengene)
100
x
64
2
MD17 (ethanol)
100
5 10
10-
MD17 (malonaldekyde)
100
1 10-1
1 10-15
64
4
MD17
(nuphialleee
100
5 X 10
X
64
2
MD17 (salicylic)
MD1 17 (toluene)
100
5 X 10 X 10
1 10-15
64
2
MD17 (2017(20) (uracit)
100
X 10
64
4
MD17 (aspirin)
100
10
10
MD1 17 (benzene)
100
5 X 10
1 10
64
2
MD17
(ethanol)
100
5 X 10
X
64
4
MD17
(malonaldekyde)
100
10
10-15
4
MD17
(nophiiolene)
100
5 x 10
1 10-15
64
2
MD17
(salicylic)
100
5 X
1
3
MD17
(toluene)
100
5 10
10
2
MD17 (uracil)
100
1 x 10
1 10-15
64
4
Protein
5 x 4
10
32
2
Table 9: Training time per epoch and inference time Motion Capture.
Walk ($29)
Run (82S)
Model
Training Time (s)
Inference Time (8)
Training Time (8)
Inference Time (8)
Linear (Satorras et al., 2021)
0.02
0.07
0.03
0.03
RF (Köhler ct al., 2019)
0.12
0,20
0.11
0.08
MPNN (Gilmer ct al., 2017)
0.09
0.16
0.10
0.07
EGNN (Satorras et al., 2021)
0.19
0,26
0,22
0.14
EGNO (Xu al., 2024)
0.33
0.36
0,30
0.22
HEGNN (Cen et al., 2024)
0.42
0.52
0.24
GF-NODE (Sun et al., 2025)
0.83
0.51
PAINET
0.20
0.26
0.21
0.15
Walk (S2T)
Run (S2T)
Model
Training Time (s)
Inference Time (s)
Training Time (s)
Inference Time (s)
EGNO (Xu al., 2024)
0.18
0.23
0.19
0.17
HEGNN (Cen et al., 2024)
0.20
0.25
0.22
0.21
GF-NODE (Sun et al., 2025)
1.10
1.26
0.76
0.68
PAINET
0.28
0.33
0.31
0.26
F.7 COMPUTATION RESOURCES
All experiments are conducted on single NVIDIA H20 GPU with 4 MiB L2 cache and 96 GiB of memory. The system is configured with NVIDIA driver version 150.144.03 and CUDA toolkit version 12.4. To evaluate the scalability of our model across datasets of varying sizes and temporal lengths, we report GPU memory usage and time cost Fig. All measurements are obtained under identical hardware and software configurations, ensuring a fair comparison across tasks.
F.8 TIME CONSUMPTION
To evaluate the model's computational efficiency, we report both training and inference time on the Motion Capture experiments. As shown in Table 9, our model consumes comparable time costs as the powerful competitors EGNO, HEGNN and GF-NODE.
F.9 PHYSICALLY MEANINGFUL METRIC
While our primary evaluation employs the standard Mean Squared Error (MSE) metric which is commonly adopted in the literature such EGNO and GF-NODE, acknowledge that MSE alone may not fully capture the physical validity the simulation. To provide a more rigorous assessment of structural fidelity, supplement analysis with Root Mean Square Deviation (RMSD). Al-
21

---

Published as a conference paper at ICLR 2026
Table 10: Test RMSD ()) for s2S (I' 1) MD17.
Model
828
aspirin (x10 9)
ethanol (x10 9)
malonaldehyde (x
Linear (Satorrus ul., 2021)
3,46 0.000
4.09 4: 0.000
2.37 = 0.000
1,46 = 0.000
MPNN (Gilmer et al., 2017)
0.053
2.39 + 0,569
2.22 + 0.084
1.14 + 0.0004
RF (Kohler et al., 2019)
3,30 0.003
3.55 + 0.028
2.15 : 0.001
1.14 + 0.006
EGNN (Satorras
2021)
0.158
2.53 0.246
2.16 0.005
0.001
ClofNet (Du et ul., 2022)
+.0229
2.19 0.010
2.15 + 0.006
1.13 = 0.001
EGNO (Xu et al., 2024)
3.07 +: 0.122
2.49 + 0.248
2.15 + 0.004
1.13 0.001
HEGNN (Con al., 2024)
3.06 0.056
2.32 0.620
0.001
0.002
GF-NODE (Sun et al., 2025)
3.02 +: 0.030
2.50 + 0.602
2.15 = 0.001
1.13 = 0.000
PAINET
2.99 0.040
2.16 + 0.022
2.08 0.004
1.12 = 0.001
Model
naphthalene (x10-9)
salicylic (x10-9)
toluene (x10-1)
uracil (x10-9)
Linear (Satorrus ul., 2021)
2.50 + 0.000
3.71 + 0.000
1.11 + 0.000
3.11 = 0.000
MPNN (Gilmer et al., 2017)
2.16 + 0.036
3.07 + 0,023
1,03 + 0.042
2.52 + 0.028
RF (Kohler et al., 2019)
1.99 0.029
3,52 : 0,105
1.04 = 0.001
2.49 : 0.050
EGNN (Satorras
2021)
1,98 0.042
3.22 0.154
0.015
2,42 = 0.082
ClofNet (Du et al., 2022)
1.97 + 11066
3.21 + 0,581
1.01 + 0.002
2.49 + 0.415
EGNO et al.,, 2024)
1,95 0.077
3.00 + 0.641
1,08 0.005
2.35 = 0.
HEGNN al., 2024)
1.93 0.037
2.89 0.021
0.004
2.30 0.047
GF-NODE (Sun et al., 2025)
1.93 0.166
2.93 + 0.254
1.02 + 0.002
2.40 = 0.017
PAINET
1.83 0.008
2.79 + 0.007
1.00 0.000
2.27 0.004
Table 11: Test RMSD (L) for S2T (7 8) MD17.
82T (T 8)
Model
aspirin (x10-²)
benzene (x10-l)
ethanol (x10-))
malonaldehyde (x10-l)
(Xu a 2024)
0.001
HEGNN (Cen et al., 2024)
3.09 + 0.024
2.65 = 0.233
2.16 + 0.001
1.14 + 0.001
GF-NODE (Sun et al., 2025)
3.04 + 0.017
2,34 : 0.228
2.15 + 0.006
1.13 + 0.000
0.001
92T
Model
naphthalene 1010-))
salicylic (x10-5)
toluene (x10-l)
uracil (x10-5)
EGNO (Xu al., 2024)
1.99 + 0.055
2.92 = 0.070
1.03 + 0.009
2.38 + 0.167
HEGNN (Cen et al., 2024)
2.07 = 0.001
3,00 = 0.024
1.03 + 0.007
2.47 = 0.062
GF-NODE (Sun et al., 2025)
2.93 + 0.009
+: 0.002
2.41 + 0.008
PAINET
1.80 0.119
2.81 # 0.142
1.00 + 0.002
2.25 + 0.011
Test RMSD (1) under multi-step prediction on Adk protein dynamics, All models are trained for single-step prediction and tested for five-step prediction.
Table 12:
82T (T 5)
Model
RMSD( 1)
RMSD(!
RMSD( 3)
RMSD
RMSD 5)
A-RMSD
EGNO (Xu et al., 2024)
1.058
1.273
1.380
1.463
1.977
1.344
HEGNN al., 2024)
1.051
GF-NODE (Sun al., 2025)
1.046
1.281
1.383
1.431
1.454
1.327
PAINET
though universally accepted accuracy thresholds for multi-body dynamics are difficult to define due to system-specifie variability, we benchmark our performance against criteria established in recent literature. For protein systems, prior work suggests that an RMSD < 2 characterizes high-quality predictions (Xu et al., 2025). Furthermore, in molecular modeling, deviations within the range of 0.01 0.1 are considered consistent with reasonable physical bond length fluctuations (Fedders & Drabold, 1996). The results in Table 10, Table 11 and Table 12 align with these standards, demon- strate that our proposed model effectively yield high-quality predictions that preserve the physical properties.
F.10 QUALITATIVE VISUALIZATION
We provide more demonstration for the trajectories predicted by PAINET in Fig. 8, Fig. 9, Fig. 10 and Fig. 11.
22

---

Published as a conference paper at ICLR 2026
Training time per qpoch (5)
Training per spoch (s)
F-MSE --- Time
F-MSE --- Time
-0.3
1.00
00.25
F-MSE
0.75
F-MSE
0.2
0.15
0.25
0.00
0.05
2
10
Nam of decoding layer
Num of decnding layers
(a) Walk (S25)
(b) Run (S2S)
Figure 5: Ablation studies w.r.t. the number of decoding layers on Motion Capture
0.7
Trainsing time per epoch (5)
0.0090
Training time por epoch (5)
F-MSE --- Time
00.225
F-MSE
Time
0,0085
F-MSS
0.6
F-MSE
+0.175
0,0080
00.25
0.5
0.150
0,0075
0.4
0.125
0.0070
Num of atteetton layers
Nurn of anttennionn layers
(a) Benzene
(b) Salicylic
Figure 6: Ablation studies w.r.t. the number of attention layers on Molecular Dynamics.
GPU menory usage (GB)
GPU menear usage (CI8)
50
GPI memony usage
40
Inferenee time @
10
GPU memory usege
25
Inferrence time (
Time
Time
40
30
20
30
15
20
10
10
15
20
25
300
400
500
600
700
800
Num time steps
Num of particlss
(a) Scalability w.r.t. time steps.
(b) Scalability w.r.t particle numbers.
Figure 7: Scalability test including inference time and GPU memory cost w.r.t. time steps and particle numbers on Proteins (Adk).
Truth
PAINET
10
1.06
1.09
3.12
1.17
GF-NODE
1.16
1.17
1.11
1.76
Figure 8: Representative snapshots of toluene molecular dynamics, initialized at snapshot 60666.0 ps. PAINET preserves structural characteristics and steady results in long time horizons.
23

---

Published as a conference paper at ICLR 2026
I82T0C0 1-3
[830033000 pat [-3
(81000 Co
(91 E
[NIIOMPO EHG
(97160 E-T
(100000) [=
Ground Truth
10"1)
0.92
GF-NODE
FMSE(X10"))
Figure 9: Representative snapshots of toluene molecular dynamics starting at 79100.0 ps.
p2 1
1392020
(82292
182022000
(1120 1-
(MENO Die
Ground Truth
FABSRO 10-1)
7,81
7,52
7:84
7.91
7.96
8.24
EMSRICTTR 10
Figure 10: Representative snapshots salicylic molecular dynamics starting 56292.0 ps.
(STi641ng)) t-3
PAINET
10"5)
7.84
7.36
T.92
F.08
R.18
GF-NODE
1D*))
8.24
8.44
9:42
$.52
8.58
5.68
R.84
Figure 11: Representative snapshots of salicylic molecular dynamics starting 77169.0 ps.
24


Reviews:
[{"review_id": "yF6Sq84uNE", "reviewer": "Reviewer_sMNn", "summary": "This paper proposes PAINET, a principled SE(3)-equivariant architecture for modeling all-pair interactions in multi-body dynamical systems. The model comprises: (1) a novel physics-inspired attention mechanism derived from the minimization trajectory of an energy function, and (2) a parallel decoder designed to maintain equivariance while enabling efficient inference. Empirical evaluation across diverse real-world benchmarks demonstrates the effectiveness of the proposed approach.", "questions": "* How was the number of trajectory steps $T$ determined for different S2T tasks (e.g., $T=5$ in Table 1)? Was this hyperparameter tuned on a validation set, and is it consistent across systems with different dynamical properties?\n* For predicting trajectories longer than the predefined $T$, what is the inference procedure? Does the model operate autoregressively, and if so, how are potential error accumulations mitigated?", "limitations": "", "rating": 2, "confidence": 4, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "* The integration of all-pair interactions with a parallel, equivariant decoder is technically well-founded and appears to be effectively implemented.\n* PAINET demonstrates strong empirical performance, outperforming existing baselines on several challenging multi-body system benchmarks.", "weaknesses": "* Several key aspects of the methodology are insufficiently motivated or explained, making it difficult to follow.\n  * The rationale for regularizing the distance between updated and current node embeddings (lines 184-185) is unclear. What specific issue does this address, and how was the regularization strength chosen?\n  * The description of the underlying physical principle (lines 192-199) is presented without a clear connection to the model's technical design. It is not evident how this principle directly informs the architecture or contributes to the reported performance.\n  * The introduction of the functions $\\phi_{ij}$ (line 211 and Eq. 7) and  $\\psi_{ij}$ (Eq. 7) appears abrupt. The logical flow and the specific roles these functions play in the overall framework need to be more clearly articulated.\n\n* The paper lacks a well-defined, physically meaningful accuracy threshold for the multi-body dynamics tasks. Without such a benchmark, it is difficult to assess whether the reported performance improvements translate to practical utility in real-world applications.", "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": "HuHTk4SRCB", "reviewer": "Reviewer_cVRm", "summary": "This paper introduces PAINET (Physics-Inspired All-Pair Interaction Network) for 3D dynamics modeling across molecular, protein, and motion capture datasets. The method formulates an explicit energy function whose gradient defines pairwise attention weights, ensuring each network layer corresponds to an energy descent step. The encoder models implicit all-pair interactions, while the decoder uses an equivariant EGNN for SE(3)-consistent trajectory prediction. The authors claim PAINET learns latent long-range interactions without explicit graph structures, outperforming prior GNN-based models such as EGNN, HEGNN, EGNO, and GF-NODE.", "questions": "1. Please clarify the differences in preprocessing, subject selection, frame rate, or hyperparameters that might explain the inconsistent EGNO and GF-NODE results.\n2. How sensitive is the model to the specific choice of potential function ρ? Did you test other learnable monotonic forms beyond the quadratic–quartic setup?\n3. Could you include or discuss comparisons removing the energy weighting, Φ/Ψ mappings, or limiting interactions to local neighbors to substantiate the claimed benefits?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The attention update is rigorously derived from an energy descent principle, providing a clear physical interpretation. The energy-derived attention mechanism replaces softmax with gradient-based weights, improving interpretability.\n2. Experiments show consistent gains across motion, molecular, and protein dynamics tasks at comparable computational cost.\n3. The paper is well-organized and readable, with clear motivation and intuitive explanations.", "weaknesses": "1. The reported results for key baselines, particularly EGNO and GF-NODE, appear inconsistent with their original publications, casting doubt on the validity of the claimed improvements. For example, in Table 1, PAINET reports EGNO achieving F-MSE = 14.2 (×$10^{-2}$) on Walk and F-MSE = 4.15 (×10) on Run—values that deviate substantially from EGNO’s original paper (F-MSE = 8.1 (×10-2) on Walk and F-MSE = 3.39 (×$10^{-1}$) on Run). Similar inconsistencies appear for GF-NODE and other baselines, likely due to preprocessing or implementation differences. These discrepancies need careful clarification.\n2. The paper lacks ablations isolating the contributions of the energy function, Φ/Ψ mappings, and all-pair interactions. It remains unclear whether performance gains stem from the proposed energy mechanism or general architectural effects.\n\nMissing related work and baselines:\n1. SE(3) Equivariant Graph Neural Networks with Complete Local Frames; ICML 2022;\n2. AlphaNet: Scaling Up Local Frame-based Atomistic Foundation Model, Npj Comput. Mater. (2025)", "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": "3dyzeUHxsL", "reviewer": "Reviewer_vf4Y", "summary": "Current state-of-the-art models for 3D dynamics prediction (e.g., GNNs, EGNNs) rely on explicitly observed structures between particles. This limits their ability to capture crucial unobserved interactions (e.g., long-range forces, dynamically forming structures), leading to errors in long-term trajectory prediction and simulation, especially in complex systems like molecular dynamics and protein folding. A principled, SE(3)-equivariant neural architecture designed to model all-pair interactions, including unobserved ones.", "questions": "1) Can you explain why the line in Figure 7 shows linear complexity with respect to particle numbers but not O(N^2)?\n2) The training and inference time should be compared among the baseline.\n3) Could you help to explain how to connect the motivated examples in the introduction with the examples in experiments? Are there prior knowledge on the all-pair interations in experiments?\n4) I am not an expert in molecular dynamics and can not justify the significance of the accuracy improvement. Therefore, I will refer other reviewers' comments on it.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Models Unobserved Interactions: This is its core innovation. Unlike previous GNN-based models that rely on fixed, observed structures (e.g., distance-based graphs), PAINET is designed to capture latent, all-pair interactions (e.g., long-range forces, dynamically forming bonds), which is crucial for accurate long-term dynamics prediction.\n\nPrincipled, Physics-Inspired Formulation: The model is not a purely black-box architecture; it is derived from the minimization of an energy function, providing a theoretical foundation for its attention mechanism and linking it to physical principles.", "weaknesses": "A more clear description on the computational cost should be described. See Questions.", "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": "9S2S5KXVtz", "reviewer": "Reviewer_FRS4", "summary": "This paper addresses a key limitation of existing GNN-based methods in 3D dynamic modeling of multi-body systems: over-reliance on explicit observed structures, which fails to capture unobserved all-pair interactions critical for complex physical behaviors. It proposes PAINET, an SE(3)-equivariant neural architecture and evaluates on three real-world benchmarks: human motion capture (CMU dataset), molecular dynamics (MD17 dataset), and large-scale protein simulation (Adk dataset), comparing with classic models (Linear, RF, MPNN, EGNN) and state-of-the-art (SOTA) models (EGNO, HEGNN, GF-NODE).", "questions": "The authors should seriously address the concerns shown in Weakness point by point to improve the quality of the paper.", "limitations": "", "rating": 4, "confidence": 5, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. model operation deisng with physical priors\n2. several benchmark evaluations on different scenarios\n3. good efficiency", "weaknesses": "1. Incomplete Theoretical Analysis​\nInsufficient explanation of the energy function’s physical meaning: While extended from Zhou et al.’s (2004) quadratic energy, its connection to real physical systems (e.g., van der Waals forces, protein hydrogen bonds) is unclear. The concavity assumption of the ρᵢⱼ function only cites Yang et al. (2021) without analyzing its validity boundaries in 3D dynamic scenarios (e.g., particle type differences, dynamic interaction changes), raising doubts about applicability in extreme cases (e.g., strong nonlinear force-driven systems).​\nThough SE(3) equivariance is fully proven, model performance under complex symmetries (e.g., mirror, scaling symmetry) common in real physics is unaddressed. Additionally, equivariance preservation of adaptive pairwise mappings Φ and Ψ in the attention network is not verified individually; while overall equivariance holds, the symmetry transfer mechanism of individual modules remains unclear.​\n2. Room for Improvement in Experimental Design​\nMolecular dynamics experiments ignore force field impacts (e.g., CHARMM, AMBER), only validating on MD17’s quantum force field, limiting generalization to classical force fields. Testing focuses on 8 common small molecules, lacking complex molecules (e.g., metal-containing molecules, polymers), restricting applicability in specialized chemical scenarios.​\nProtein simulation relies solely on Adk protein equilibrium trajectories, excluding other protein types (e.g., membrane proteins, antibodies) or non-equilibrium dynamics (e.g., protein folding, ligand binding), insufficiently evaluating performance in complex conformational changes. Multi-step prediction only tests T=5, without exploring error accumulation at longer steps (e.g., T=10, T=20), inadequately validating long-term prediction ability.\n3. More evalutions on large molecules, such as MD22.", "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)