SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Xiao Luo $^{1}$ , Haixin Wang $^{2}$ , Zijie Huang $^{1}$ , Huiyu Jiang $^{3}$ , Abhijeet Sadashiv Gangan $^{1}$ , Song Jiang $^{1}$ , Yizhou Sun $^{1}$
<sup>1</sup>University of California, Los Angeles, <sup>2</sup>Peking University,
<sup>3</sup>University of California, Santa Barbara
{xiaoluo,yzsun}@cs.ucla.edu，wang.hx@stu.pku.edu.cn
Modeling interacting dynamical systems, such as fluid dynamics and intermolecular interactions, is a fundamental research problem for understanding and simulating complex real-world systems. Many of these systems can be naturally represented by dynamic graphs, and graph neural network-based approaches have been proposed and shown promising performance. However, most of these approaches assume the underlying dynamics does not change over time, which is unfortunately untrue. For example, a molecular dynamics can be affected by the environment temperature over the time. In this paper, we take an attempt to provide a probabilistic view for time-varying dynamics and propose a model Context-attended Graph ODE (CARE) for modeling time-varying interacting dynamical systems. In our CARE, we explicitly use a context variable to model time-varying environment and construct an encoder to initialize the context variable from historical trajectories. Furthermore, we employ a neural ODE model to depict the dynamic evolution of the context variable inferred from system states. This context variable is incorporated into a coupled ODE to simultaneously drive the evolution of systems. Comprehensive experiments on four datasets demonstrate the effectiveness of our proposed CARE compared with several state-of-the-art approaches.
Modeling interacting dynamical systems is a fundamental machine learning problem [66, 65, 23, 59, 37] with a wide range of applications, including social network analysis [12, 17, 36] and particle-based physical simulations [46, 43, 35]. Geometric graphs [27] are utilized to formalize these interactions between objects. For example, in particle-based physical systems, edges are constructed based on the geographical distance between particles, which represents the transfer of energy.
In the literature, numerous data-driven approaches have been proposed for understanding interacting dynamical systems [2, 25, 49]. Among them, graph neural networks [26, 62, 72, 33, 13] (GNNs) are widely utilized to predict trajectories at the next timestamp due to their strong capacity to capture interactions in graph-structured data. In particular, each object is considered as a graph node, and edges represent interactions between neighboring objects. Given the observations and their corresponding graph structure, these methods forecast states in the next timestamp using the message passing mechanism. This process involves aggregating information from the neighbors of each node to update its representation in an iterative fashion, effectively capturing the dynamics of the system.
Although impressive progress has been achieved on GNN-based approaches, capturing long-term dependency in interacting dynamical systems is still a practical but underexplored problem. Existing next-step predictors [42, 46, 45] can send the predictions back to generate rollout trajectories, which could suffer from serious error accumulation for long-term predictions. More importantly, system
environments and relational structures could be changeable [69] (e.g., unsteady flow [11, 16]), which implies the potential temporal distribution variation during the evolution. In particular, in physical systems, there are various potential factors which can influence the trajectories extensively. For example, high temperatures [69] or pressure [68] could speed up molecular movement. Their continuous variation would make understanding interacting dynamic systems more challenging. First, temporal environmental variation would indicate different data distributions over the time [52], which requires the model equipped with superior generalization capability. In contrast, existing methods typically focus on in-distribution trajectories [7, 1, 58, 3, 21, 66], which would perform worse when it comes to out-of-distribution data. Second, recent out-of-distribution generalization methods [50, 39, 57, 64, 44] usually focus on vision and text with discrete shift across different domains. However, our scenarios would face the continuous distribution variation, which is difficult to capture in interacting dynamical systems.
In this paper, we propose a novel method named Context-attended Graph ODE (CARE) to capture interacting system dynamics. The core of our CARE approach is to characterize the temporal environmental variation by introducing the context variable. In particular, we propose a probability model to depict the interaction between the context variable and trajectories. Based on the probabilistic decomposition, we divide each training sequence into two parts for initializing embeddings and predictions. Here, we first construct a temporal graph and then leverage an attention-based encoder to generate node representations and the context representation from spatial and temporal signals simultaneously. More importantly, we introduce coupled ODEs to model the dynamic evolution of node representations and the context variable. On the one hand, we adopt a graph-based ODE system enhanced with context information to drive the evolution. On the other hand, the context information can also be updated using summarized system states and the current context. We also provide a theoretical analysis that indicates, at least locally, the future system trajectory and context information are predictable based on their historical values. Finally, we introduce efficient dynamical graph updating and robust learning strategies to enhance the generalization capability and efficiency of the framework, respectively. Extensive experiments on various dynamical systems demonstrate the superiority of our proposed CARE compared with state-of-the-art approaches.
To summarize, in this paper we make the following contributions: (1) Problem Formalization. We formalize the problem of temporal environmental variation in interacting dynamics modeling. (2) Novel Methodology. We analyze the problem under a probabilistic framework, and propose a novel approach CARE, which incorporates the continuous context variations and system states into a coupled ODE system. (3) Extensive Experiments. Extensive experiments conducted on four datasets validate the superiority of our CARE. The performance gain of our proposed CARE over the best baseline is up to  $36.35\%$ .
Interacting Dynamics Modeling. Deep learning approaches have been extensively used in recent years to model interacting systems across various fields [9, 53, 29, 20], including molecular dynamics and computational fluid dynamics. Early efforts focus on incorporating convolutional neural networks to learn from interacting regular grids [41]. To address more generalized scenarios, graph neural network (GNN) methods have been developed [42, 46, 45], leveraging message passing mechanisms to extract complex spatial signals. However, these methods often fail to account for environmental fluctuations, which hinders their ability to make reliable long-term predictions. In contrast, our CARE adopts a context-attended ODE architecture to explicitly represent both the observations and the underlying environment, enabling the generation of accurate future trajectories.
Neural Ordinary Differential Equations (ODEs). Drawing inspiration from the approximation of ResNet [5], neural ODEs equip neural networks with a continuous-depth architecture by parameterizing the derivatives of hidden states. Several attempts have been made to increase the expressiveness of neural ODEs [65], including adding regularization [8] and designing high-order ODEs [61]. Neural ODEs have also been incorporated into GNNs, producing continuous message passing layers to avoid oversmoothing [60] and increase the interpretability of predictions [70]. In this study, we employ a graph ODE architecture to capture the continuous nature of interacting system dynamics, relieving potential error accumulation caused by discrete prediction models.
Figure 1: An overview of the proposed CARE. To begin, we construct a temporal graph and utilize an encoder to initialize both the context variable and node representations from historical trajectories. Then, a coupled ODE model simulates the evolution of both nodes and context. Finally, CARE feeds node representations into decoders, which output the predicted trajectories at any timestamp.
Out-of-distribution (OOD) Generalization. OOD generation aims to make models more effective when the training and testing distributions diverge [50, 39, 57, 63]. This problem has drawn considerable attention in several areas, including text and vision [48]. One effective solution is to learn domain-invariant representations in the hidden space [31, 55, 30], which has been achieved under the guidance of invariant learning theory [6, 32]. Additionally, uncertainty modeling [34], causal learning [54, 14], and model selection [38, 56] are employed to improve the performance. Interacting systems inherently exhibit dynamic distribution variation caused by environmental changes, an aspect that remains underexplored in the literature. To address this, our paper proposes a novel approach named CARE to model context information from the perspective of a probabilistic model.
In a multi-agent dynamical system, the state at time  $t$  is represented by  $G^{t} = (V,E^{t},X^{t})$ , where each node in  $V$  corresponds to an object,  $E^{t}$  denotes the current edge set, and  $X^{t}$  signifies the node attribute matrix. Specifically, the state vector for each  $i\in V$  is given by  $\pmb{x}_i^t = [\pmb {p}_i^t,\pmb {q}_i^t,\pmb {a}_i]$ , with  $\pmb {p}_i^t\in \mathbb{R}^3$  and  $q_{i}^{t}\in \mathbb{R}^{3}$  representing the position and velocity, respectively, and  $a_i$  representing static attributes. We are given the sequence  $\{G^0,G^1,\dots ,G^t\}$  and aim to learn a model that produces the target dynamic states  $Y^{s}(s > t)$  (e.g., velocity), which are part of  $X^{s}$  at the corresponding time. The temporal environmental variation would result in data distribution changes during the evolution of interacting systems. If we utilize  $C^{0:t}$  to indicate the dynamical environment factor till timestamp  $t$ , we have data from variable distribution, i.e.,  $(G^{0:t},Y^{s})\sim P(G^{0:t},Y^{s}|C^{0:t})$ . Thus, we must take these changes into account for accurate trajectory predictions.
3.2 GNNs for Modeling Dynamical Systems
Graph neural networks (GNNs) are extensively employed in dynamical system modeling to investigate the interactive relationships between objects [42, 46, 45]. These methods typically use the current states to predict the states of nodes at the next timestamp. Specifically, omitting the time notation, GNNs first initialize node representations and edge representations using encoders:
$$
\boldsymbol {v} _ {i} ^ {(0)} = f ^ {v} \left(\boldsymbol {x} _ {i}\right), \quad \boldsymbol {e} _ {i j} ^ {(0)} = f ^ {e} \left(\boldsymbol {x} _ {i}, \boldsymbol {x} _ {j}\right), \tag {1}
$$
where  $f^v(\cdot)$  and  $f^e(\cdot)$  are two encoders for node and edge representations, respectively. Then, they utilize two propagation modules to update these representations at the  $l$ -th layer, i.e.,  $\pmb{v}_i^{(l)}$  and  $e_{ij}^{(l)}$  following the message passing mechanism:
$$
\boldsymbol {e} _ {i j} ^ {(l + 1)} = \phi^ {e} \left(\boldsymbol {v} _ {i} ^ {(l)}, \boldsymbol {v} _ {j} ^ {(l)}, \boldsymbol {e} _ {i j} ^ {(l)}\right), \quad \boldsymbol {v} _ {i} ^ {(l + 1)} = \phi^ {v} \left(\boldsymbol {v} _ {i} ^ {(l)}, \sum_ {j \in \mathcal {N} _ {i}} \boldsymbol {e} _ {i j} ^ {(l + 1)}\right), \tag {2}
$$
where  $\mathcal{N}_i$  collects the neighbours of node  $i$ .  $\phi^e (\cdot)$  and  $\phi^v (\cdot)$  are two functions for representation updating. Finally, they generate the target velocity vectors at the next timestamp using a decoder.
In this paper, we propose a novel method called CARE for modeling interacting dynamics under temporal environmental variation. We start by formalizing a probabilistic model to understand the relationships between trajectories and contexts. Based on this foundation, we construct a spatiotemporal encoder to initialize the representations of nodes and contexts. Then, to simultaneously model their evolution, a coupled graph ODE is introduced where node representations are evolved with the guidance of contexts and the states of the context variable are inferred from current trajectories. Additionally, we introduce a regularization term and dynamic graph updating strategies to enhance our framework. An illustration of our CARE can be found in Figure 1.
4.1 Probabilistic Model for System Dynamics under Temporal Distribution Drift
In this work, to tackle the challenge brought by temporal distribution drift, we inject a context variable  $c^t$  in our dynamic system modeling, which indicates the environment state at timestamp  $t$ . For example, the context variable could indicate flow speed, density and viscosity in fluid dynamics.
Here, we make two basic assumptions in our probabilistic model.
Assumption 4.1. (Independence-I) The context variable is independent of the sequences before the last observed timestamp, i.e.,  $P(\pmb{c}^t | \pmb{c}^{t-k}, G^{0:t}) = P(\pmb{c}^t | \pmb{c}^{t-k}, G^{t-k:k:t})$ , where  $t - k$  is the last observed timestamp.
Assumption 4.2. (Independence-II) Given the current states and contexts, the future trajectories are independent of the previous trajectories and contexts, i.e.,  $P(\mathbf{Y}^{t - k:t + l}|G^{0:t - k},\pmb{c}^{0:t - k}) = P(\mathbf{Y}^{t - k:t - k + l}|G^{t - k},\pmb{c}^{t - k})$  where  $l$  is the length of the prediction.
Then, we can have the following lemma:
Lemma 4.1. With Assumptions 4.1 and 4.2, we have:
$$
\mathrm {P} \left(\boldsymbol {Y} ^ {t} \mid G ^ {0: t - 1}\right) = \int \mathrm {P} \left(\boldsymbol {Y} ^ {t} \mid \boldsymbol {c} ^ {t - 1}, G ^ {t - 1}\right). \tag {3}
$$
$$
\mathrm {P} \left(\boldsymbol {c} ^ {t - 1} \mid \boldsymbol {c} ^ {t - k}, G ^ {t - k: t - 1}\right) \cdot \mathrm {P} \left(\boldsymbol {c} ^ {t - k} \mid G ^ {0: t - k}\right) d \boldsymbol {c} ^ {t - 1} d \boldsymbol {c} ^ {t - k}.
$$
The proof of Lemma 4.2 can be found in Appendix. From Lemma 4.1, we decompose the probability  $\mathrm{P}\left(\mathbf{Y}^t\mid G^{0:t - 1}\right)$  into three terms. Specifically, the last term necessitates encoding context information based on the historical trajectory  $G^{0:t - k}$ . The second term aims to update the context vector according to the recent trajectory  $G^{t - k:t - 1}$ . The first term suggests using the context variable in conjunction with the current states to make predictions for the next timestamp. Besides making a single next-time prediction, our model can also predict trajectories  $(\mathbf{Y}^{t - k},\mathbf{Y}^{t - k + 1},\dots ,\mathbf{Y}^t)$  by modifying Eqn. 14.
Consequently, we divide each training sequence into two parts, namely  $[0,t - k]$  and  $(t - k,t]$  as in [19, 18]. The first part is used to encode contexts and nodes in the system, while the second part serves for updating and prediction purposes.
4.2 Context Acquisition from Spatio-temporal Signals
In this part, our goal is to acquire knowledge from the historical trajectory, i.e.,  $\{G^0,\dots ,G^{t - k}\}$  to encode contexts and nodes for initialization. To be specific, we first construct a temporal graph to capture spatial and temporal signals simultaneously. Subsequently, we employ the attention mechanism to update temporal node representations, which will be used to initialize the context representation and node representations for  $\{G^{t - k},\dots ,G^T\}$ .
To begin, we construct a temporal graph containing two types of edges, i.e., spatial and temporal edges. Spatial edges are built when the distance between two nodes at the same timestamp is less than a threshold while temporal edges are between every two consecutive observations for each node. Specifically, in the constructed temporal graph  $G^{H}$ , there are  $N(t - k + 1)$  nodes  $\{i_s\}^{s \in [0:t - k], i \in V}$
in total. The adjacent matrix  $A$  contains both spatial and temporal edges as follows:
$$
\boldsymbol {A} \left(i ^ {s}, j ^ {s ^ {\prime}}\right) = \left\{ \begin{array}{l l} e x p \left(- d ^ {s} (i, j)\right) & s = s ^ {\prime}, e x p \left(- d ^ {s} (i, j)\right) <   \tau , \\ 1 & i = j, s ^ {\prime} = s + 1, \\ 0 & \text {o t h e r w i s e}, \end{array} \right. \tag {4}
$$
where  $d^{s}(i,j)$  denotes the distance between particles  $i$  and  $j$  at timestamp  $s$  and  $\tau$  is the predefined threshold. Then, we utilize an attention-based GNN to extract spatio-temporal relationships into node representations. Here, we first compute the interaction scores between each node in  $G^{H}$  with its neighboring nodes, and then aggregate their embeddings at the previous layer. Let  $d$  represent the hidden dimension, the interaction score between  $i^{s}$  and  $j^{s^{\prime}}$  at layer  $l$  is:
$$
w ^ {(l)} \left(i ^ {s}, j ^ {s ^ {\prime}}\right) = \boldsymbol {A} \left(i ^ {s}, j ^ {s ^ {\prime}}\right) \left(\boldsymbol {W} _ {\text {q u e r y}} \boldsymbol {h} _ {i} ^ {s, (l)}\right) \star \left(\boldsymbol {W} _ {\text {k e y}} \boldsymbol {h} _ {j} ^ {s ^ {\prime}, (l)}\right), \tag {5}
$$
where  $\mathbf{W}_{query} \in \mathbb{R}^{d \times d}$  and  $\mathbf{W}_{key} \in \mathbb{R}^{d \times d}$  are two matrices to map temporal node representations into different spaces.  $\star$  computes the cosine similarity between two vectors. With the interaction scores, we can compute temporal node representations at the next layer:
$$
\boldsymbol {h} _ {i} ^ {s, (l + 1)} = \boldsymbol {h} _ {i} ^ {s, (l)} + \sigma \left(\sum_ {j ^ {s ^ {\prime}} \in \mathcal {N} (i ^ {s})} w ^ {(l)} \left(i ^ {s}, j ^ {s ^ {\prime}}\right) \boldsymbol {W} _ {\text {v a l u e}} \boldsymbol {h} _ {j} ^ {s ^ {\prime}, (l)}\right), \tag {6}
$$
where  $W_{value}$  is a weight matrix,  $\mathcal{N}(i^s)$  collects all the neighbours of  $i^s$  and  $\sigma(\cdot)$  is an activation function. After stacking  $L$  layers, we add temporal encoding and then summarize all these temporal node representations to initialize node representations for the upcoming ODE module:
$$
\boldsymbol {q} _ {i} ^ {s} = \boldsymbol {h} _ {i} ^ {s, (L)} + \operatorname {T E} (s), \quad \boldsymbol {u} _ {i} ^ {t - k} = \frac {1}{t - k + 1} \sum_ {s = 0} ^ {t - k} \sigma \left(\boldsymbol {W} _ {\text {s u m}} \boldsymbol {q} _ {i} ^ {s}\right), \tag {7}
$$
where  $\mathrm{TE}(s)[2i] = \sin \left(\frac{s}{10000^{2i / d}}\right)$ ,  $\mathrm{TE}(s)[2i + 1] = \cos \left(\frac{s}{10000^{2i / d}}\right)$  and  $W_{sum}$  denotes a projection matrix. The initial context variable  $c^{t - k}$  is driven by summarizing all node representations:
$$
\beta_ {i} ^ {t} = \operatorname {t a n h} \left(\left(\frac {1}{| V |} \sum_ {i ^ {\prime} \in V} \boldsymbol {u} _ {i ^ {\prime}} ^ {t - k}\right) \boldsymbol {W} _ {\text {c o n t e x t}}\right) \cdot \boldsymbol {u} _ {i} ^ {t - k}, \quad \boldsymbol {c} ^ {t - k} = \sum_ {i \in V} \beta_ {i} ^ {t} \boldsymbol {u} _ {i} ^ {t - k}, \tag {8}
$$
where  $W_{\text{context}}$  is a learnable matrix and  $\beta_i^t$  calculates the attention score for each node.
4.3 Context-attended Graph ODE
In this module, to model continuous evolution, we incorporate an ODE system into our approach. The precondition requires assuming that both the context variable and node representations are continuous to fit neural ODE models, which inherently holds for common dynamical systems in practice. We then introduce coupled ODEs to model the dynamic evolution of node representations and the context variable. Specifically, the context variable can be inferred during the evolution of node representations, which in turn drives the evolution of the system. We first introduce the assumption:
Assumption 4.3. (Continuous) We assume that both context variable  $c^s$  and node representations  $v_i^s$  are continuously differentiable with respect to  $s$ .
Then, to utilize the context variable and the current state for making future predictions, we introduce a graph ODE model. Let  $\hat{A}^s$  denote the adjacency matrix at timestamp  $s$  with self-loop, we have:
$$
\frac {d \boldsymbol {v} _ {i} ^ {s}}{d s} = \Phi \left(\left[ \boldsymbol {v} _ {1} ^ {s}, \dots , \boldsymbol {v} _ {N} ^ {s}, \boldsymbol {c} ^ {s} \right]\right) = \sigma \left(\sum_ {j \in \mathcal {N} ^ {s} (i)} \frac {\hat {\boldsymbol {A}} _ {i j} ^ {s}}{\sqrt {\hat {\boldsymbol {D}} _ {i} ^ {s} \cdot \hat {\boldsymbol {D}} _ {j} ^ {s}}} \boldsymbol {v} _ {j} ^ {s} \boldsymbol {W} _ {1} + \boldsymbol {c} ^ {s} \boldsymbol {W} _ {2}\right), \tag {9}
$$
where  $\mathcal{N}^s (i)$  denotes the neighbours of node  $i$  at timestamp  $s$  and  $\hat{D}_i^s$  represents the degree of node  $i$  according to  $\hat{A}^s$ . The first term in Eqn. 9 aggregates information from its instant neighbors and the second term captures information from the current context information.
The next question is how to model the evolution of  $c^s$ . Notice that we have:
$$
\mathrm {P} \left(\boldsymbol {c} ^ {t} \mid \boldsymbol {c} ^ {t - k}, G ^ {t - k: t}\right) = \int P \left(\boldsymbol {c} ^ {t} \mid \boldsymbol {c} ^ {t - \Delta t}, G ^ {t - \Delta t: t}\right) \dots \tag {10}
$$
$$
P \left(\boldsymbol {c} ^ {t - k + \Delta t} \mid \boldsymbol {c} ^ {t - k}, G ^ {t - k: t - k + \Delta t}\right) d \boldsymbol {c} ^ {t - k + \Delta t} \dots d \boldsymbol {c} ^ {t - \Delta t},
$$
where  $\Delta t$  denotes a small time interval. With Assumption 4.3, we can simplify  $P(c^{t - k + \Delta t}|c^{t - k},G^{t - k:t - k + \Delta t})$  into  $P(c^{t - k + \Delta t}|c^{t - k},V^{t - k},dV^{t - k})$  where  $V^{t - k}$  denotes the node embedding matrix at timestamp  $t - k$  and  $dV^{t - k}$  is its differentiation. On this basis, we introduce another ODE to update the context variable as:
$$
\frac {d \boldsymbol {c} ^ {s}}{d s} = \Phi^ {c} (\operatorname {A G G} \left(\left\{\boldsymbol {v} _ {i} ^ {s} \right\} _ {i \in V}\right), \operatorname {A G G} \left( \right.\left\{\frac {d \boldsymbol {v} _ {i} ^ {s}}{d s} \right\} _ {i \in V}), \boldsymbol {c} ^ {s} \left. \right]\left. \right), \tag {11}
$$
where  $\Phi^c$  is an MLP with the concatenated input and  $\mathrm{AGG}(\cdot)$  is an operator to summarize node representations such as averaging and sum. Compared to previous methods, the key to our CARE is to take into account the mutual impact between the environment and the trajectories, and model their evolution simultaneously by coupling Eqn. 9 and Eqn. 11. We also provide a theoretical analysis of the uniqueness of the solution to our system. To simplify the analysis, we set  $\mathrm{AGG}(\cdot)$  to summation and rewrite Eqn. 11 with learnable matrices  $\mathbf{W}_3$ ,  $\mathbf{W}_4$  and  $\mathbf{W}_5$  as:
$$
\frac {d \boldsymbol {c} ^ {s}}{d s} = \sigma \left(\sum_ {i = 1} ^ {N} \left(\boldsymbol {v} _ {i} ^ {s} \boldsymbol {W} _ {3} + \frac {d \boldsymbol {v} _ {i} ^ {s}}{d s} \boldsymbol {W} _ {4}\right) + \boldsymbol {c} ^ {s} \boldsymbol {W} _ {5}\right). \tag {12}
$$
Then, we introduce the following assumption:
Assumption 4.4. All time-dependent coefficients in Eqn. 9, i.e  $A_{ij}^{t}, \hat{D}_{i}^{t}$  are continuous with respect to  $t$  and bounded by a constant  $C > 0$ . All parameters in the weight matrix are also bounded by a constant  $W > 0$ .
With Assumption 4.4, we can deduce the following lemma:
Lemma 4.2. Given the initial state  $(t_0, \pmb{v}_1^{t_0}, \dots, \pmb{v}_N^{t_0}, \pmb{c}^{t_0})$ , we claim that there exists  $\varepsilon > 0$ , s.t. the ODE system 9 and 12 has a unique solution in the interval  $[t_0 - \varepsilon, t_0 + \varepsilon]$ .
The proof of Lemma 4.2 can be found in Appendix. Our theoretical analysis indicates that at least locally, the future system trajectory and context information are predictable based on their historical values [51], which is also an important property for dynamical system modeling [5, 28].
4.4 Decoder and Optimization