Structural Causal Model for Molecular Dynamics Simulation
Abstract
Molecular dynamics (MD) simulations describe the mechanical behaviors of molecular systems through empirical approximations of interatomic potentials.
Machine learning-based approaches can improve such potentials with better transferability and generalization.
Among them, graph neural networks have prevailed as they incorporate the graph structure prior while learning the interatomic interactions.
Nevertheless, the simple design choices and heuristics in devising graph neural networks make them lack an explicitly interpretable component to identify the true physical interactions within the underlying system.
On the other extreme, physical models can give a rather comprehensive description of a system but are hard to specify.
Causal modeling lies in between these two extremes, and can provide us with more modeling flexibility.
In this paper, we propose a structural causal molecular dynamics model (SCMD), the first causality-based framework to model interatomic and dynamical interactions in molecular systems by inferring causal relationships among atoms from observational data.
By concerning the underlying data generation process, one can enjoy a robust and flexible MD simulation model to explicitly capture the long-range and time-dependent movement dynamics.
We demonstrate the efficacy of SCMD through empirical validations on the complex molecular system (i.e., single-chain coarse-grained polymers in implicit solvent) for long-duration simulation and dynamical property prediction.
Introduction
Molecular dynamics (MD) simulations aim to analyze the physical movements of atoms and molecules.
It has become an indispensable computational tool for studying molecular system behaviors in a variety of fields, from material science and catalysis to biological processes [Boero et al., 1998, Massobrio et al., 2015, Karplus and McCammon, 2002, Holdijk et al., 2022].
However, molecular dynamics simulation often suffers from a trade-off between accuracy and efficiency.
Highly accurate simulations require resource- and time-intensive ab initio quantum-mechanical calculations, while empirical force fields model interatomic potentials by simple approximations and thus cannot guarantee accuracy and transferability [Zhang et al., 2018].
In recent years, machine learning-based approaches have been introduced to model the molecular system interactions and serve as neural potential functions for MD simulation [Unke et al., 2021].
Among them, graph neural networks (GNNs) are commonly used to learn the interactions between atoms in the system [Chmiela et al., 2017, Unke and Meuwly, 2019, Sch¨utt et al., 2017, Gasteiger et al., 2019, 2021, Batzner et al., 2022].
Although these approaches have achieved remarkable performance for energy and force predictions, the heuristics for considering atomic interactions in the system are rather simple, which often only include bonded interactions or set a distance threshold for atomic interactions.
In contrast, none of these challenges exists in the physics-based approaches, since the approximations of parameters are given to simulate the dynamics according to different types of interaction, atoms and bonds.
However, the efficiency and transferability may be limited compared to ML-based methods.
In view of the advantages and challenges of physical and ML-based approaches, we aim to strike a better trade-off between simulation accuracy and computational efficiency for molecular dynamics through the lens of causal discovery and learning.
As argued in [Sch¨olkopf et al., 2021, Peters et al., 2017] and summarized in Table 1, causal discovery and learning aim to arrive at physical models in a data-driven manner, without manually designed parameters by the physical models, which lie in between those two extremes.
Therefore, instead of heavily relying on physical functions or purely applying black-box prediction models, this work explores the identification of the causal relationship among atoms in molecular systems by only using observational data.
The causal connections between two atoms exist when there is a mechanistic function between them in the physical model so that we can uncover the physical interactions solely from the observed dynamics trajectories, enabling interpretability and efficiency of the modeling simultaneously.
Specifically, we proposed SCMD, which uses the structural causal model (SCM)
[Pearl, 2009,
Sch ¨ olkopf et al., 2021] to model the dynamical behavior of molecular dynamics.
SCM can describe the data generation process of the observed molecular dynamics trajectories, and the corresponding causal relationships can be represented via the dynamic Bayesian networks (DBNs)
[Pearl, 2011].
In DBNs, each node denotes an atom in the system, and the associated edges represent the causal dependencies between atoms.
We can learn the DBN from data by parameterizing the connections using GNNs with recurrent designs to model the spatial and temporal dynamics.
Since the molecules have different causal graphs but share the same underlying physical dynamic laws, we use an autoencoder to learn the generic dynamics laws among all molecules, enabling the zero-shot generalized capability on unseen molecules.
We can infer the causal graphs via encoders and predict the future forces or positions of all atoms using the decoder.
In short, SCMD can learn physics-related insights, which are equipped with the capability to generalize to unseen systems.
Table 1 compares the model attributes of our approach with ML-based and physical models.
In summary, our main contributions are four-fold: (i) .
Providing a systematic framework for leveraging causal discovery to better model the physical and causal dependencies among atoms in the MD system; (ii) .
Developing an efficient method to learn the components of the DBNs that describes the causal model of the MD system based on the neural relation inference; (iii) .
Presenting a more accurate and robust simulation performance on MD simulation with different simulated duration.
Especially, our model can predict long-duration trajectory ( 350 steps) with only observable states on the initial step; and (iv) .
Providing findings on how to leverage causality to better understand the physical and chemical systems.
Problem definitions and preliminaries 2.1 Problem definition
Notations The design target for the MD simulation is to model the evolution of the system of the atoms in an accurate and cost-effective manner.
We use x i = (cid:0) x 1 i , x 2 i , . . .
, x Ti (cid:1) and i ∈ { 1 , 2 , . . .
, N } , where T and N indicate the number of total time steps and atoms in the system.
Learning target Our target is to predict future k -step movements for each atom ˆ x t +1: t
+ k = (cid:0) ˆ x t +1 , ˆ x t +2 , . . .
, ˆ x t + k (cid:1) given the current states of the atoms x t 2 .
The objective function is to minimize the distance between the simulated trajectories and the corresponding ground truth generated by physical simulation estimator, D (ˆ x t +1:
t + k , x t +1: t + k ) .
[Table]
++
[Table]
✓
[Table]
✓ +
2.2 Structure causal molecular dynamics (SCMD) 2.2.1 Causality and structural causal model
Causal inference and discovery tools are designed to identify the underlying causal relations among variables with the interventional or observed data [Pearl, 2009].
The structural causal model (SCM)
[Pearl, 2009] and potential outcome model (POM)
[Rubin, 2005] are two common causal inference frameworks.
In this paper, we choose to leverage SCM to model the data generation process of the molecular dynamic system.
An SCM consists of sets of endogenous ( X ) and exogenous ( U ) variables connected by certain functions ( F ) that determine the values of the variables in X based on the values of the variables in U [Pearl, 2009].
The structure can be represented by a directed acyclic graph (DAG)
G :=
( V , E ) , encoding the direct causal dependencies (edges E ) among variables (node set V ).
For an MD simulation, we assume that the nodes denote atoms and edges denote interactions among atoms.
Unlike ordinary SCM modeling, which excludes time evolution, in this work, we think of the SCMs as the state abstractions of the underlying temporal process [Peters et al., 2017].
We give the SCM and data generation process of MD system below.
Data generation process of MD In the MD system, the movement of each atom will be affected by a group of atoms through physical links ( e.g. , chemical bounds).
Therefore, at time step t , the observed trajectory X t = ( x t 1 , x t 2 , . . .
, x tN ) is generated by the following dynamical process, modeled by an SCM:
[Equation]
where PA ( x ti ) =
(cid:8) x t − 1 j |∀ x j = PA ( x i ) in G (cid:9) , where G is the skeleton graph of the SCM.
The nodes are denoted by a set of atom indexes { x 1 , x 2 , . . .
, x N } the corresponding adjacent matrix
A can be estimated by the observed data.
We assume that the movement of the atoms at t + 1 is only affected by those at time step t , similarly to those works on causal dynamic systems [Lippe et al., 2022, Wang et al., 2022, L ¨ owe et al., 2022].
In other words, there are no instantaneous temporal effects in the SCM.
Fig.
1(a) indicates the data generation process, including the observed trajectories and the underlying SCM.
Notably, we assume that there exists a self-loop for all atoms.
This is because the movements of each atom are affected by their historical movement variables ( e.g. , locations).
Hence, the SCM (Fig. 1(a)) may contain self-loops and cycles in the graph.
However, this will not affect the intervention and identifiability theory of causal discovery since the timedependent process ( i.e. , full-time graph) does not have cycles (Fig. 1(b))
[Peters et al., 2017].
To conclude, the SCM can extract causal relations among different atoms, expressing the underlying data generation process that takes place over time.
2.2.2 Dynamic Bayesian networks
A Bayesian network (BN) [Dean and Kanazawa, 1989] over the set X is defined as B =
( G X , P ) , where G X is the graph structure of the network, giving the dependency connectivity among the variables, and the distribution P gives the conditional probabilities of the variables in X :
[Equation]
DBNs [Murphy, 2002] are one of the extended versions of BN, modeling the time-dependent process of variables and their conditional dependence.
The unfolded DBNs can be represented as BNs.
The variables in DBNs are in discrete time steps and depend on variables only from the same and previous time slices.
The network structure and conditional dependence are the same as in other time
[Figure]
t t + steps, which assume the stationary property of the system.
We can infer the SCM by approximating the distribution over the DBN in Equation 2.
In other words, we use DBN to model the SCM in a probabilistic way.
To model the probabilistic SCM with temporal effects, each conditional dependence in the
DBN should represent the causal relationship.
We assume B is the DBN model over (cid:8) x 11 , x 12 , . . .
, x tN − 1 , x tN (cid:9) .
Thus, the DBN is the unfolded causal graph with connections beyond time steps, modeling the temporal effects among atom pairs.
Fig. 1(b) displays the visualized DBN with two-time slices.
In the present paper, we learn the DBN using the auto-encoder model, similar to [L ¨ owe et al., 2022], where the causal dependence structure is inferred via the latent dimensions, and the temporal dynamics are estimated by the decoder.
2.2.3 Neural relation inference
We use the neural relation inference (NRI) framework [Kipf et al., 2018, L ¨ owe et al., 2022] to learn the DBN.
In NRI, the interactions among entities are learned via the variational auto-encoder (VAE) model, where the interactions are modeled as the latent codes inferred by the encoder, and the decoder reconstructs the original trajectories.
The evidence lower bound (ELBO) of the VAE is given by:
[Equation]
where, the latent vector z denotes the inferred interaction among entities X = { x 1 , x 2 , . . .
, x N } and we have observed trajectories x i = (cid:0) x 1 i , . . .
, x Ti (cid:1) of all entities.
N and T are the number of entities and time steps, respectively.
p ( z ) is the prior distribution of the latent space.
ϕ
and θ denote the encoder and decoder, respectively.
The amortized causal discovery [L¨owe et al., 2022] further proposes to leverage the latent space to learn causal relations between entities.
Specifically, the work extends the NRI with an additional edge type defining the non-causal relations in the latent space.
In our work, we follow the designed structure of the amortized causal discovery, where we learn the DBN using the VAE model and estimate the temporal dynamics by the decoder.
In the original amortized causal discovery, the applied scenarios are relatively simple and the datasets are all with low dimensions ( e.g. , 1-D time-series, particles, and simulated fMRI data).
We explore how to leverage the framework to model the complex MD system with possible noisy and high-dimensional data.
In particular, Zhu et al.
[2022] also utilize the NRI framework to infer the latent interactions of protein allosteric communications.
Different from [Zhu et al., 2022], we aim to propose a generic framework for discovering the causal links in a molecular dynamics system.
Structural Causal Molecular Dynamics
We utilize a generative model, named sequential causal VAE (SC-VAE), to learn the graph structure and temporal causal dependencies of DBN in an end-to-end fashion.
The overall pipeline is illustrated in Fig.
In SC-VAE, the atom representation is learned via several GNN layers and serves as the input of the network.
The encoder infers the distribution of causal graphs given the atom representation.
Specifically, we model the adjacent matrix of the causal graph as the latent factors.
Hence, the produced causal graphs are atom-specific, making the framework generic and flexible.
The decoder generates future movements by modeling the temporal interactions among atoms with the sampled causal graph instance from the latent space and the learned atom representation.
The supervision signals include the sequential prediction error and the KL-divergence between the posterior and prior distribution of the causal graph.
We discuss each component below in our proposed framework.
3.1 Learning the generative process of SCMD
[Figure]
Learning the edge representation H
Following NRI, we use the neural message passing to obtain the edge embeddings as the input to estimate the causal dependencies.
The edges embeddings
H are computed based on the message passing on a fully-connected graph, where each atom will propagate the message to all the remaining in the molecular system.
Thus, the edge embedding is H
=
{∀ i, j ∈ { 1 , . . .
, N }
and i ̸
=
j | h ( x i , x j ) } , where h ( x i , x j ) is the embedding of edge connected with x i and x j .
Specifically, to calculate h ( x i , x j ) , we follow the steps below:
Step 1: Atom-wise message
passing
[Equation]
where f emb is composed of several stacking MLP layers, deriving the initial embedding representation of atoms.
f e 1 is to compute the message passing between atom pairs, implemented by MLP layers.
The output of step 1 is the initial edge embedding h 0 .
Step 2: Edge-wise message passing
[Equation]
where f v is the edge message passing function, implemented by the stacking MLP layers.
f j is the updated atom representation with the messages passed from all the edges.
Step 3: Final Edge embedding
[Equation]
where f e 2 is another module to compute the message passing between atoms, similar to f e 1 but with different parameters.
The final edge embedding between x i and x j is then derived as h ( x i , x j ) .
The edge embeddings in H , are then used to estimate the distribution of the causal graph and the temporal dynamics.
Learning the causal graph structure A Following the causal discovery with NRI framework [Kipf et al., 2018, L¨owe et al., 2022], we utilize an encoder to learn the sample-agnostic causal graph.
This is particularly appropriate because most molecule samples with the same chemical space share the same underlying physical dynamics given some sound conditions ( e.g. , same media), while the causal graphs for describing the atom movements are sample-specific due to the different molecular structures.
The encoder q ϕ , hence, to learn the structure of the causal graphs among all samples.
Given a trajectory data sample x t at time step t , the encoder output the distribution of the adjacent matrix of the sample q ϕ
( A | g ( x t )) , where A is the adjacent matrix and q ϕ is parameterized via several dense layers empirically.
We have three edge types, indicating different causal dependencies below:
Type I : Non-causal edges a 0 ; Type II : Causal edges with bounded links a 1 ; Type III : Causal edges without bounded links a 2 .
Hence, q ϕ ( A | g ( x t )) will generate the probability distribution of each type edge based on the embedding and sparse constraints.
[Equation]
where f c is a dense layer to classify the edge type.
a ij ∈ R 3 , indicating the likelihoods of which types of casual links the edge belong to.
We have the prior beliefs that the edge types are uniformly distributed.
Hence, the prior distribution of p ( A ) is a uniform distribution and
L KL = KL ( q ϕ ( A | x ) ∥ p ( A )) .
In the meantime, we also have the regularization terms based on the sparse causal mechanism [Sch¨olkopf et al., 2021] of A :
[Equation]
p ( a 1 ) and p ( a 2 ) correspond to the total output likelihoods of causal edges of all atoms.
Learning the temporal dynamic model p θ
With the learned causal structure A , we have the skeleton of the DBN and the only thing to do is estimating the edge strength of each path in DBN.
Formally, the path strength links to the dynamics model p θ ( x t +1 | g ( x t ) , A ) .
We use g to represent all functions involved to compute the edge embeddings H .
We utilize two GNNs to parameterize p θ .
Specifically, one GNN f dyn1 is trained to model the interaction between those atoms with edge
Type III, and another f dyn2 is to model those atoms linked with edge Type II.
The loss for the temporal dynamic prediction is L dyn =
MSE ( x t +1 , p θ ( x t +1 | g ( x t ) , A )) , where MSE denotes the mean square error.
Total loss In general, the total loss is given below: L total =
α 1 L dyn + α 2 L KL + α 3 L reg , where { α 1 , α 2 , α 3 } are tunable hyper-parameters.
Experiments
[Figure]
4.1 Setup
We validate the effectiveness of our method on long-range and long-duration MD simulations.
Specifically, we aim to simulate x 2: T by only observing the x 1 of the molecular system.
Training To achieve this, during the training phase, we have the collected MD trajectories D tr , each trajectory corresponds to a specific molecule with the number of time steps T .
We train the SCVAE using D tr and fix the trained parameters q ϕ and p θ .
We observe the initial states ( i.e. , locations
[Figure]
and velocities)
x 1 of all atoms and predict the future steps with the inferred causal relations and temporal dynamics.
At t = 2 , we use the initial value x 1 of the system as input.
To predict the states of atoms at time step t ( t > 2 ), we use the estimated states ˆx t − 1 at time step t − 1 .
Simulation During the simulation phase, we use the trained model to predict the whole trajectories by observing the initial states (see the illustration at the r.h.s.
of Fig.
3(A)).
Dataset We validate the performance of SC-VAE over a single-chain coarse-grained polymer dataset simulated with LAMMPS
[Thompson et al., 2022].
The simulation setup, such as force field parameters, is described in [Webb et al., 2020] and the dataset is simulated by [Fu et al., 2022].
As shown in Fig.
3(B), each polymer is composed of 4 types of constitutional units (CUs) with different numbers ( n ) of the such repeated pattern.
And the dynamics of each polymer corresponds to the trajectory of 50k
τ timesteps with a recording frequency of 5 τ , where τ is the unit time.
We have a total of 100 polymers, 90 of which are used as the training set, and the remaining 10 polymers are used for simulation.
Baseline We compare our method with MLCGMD [Fu et al., 2022], the first work exploring the coarse-graining simulation of MD using ML tools.
MLCGMD utilizes a score-based GNN to model atomic interactions, which embeds and coarse-grains an MD system to a coarse-level graph via clustering.
At time step t , the historical states x t − k : t − 1 are treated as the input and the framework generates the simulated state at time t (see the l.h.s.
of Fig.
3(A)).
4.1.1 Implementation
Training and simulation Our model generates the simulated trajectory with a duration of T steps.
That is, we split the whole trajectory into ⌊ 10 , 000 /T ⌋ disjoint sequences and predict the atom movements in each sequence by observing the first state.
To test the robustness of our method in terms of the sequence length, we vary T ∈ { 10 , 50 , 100 , 150 , 200 , 250 , 300 , 350 } .
Different from ours, MLCGMD utilizes the input with 30 history states, and generates the next step each time.
Experimental details
The hyper-parameters we used are summarized in Appendix Table B and
C. All methods are implemented using PyTorch [Paszke et al., 2019] toolbox with an Intel Core
i
CPU and 4 Nvidia RTX 3090 GPUs.
4.2 Results 4.2.1 Evaluation metrics
Following [Fu et al., 2022], we use the estimated surrogate property of the radius of gyration ( R 2 g ) to evaluate the precision of the simulation models.
We use two metrics in terms of R 2 g :
Earth mover’s distance Earth mover’s distance (EMD) is a measure of the distance between two probability distributions over a specific region.
Here, we use EMD to measure the distance between simulated trajectories and the ground truth.
Hence, a small EMD leads to a more accurate simulation.
Mean absolute error Mean absolute error (MAE) is also adopted to measure the distance between the simulated movements and the ground truth for each time step.
4.2.2 Results and analysis
Fig.
4 gives the comparison between our method (with T = 350 and 10 ) and MLCGMD, where the plots indicate the simulated R 2 g for all time steps.
From the plots, we can find our model leads to a more accurate dynamic simulation.
We give a detailed analysis in terms of the robustness of the simulation duration and sample efficiency below.
Robustness against different simulation duration T .
We vary the simulation duration T ∈ { 10 , 50 , 100 , 150 , 200 , 250 , 300 , 350 } .
The full results are given in Appendix Fig.
A.
Though the accuracy is higher when the duration is smaller, our model can still result in a better simulation than the baseline when the duration is extremely long ( e.g. , T = 30 ).
The results verify that our model
[Figure] [Figure]
Time
[Figure]
R a d i u s o f g y
r
a t
i o n ( æ 2 )
Time ( ø )
R a d i u s o f g y
r
a t
i o n ( æ 2 )
[Figure]
MLCGMD
[Figure]
can predict a very long sequence with only one truth observation of the initial step.
However, the baseline model requires 30 steps of historical values to predict a future step.
This makes our model robust and potentially be applied to many down-streaming MD applications (e.g., protein modeling).
Better simulation with fewer samples.
Table A gives the results of accuracies versus the required samples.
We follow the original settings in MLCGMD to use only 30 samples at the beginning of the trajectories to train the model.
We find that our model can learn a better simulation model (lower
MAE and EMD) with fewer samples ( 28 versus 30 ) compared with MLCGMD.
[Figure]
Conclusions
We propose the Structural Causal Molecular Dynamics (SCMD), exploring the potential to take advantage of causality to learn a more accurate and robust molecular dynamic simulation.
We formalize the MD system as a structural causal model and utilize a dynamic Bayesian network to estimate the causal model.
The causal model does not model the specific functions of MD, but instead includes the interaction relations and strength analogously to the true physical dynamics.
To learn the model, we apply and extend the neural relation inference framework to the MD system to learn the DBN in an end-to-end fashion.
The empirical results verify the effectiveness and efficiency of our model with different prediction durations.
Based on both algorithmic and empirical aspects, we believe that our model can be applied to many downstream MD applications, e.g., drug discovery.
Acknowledgement
We thank Ge Zhang, Yiming Liang, Zekun Wang for helping with experiments, and Yucheng Zhang,
Chenqing Hua, Sitao Luan, Shuangjia Zheng, Jianmin Wang for helpful discussions.
A Limitations and future works
In this work, we only use the sparse constraints to learn the causal graph from the observational data.
Though it is a normal setup in many works on causal representation learning for dynamic systems [Wang et al., 2021], it lacks some grounded causal discovery foundations.
In future work, we plan to explore more aspects of causality, including intervention and counterfactual learning [Peters et al., 2017], to learn the causal graph and its applicability to complex MD simulation systems.
B Related work B.1 Machine Learning for Molecular dynamics
Owing to its efficiency and transferrability, machine learning methods have been extensively studied for molecular dynamics [No´e et al., 2020].
Specifically, there are two main ways that machine learning methods have been used to improve the accuracy and/or efficiency of molecular dynamics simulations.
The first way is to leverage machine learning models to fit hand-crafted parameters in molecular dynamics simulations.
For example, traditional molecular dynamics rely heavily on hand-crafted force field parameters.
Machine learning methods are used to learn from ab initio quantum-mechanical data and serve as a machine learning force field for molecular dynamics with ab initio-level accuracy and increased efficiency [Smith et al., 2017, Zhang et al., 2018, Doerr et al., 2021, Batzner et al., 2022].
Another line of work leverages deep generative models to learn to sample probability distribution from the observational data.
Boltzmann generator proposes a normalizing flow that generates samples from the equilibrium distribution of the molecular system [No ´ e et al., 2019].
This approach can be used to avoid incredibly long molecular dynamics simulations for studying complex dynamical behaviors such as protein folding.
B.2 Causal Dynamical System
Dynamical systems have been widely used to describe the dynamical behaviors of systems in various fields, including physical systems [Nelles, 2001, Sanchez-Gonzalez et al., 2018, Nelles, 2001], neuroscience [Izhikevich, 2007, Van Bavel et al., 2015], robotics [Kober and Peters, 2009, Girard et al., 2008].
Recently, a range of works in the causality literature aim to uncover the causal semantics of the dynamic system.
Bongers et al.
[2018] provide the formal framework named structural dynamical causal models (SDCMs) to model the dynamics system as the structured casual model with differential equations.
Mooij et al.
[2013], Blom and Mooij [2018] describe the asymptotic solution of the dynamic systems as a causal model and consider the interventions on ODE systems.
Peters et al.
[2022] propose causal kinetic models to model the continuous time models with ODE functions by intervening on the dynamic systems.
Based on these theoretical frameworks, several works have been proposed to learn the causal representation of dynamic systems under real-world or simulated scenarios.
Yao et al.
[2022] learn latent causal dynamic systems under modular distribution shifts and provide the conditions for identifiability guarantees on causal graphs.
Lippe et al.
[2022] learn the causal factors from the sequences of observations and design interventions on dynamic systems.
Weichwald et al.
[2022] propose a general framework to use causality for the control tasks in dynamic systems.
Huang et al.
[2021], Feng et al.
[2022] learn the factored and causal structure of the dynamic systems in RL under heterogeneous and non-stationary environments.
In this work, we first leverage the causal dynamic system to model molecular dynamics, learning how the atoms evolve and interact temporally in the system.
C Full results and hyper-parameters
Fig.
A and Table A give the full results with different simulated durations and required label samples.
Table B and C list the chosen architectures and hyper-parameters, respectively.
[Figure]
a b c d e f g h
[Figure]
[Table] [Table]
MLPs (atom)
[Table]
ELU - ELU + BatchNorm
[Table]
(edge)
[Table]
ELU - ELU + BatchNorm
MLP (atom)
[Table]
ELU - ELU + BatchNorm
MLP (edge)
[Table]
ELU - ELU + BatchNorm f dec
MLP (static atom)
ReLU+LayerNorm
MLP (static edge)
ReLU+LayerNorm
MLP (static edge)
ReLU+LayerNorm
MLP (static node)
ReLU+LayerNorm
MLP (static edge) 64ReLU+LayerNorm - ReLU+LayerNorm
MLP (static node)
[Table]
ReLU+LayerNorm - ReLU+LayerNorm
MLP (temporal atom)
ReLU+LayerNorm
MLP (temporal edge)
ReLU+LayerNorm
MLP (temporal edge) [ 128 - 128 ] ×[ReLU+LayerNorm - ReLU+LayerNorm] ×
MLP (temporal node)
[Table]
ReLU+LayerNorm - N/A
[Table]