TESTAM: A Time-Enhanced Spatio-Temporal Attention Model with Mixture of Experts

Let us define a road network as $\mathcal{G}=(\mathcal{V},\mathcal{E},\mathcal{A})$, where $\mathcal{V}$ is a set of all roads in road networks with $|\mathcal{V}|=N$, $\mathcal{E}$ is a set of edges representing the connectivity between roads, and $\mathcal{A}\in\mathbb{R}^{N\times N}$ is a matrix representing the topology of $\mathcal{G}$. Given road networks, we formulate our problem as a special version of multivariate time series forecasting that predicts future $T$ graph signals based on $T^{\prime}$ historical input graph signals:

where $X_{\mathcal{G}}^{(i)}\in\mathbb{R}^{N\times C}$, $C$ is the number of input features. We aim to train the mapping function $f(\cdot):\mathbb{R}^{T^{\prime}\times N\times C}\rightarrow\mathbb{R}^{T\times N\times C}$, which predicts the next $T$ steps based on the given $T^{\prime}$ observations. For the sake of simplicity, we omit $\mathcal{G}$ from $X_{\mathcal{G}}$ hereinafter.

To effectively forecast the traffic signals, we first discuss spatial modeling, which is one of the necessities for traffic data modeling. In traffic forecasting, we can classify spatial modeling methods into four categories: 1) with identity matrix (i.e., multivariate time-series forecasting), 2) with a pre-defined adjacency matrix, 3) with a trainable adjacency matrix, and 4) with attention (i.e., dynamic spatial modeling without prior knowledge). Conventionally, a graph topology $\mathcal{A}$ is constructed via an empirical law, including inverse distance (Li et al., 2018; Yu et al., 2018) and cosine similarity (Geng et al., 2019). However, these empirically built graph structures are not necessarily optimal, thus often resulting in poor spatial modeling quality. To address this challenge, a line of research (Wu et al., 2019; Bai et al., 2020; Jiang et al., 2023) is proposed to capture the hidden spatial information. Specifically, a trainable function $g(\cdot,\theta)$ is used to derive the optimal topological representation $\tilde{\mathcal{A}}$ as:

where $g(X^{(t)},\theta)\in\mathbb{R}^{N\times e}$, and $e$ is the embedding size. Spatial modeling based on Eq. 1 can be classified into two subcategories according to whether $g(\cdot,\theta)$ depends on $X^{(t)}$. Wu et al. (2019) define $g(\cdot,\theta)=E\in\mathbb{R}^{N\times e}$, which is time-independent and less noise-sensitive but less in-situ modeling. Cao et al. (2020); Zhang et al. (2020) propose time-varying graph structure modeling with $g(H^{(t)},\theta)=H^{(t)}W$, where $W\in\mathbb{R}^{d\times e}$, projecting hidden states to another embedding space. Ideally, this method models dynamic changes in graph topology, but it is noise-sensitive.

To reduce noise-sensitivity and obtain a time-varying graph structure, Zheng et al. (2020) adopt a spatial attention mechanism for traffic forecasting. Given input $H_{i}$ of node $i$ and its spatial neighbor $\mathcal{N}_{i}$, they compute spatial attention using multi-head attention as follows:

where $W^{O}$ is a projection layer, $d_{k}$ is a dimension of key vector, and $f^{(k)}_{q}(\cdot),f^{(k)}_{k}(\cdot),$ and $f^{(k)}_{v}(\cdot)$ are query, key, and value projections of the $k$-th head, respectively. Although effective, these attention-based approaches still suffer from irregular spatial modeling, such as less accurate self-attention (i.e., from node $i$ to $i$) (Park et al., 2020) and uniformly distributed uninformative attention, regardless of its spatial relationships (Jin et al., 2023).

Since temporal features (e.g., time of day) work as a global position with a specific periodicity, we omit position embedding in the original transformer architecture. Furthermore, instead of normalized temporal features, we utilize Time2Vec embedding (Kazemi et al., 2019) for periodicity and linearity modeling. Specifically, for the temporal feature $\tau\in\mathbb{N}$, we represent $\tau$ with $h$-dimensional embedding vector $v(\tau)$ and the learnable parameters $w_{i},\phi_{i}$ for each embedding dimension $i$ as below:

where $\mathcal{F}$ is a periodic activation function. Using Time2Vec embedding, we enable the model to utilize the temporal information of labels. Here, temporal information embedding of an input sequence is concatenated with other input features and then projected onto the hidden size $h$.

As temporal attention in TESTAM is the same as that of transformers, we describe the benefits of temporal attention. Recent studies (Li et al., 2018; Bai et al., 2020) have shown that attention is an appealing solution for temporal modeling because, unlike recurrent unit-based or convolution-based temporal modeling, it can be used to directly attend to features across time steps with no restrictions. Temporal attention allows parallel computation and is beneficial for long-term sequence modeling. Moreover, it has less inductive bias in terms of locality and sequentiality. Although strong inductive bias can help the training, less inductive bias enables better generalization. Furthermore, for the traffic forecasting problem, causality among roads is an unavoidable factor (Jin et al., 2023) that cannot be easily modeled in the presence of strong inductive bias, such as sequentiality or locality.

In this work, we leverage three spatial modeling layers for each expert, as shown in the middle of Fig. 1: spatial modeling with an identity matrix (i.e., no spatial modeling), spatial modeling with a learnable adjacency matrix (Eq. 1), and spatial modeling with attention (Eq. 2 and Eq. 3). We calculate spatial attention using Eqs. 2 and 3. Specifically, we compute attention with $\forall_{i\in\mathcal{V}},\mathcal{N}_{i}=\mathcal{V}$, which means attention with no spatial restrictions. This setting enables similarity-based attention, resulting in better generalization.

Inspired by the success of memory-augmented graph structure learning (Jiang et al., 2023; Lee et al., 2022), we propose a modified meta-graph learner that learns prototypes from both spatial graph modeling and gating networks. Our meta-graph learner consists of two individual neural networks with a meta-node bank $\mathbf{M}\in\mathbb{R}^{m\times e}$, where $m$ and $e$ denote total memory items and a dimension of each memory, respectively, a hyper-network (Ha et al., 2017) for generating node embedding conditioned on $\mathbf{M}$, and gating networks to calculate the similarities between experts’ hidden states and queried memory items. In this section, we mainly focus on the hyper-network. We construct a graph structure with a meta-node bank $\mathbf{M}$ and a projection $W_{E}\in\mathbb{R}^{e\times d}$ as follows:

By constructing a memory-augmented graph, the model achieves better context-aware spatial modeling than that achieved using other learnable static graphs (e.g., graph modeling with $E\in\mathbb{R}^{N\times d}$). Detailed explanations for end-to-end training and meta-node bank queries are provided in Sec. 3.3.

To eliminate the error propagation effects caused by auto-regressive characteristics, we propose a time-enhanced attention layer that helps the model transfer its domain from historical $T^{\prime}$ time steps (i.e., source domain) to next $T$ time steps (i.e., target domain). Let $\mathbf{\tau}^{(t)}_{label}=[\tau^{(t+1)},\dots,\tau^{(t+T)}]$ be a temporal feature vector of the label. We calculate the attention score from the source time step $i$ to the target time step $j$ as:

where $d_{k}=d/K$, $K$ is the number of heads, and $W_{q}^{(k)},W_{k}^{(k)}$ are linear transformation matrices. We can calculate the attention output using the same process as in Eq. 2, except that time-enhanced attention attends to the time steps of each node, whereas Eq. 2 attends to the important nodes at each time step.

To enable fine-grained routing that fits the regression problem, we adopt two classification losses: a classification loss to avoid the worst routing and another loss function to find the best routing. Inspired by SMoEs, we define the worst routing avoidance loss as the cross entropy loss with pseudo label $l_{e}$ as shown below:

where $\hat{y}$ is the output of the selected expert, and $q$ is an error quantile. If an expert is incorrectly selected, its label becomes zero and the unselected experts have the pseudo label $1/(E-1)$, which means that there are equal chances of choosing unselected experts.

We also propose the best-route selection loss for more precise routing. However, as traffic data are noisy and contain many nonstationary characteristics, the best-route selection is not an easy task. Therefore, instead of choosing the best routing for every time step and every node, we calculate node-wise routing. Our best-route selection loss is similar to that in Eq. 6, except that it calculates node-wise pseudo labels and the routing probability, and the condition for pseudo labels is changed from “$L(y,\hat{y})$ is greater/smaller than $q$-th quantile” to “$L(y,\hat{y})$ is greater/smaller than $(1-q)$-th quantile.” Detailed explanations are provided in Appendix A.

It uses only the output of the attention experts without ensembles or any other gating mechanism. Memory items are not trained because there are no gradient flows for the adaptive expert or gating networks. This setting results in an architecture similar to that of GMAN.

Instead of using MoEs, the final output is calculated with the weighted summation of the gating networks and each expert’s output. This setting allows the use of all spatial modeling methods but no in-situ modeling.

It excludes the loss for guiding best route selection. The exclusion of this relatively coarse-grained loss function is based on the fact that coarse-grained routing tends not to change its decisions after initialization (Dryden & Hoefler, 2022).

It does not exclude any components. Instead, it replaces identity expert with a GCN-based adaptive expert, reducing spatial modeling diversity. The purpose of this setting is to test the hypothesis that in-situ modeling with diverse graph structures is helpful for traffic forecasting.

It replaces temporal information embedding (TIM) with simple embedding vectors without periodic activation functions.

It replaces time-enhanced attention with basic temporal attention as we described in Sec. 3.2.

The experimental results shown in Table 2 connote that our hypotheses are supported and that TESTAM is a complete and indivisible set. The results of “w/o gating” and “ensemble” suggest that in-situ modeling greatly improves the traffic forecasting quality. The “w/o gating” results indicate that the performance improvement is not due to our model but due to in-situ modeling itself since this setting lead to performance comparable to that of GMAN (Zheng et al., 2020). “worst-route avoidance only” results indicate that our hypothesis that both of our routing classification losses are crucial for proper routing is valid. Finally, the results of “replaced,” which indicate significantly worse performance even than “worst route avoidance only,” confirm the hypothesis that diverse graph structures is helpful for in-situ modeling. Additional qualitative results with examples are provided in Appendix C.