Newton–Cotes Graph Neural Networks:
On the Time Evolution of Dynamic Systems

We suppose that a system comprises $N$ particles $p_{1},p_{2},...,p_{N}$, with the velocities ${\mathbf{v}}_{1},{\mathbf{v}}_{2},...,{\mathbf{v}}_{N}$ and the coordinates ${\mathbf{x}}_{1},{\mathbf{x}}_{2},...,{\mathbf{x}}_{N}$ as states. The particles can also carry various scalar features (e.g., mass and charge for atoms) which will be encoded as embeddings. We follow the corresponding existing works [17, 18, 19] to process them and do not discuss the details in this paper.

Although the existing EGNN methods have great advantage in computational cost over the conventional simulation tools, the potential prediction error that a neural model needs to compensate for is also huge. Take the standard EGNN [18] as an example, the equivariant graph convolution can be written as follows:

where the aggregation function $\phi_{e}$ takes three types of information as input: the feature embeddings ${\mathbf{h}}_{i}$ and ${\mathbf{h}}_{j}$ for the input particle $p_{i}$ and an arbitrary particle $p_{j}$, respectively; the Euclidean distance $||{\mathbf{x}}_{i}^{0}-{\mathbf{x}}_{j}^{0}||^{2}$ at time $t=0$; and the edge attribute $e_{ij}$. Then, the predicted velocity and future coordinate can be calculated by the following equation:

where $\phi_{v}:{\mathbb{R}}^{D}\rightarrow{\mathbb{R}}^{1}$ is a learnable function. From the above equations, we can find that the predicted velocity $\hat{{\mathbf{v}}}_{i}^{0}$ is also correlated with three types of information: 1) the initial velocity ${\mathbf{v}}_{i}^{0}$ as a basis; 2) the pair-wise Euclidean distance $||{\mathbf{x}}_{i}^{0}-{\mathbf{x}}_{j}^{0}||^{2}$ (i.e., ${\mathbf{m}}_{ij}^{0}$) to determine the amount of force (analogous to Coulomb’s law); and the pair-wise relative coordinate difference $({\mathbf{x}}_{i}^{0}-{\mathbf{x}}_{j}^{0})$ to determine the direction. For multi-layer EGNN and other EGNN models, $\hat{{\mathbf{v}}}_{i}^{0}$ is still determined by these three types of information, which we refer interested readers to Appendix A for details.

As the model makes prediction only based on the state and velocity at $t=0$, we assume that ${\mathbf{v}}^{t^{*}}$ fluctuates around ${\mathbf{v}}^{0}$ and follows a normal distribution denoted by ${\mathcal{N}}_{NC(0)}=({\mathbf{v}}^{0},\sigma^{2}_{NC(0)})$. Then, the variance term $\sigma^{2}_{NC(0)}$ reflects how difficult to train a neural model on data sampled from this distribution. In other words, the larger the variance is, the worse the expected results are.

In comparison with the $0$ degree polynomial approximation $\text{NC}(0)$, the high order $\text{NC}(k)$ is more accurate and yields only a little extra computational cost.

Suppose that we now have $K+1$ (usually $K\leq 8$ due to the catastrophic Runge’s phenomenon [42]) points ${\mathbf{x}}_{i}^{0},{\mathbf{x}}_{i}^{1},...,{\mathbf{x}}_{i}^{K}$ equally spaced on time interval $[0,T]$, then the prediction for the future coordinate $\hat{{\mathbf{x}}}_{i}^{T}$ in Equation (4) can be re-written as:

where $\{w^{k}\}$ is the coefficients of Newton-Cotes formulas and $\hat{{\mathbf{v}}}_{i}^{k}$ denotes the predicted velocity at time $t^{k}$. To obtain $\hat{{\mathbf{v}}}_{i}^{k},k\geq 1$, we simply regard EGNN as a recurrent model [43, 44] and re-input the last output velocity and coordinate. Some popular sequence models like LSTM [45] or Transformer [13] may be also leveraged, which we leave to future work. Then, the equation of updating the predicted velocity $\hat{{\mathbf{v}}}_{i}^{k}$ at $t^{k}$ can be written as:

where ${\mathbf{v}}_{i}^{k-1}$, ${\mathbf{x}}_{i}^{k-1}$ denote the velocity and coordinate at $t^{k-1}$, respectively. Note that, Equation (13) is different from the form used in a multi-layer EGNN. The latter always uses the same input velocity and coordinate as constant features in different layers. By contrast, we first predict the next velocity and coordinate and then use them as input to get the new trainable velocity and coordinate. This process is more like training a language model [13, 43, 46, 45, 47].

From the interpolation theorem [26, 48], there exists a unique polynomial $\sum_{k=0}^{K}{C^{k}t^{(k)}}$ of degree at most $K$ that interpolates the $K+1$ points, where ${C^{k}}$ is the coefficients of the polynomial. To avoid ambiguity, we here use $t^{(k)}$ to denote $t$ raised to the power of $k$, whereas $t^{k}$ to denote the $k$-th time point. $\sum_{k=0}^{K}{C^{k}t^{(k)}}$ thus can be constructed by the following equation:

where $[\cdot]$ denotes the divided difference. In our case, the $K+1$ points are equally spaced over the interval $[0,T]$. According to Newton-Cotes formulas, we will have:

Particularly, $\{w^{k}\}$ are Newton-Cotes coefficients that are irrelevant to ${\mathbf{v}}(t)$. With additional observable points, the estimation for the integration can be much more accurate. The objective thus can be written as follows:

In the previous formulation, in addition to the initial and terminal points, we also need the other $K-1$ points for supervision which yields a regularization loss:

We denote NC (k) with the above loss by NC⁺ (k). Minimizing $\mathcal{L}_{r}$ is equivalent to minimizing the difference between the true velocity and predicted velocity at each time point $t^{k}$ for each particle. $\|\cdot\|$ can be arbitrary distance measure, e.g., L2 distance. Although the $K-1$ points of data can be easily accessed in practice, we argue that a loose version (without $\mathcal{L}_{r}$) NC (k) can also learn promising estimation of the integration $\int_{0}^{T}{\mathbf{v}}(t)dt$. Specifically, we still use $K+1$ points $\{\hat{{\mathbf{v}}}^{k}\}$ to calculate the integration over $[0,T]$, except that the intermediate $K-1$ points are unbounded. The model itself needs to determine the proper values of $\{\hat{{\mathbf{v}}}^{k}\}$ to optimize the same objective defined in Equation 16. In Section 4.5, we design an experiment to investigate the difference between the true velocities $\{{\mathbf{v}}^{k}\}$ and the predicted velocities $\{\hat{{\mathbf{v}}}^{k}\}$.

In comparison with EGNN, NC (k) does not involve additional neural layers or parameters. Intuitively, we recurrently use the last output of EGNN to produce new predicted velocities at next time point, the corresponding computational cost should be $K$ times that of the original EGNN. However, our experiment in Section 4.4 shows that the training time did not actually linearly increase w.r.t. $K$. One reason may be the gradient in back-propagation is still computed only once rather than $K$ times.

We present an implementation of NC⁺ in Algorithm 1. We first initialize all variables in the model and then recurrently input the last output to the backbone model to collect a series of predicted velocities. We then calculate the final predicted integration and plus the initial coordinate as the final output (i.e., Equation (12)). Finally we compute the losses and update the model by back-propagation.

We selected three state-of-the-art methods as our backbone models: RF [17], EGNN [18], and GMN [19]. To ensure a fair comparison, we followed their best parameter settings (e.g., hidden size and the number of layers) to construct NC. In other words, the main settings of NC and the baselines were identical. The results reported in the main experiments were based on a model of NC (2) for a balance of training cost and performance.

We reported the performance of NC w/ additional points (denoted by NC⁺) and w/o additional points (the relaxed version, denoted by NC).

We leveraged three datasets towards different aspects of reasoning dynamics to exhaust the performance of NC. We used the N-body simulation benchmark proposed in [19]. Compared with the previous ones, this benchmark has more different settings, e.g., number of particles and training data. Furthermore, it also considers some physical constraints (e.g., hinges and sticks) widely existing in the real world. For molecular dynamics, We used MD17 [28] that consists of the molecular dynamics simulations of eight different molecules as a scenario. We followed [19] to construct the training/testing sets. We also considered reasoning human motion. Follow [19], the dataset was constructed based on $23$ trials in CMU Motion Capture Database [27].

As NC⁺ requires additional supervised data, we accordingly extracted more intermediate data points between the input and target from each dataset. It is worth noting that this operation was tractable as the sampling frequency of the datasets was much higher than our requirement.

We reported the main experimental results w.r.t. three datasets in Tables 1, 2, and 3, respectively. Overall, NC⁺ significantly improved the performance of all baselines on almost all datasets and across all settings, which empirically demonstrates the generality as well as the effectiveness of the proposed method. Remarkably, NC (without extra data) achieved slightly low results but still had great advantages over performance in comparison with three baselines.

Specifically, The average improvement on N-body datasets was most significant, where we observe a near 20% decrease in prediction error for all three baseline methods. For example, NC (RF) and NC⁺ (RF) greatly outperformed the original RF and even had better results than the original EGNN under some settings (e.g., 3,2,1). This phenomenon also happened on the molecular dynamics datasets, where we find that NC⁺ (EGNN) not only defeated its original version, but also the original GMN and even NC⁺ (GMN) on many settings.

We conducted experiments to explore how the performance changes with respect to the estimation step $k$. The experimental results are shown in Figure 2.

For all three baselines, the performance improved rapidly as the growth of step $k$ from $0$ to $2$, but this advantage gradually vanished as we set larger $k$. Interestingly, the curve of training time showed an inverse trend, where we find that the total training time of NC (5) with GMN was almost two times slower than the NC (1) version. Therefore, we set $k=2$ in previous experiments to get the best balance between performance and training cost.

By comparing the prediction errors of different baselines, we find that GMN that considers the physical constraints (e.g., sticks and hinges) usually had better performance than the other two methods, but the gap significantly narrowed with the increase of step $k$. For example, on MD17 and motion datasets, EGNN obtained similar prediction errors for $k\geq 2$ while demanding less training time. This phenomenon demonstrates that predicting the integration with multiple estimations lowers the requirement for model capacity.

Overall, learning with NC significantly reduced the prediction errors of different models. The additional training time was also acceptable in most cases.

The only difference between NC and NC⁺ was the use of the regularization loss in NC⁺ to learn the intermediate velocities. Therefore, we designed experiments to compare the intermediate results and the final predictions of NC and NC⁺ with GMN as the backbone model.

We first tracked the average intermediate velocity prediction errors on valid datasets in terms of training epochs. The results are shown in Figure 3, from which we observe that the prediction errors were not huge for both two methods, especially NC⁺ obtained highly accurate prediction for the intermediate velocities with different $k$. For NC, we actually find that it implicitly learned to model the intermediate velocities more or less, even though its prediction errors gradually improved or dramatically fluctuated during training. This observation empirically demonstrates the effectiveness of the intermediate velocities in predicting the final state of the system. We also reported the results on MD17 and N-body datasets in Appendix C.

We then visualized the intermediate velocities and the final predictions in Figure 4. The intermediate velocities were visualized by adding the velocity with the corresponding coordinate at the intermediate time point. From Figure 4, we can find that NC⁺ (green ones) learned very accurate predictions with only a few non-overlapped nodes to the target nodes (red ones). For the loose version NC, it learned good predictions for the backbone part, but failed on the limbs (especially the legs). The motion of the main backbone was usually stable, while that of the limbs involved swing and twisting. We also conducted experiments on the other two types of datasets, please see Figure 8 in Appendix C.

In previous experiments, we show that NC⁺ was capable of learning highly accurate predictions of intermediate velocities. This inspired us to realize the full potential of NC⁺ by producing a series of predictions given only the initial state as input. For each time point (including the intermediate time points), we used the most recent $k+1$ points of averaged data as input to predict the next velocity and coordinate. This strategy produced much more stable predictions than the baseline NC (0) – directly using the last output as input to predict the next state.

The results are shown in Figure 5, from which we can find that both the baseline (blue lines) and NC⁺ (green lines) had good predictions at the initial steps. However, due to the increasing acculturated errors, the baseline soon collapsed and only the main backbone was roughly fit to the target. By contrast, NC⁺ was still able to produce promising results. The deviations mainly concentrated on the limbs. Therefore, we argue that NC⁺ can produce not only the highly accurate one-step prediction, but also the relatively reliable consecutive predictions. We also uploaded the .gif files of examples on all three datasets in Supplementary Materials, which demonstrated the consistent conclusion.

We also compared our method with the EGNN methods for multi-step prediction [49, 50, 51, 52]. Specifically, multi-step prediction takes a sequence of states with fixed interval as input and predicts a sequence of future states with same length and interval. As recent multi-step prediction methods are still based on sequence or encoder-decoder architectures, it would be interesting to examine the applicability of NC to them. To this end, we adapted NC ($k$=2) to the source code of the best-performing method EqMotion [49] and used identical settings for the hyper-parameters.

The experimental results are shown in Table 4, where we can observe that NC (EqMotion) outperformed the base model EqMotion on all datasets and across all metrics. It is worth noting that the archiecture and parameter-setting of the EGNN module were identical to EqMotion and NC (EqMotion), which clearly demonstrates the effectiveness of the proposed method. We have also released the source code of NC (EqMotion) in our GitHub repository.

We conducted ablation studies to investigate the effectiveness of the proposed Newton-Cotes integration. We implemented two alternative methods: NC (coeffs$=1$), where we set all the coefficients for summing the predicted velocities as $1$; and depth-varying ODE, where we employed common neural ODE solvers from [40, 41] to address our problem. We used an identical GMN as the backend EGNN model and conducted a simple hyper-parameter search for the learning rate and ODE solver.

The results are presented in Table 5. The two alternative methods methods generally performed better than the baseline GMN. The performance of depth-varying GDE was comparable to that of NC (coeffs$=1$), but both were lower than that of the original NC. This observation further supports the effectiveness of our method.