ST-LINK: Spatio-Temporal Forecasting
- ST-LINK is a spatio-temporal forecasting framework that extends pre-trained transformers by incorporating SE-Attention and MRFFN to capture spatial dependencies and long-term temporal dynamics.
- It replaces standard attention and feed-forward layers with spatially-enhanced modules and a retrieval-augmented mechanism, optimizing prediction accuracy for traffic data.
- Empirical results on demand and speed forecasting benchmarks demonstrate ST-LINKâs efficiency gains and lower MAE and RMSE compared to conventional deep learning and LLM approaches.
ST-LINK is a framework for spatio-temporal forecasting that builds on a pre-trained transformer LLM and modifies its trainable layers to better capture spatial dependencies and long-term temporal structure in traffic data. It is introduced for traffic forecasting within intelligent transportation systems, where standard LLMs are promising but are limited by their sequential token design and their architectural incompatibility with graph-structured spatial data. ST-LINK addresses these issues through two componentsâSpatially-Enhanced Attention (SE-Attention) and the Memory Retrieval Feed-Forward Network (MRFFN)âand is evaluated on demand and speed forecasting benchmarks, where it is reported to surpass conventional deep learning and LLM approaches (Jeon et al., 17 Sep 2025).
1. Architectural organization
ST-LINK builds on a pre-trained transformer LLM, using GPT-2 in the reported experiments, and partitions its layers into two groups. The bottom layers are frozen, with all parameters fixed except LayerNorm, while the top layers are trainable. In each of these trainable layers, the ordinary multi-head attention and feed-forward network are replaced by SE-Attention and MRFFN, respectively (Jeon et al., 17 Sep 2025).
The overall pipeline is specified as follows:
$[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$
After the last block, a small regression head maps hidden states at each node and time step to future traffic values (Jeon et al., 17 Sep 2025).
This organization preserves the LLMâs inherent sequential processing structure while inserting spatial and retrieval-based mechanisms directly into the trainable layers. A plausible implication is that the design aims to retain the benefits of a pre-trained language-model backbone while limiting the amount of task-specific fine-tuning.
2. Spatially-Enhanced Attention
SE-Attention extends rotary positional embeddings to encode both temporal and graph-structured spatial information. Standard RoPE applies a rotation to each 2-dimensional subvector of queries and keys using a single time index :
with
SE-Attention introduces a second rotation for each token that depends on its node index. For the -th spatial location, with learnable scalar embedding , the spatial angles are defined as
0
Two 1-dimensional query vectors are then formed, and similarly for keys: a temporally encoded vector 2 using the time index 3, and a spatially encoded vector 4, which applies the rotation to the first 5 dimensions of 6 and leaves the rest unchanged. These are concatenated and projected back to 7 dimensions:
8
The same construction is used for 9 (Jeon et al., 17 Sep 2025).
The attention computation replaces the standard query and key representations with these fused forms:
0
Because 1 fuses two orthogonal rotations, the resulting dot product depends jointly on 2 and on relative node relationships induced by the learned 3. In effect, the model can infer graph-structured spatial dependencies alongside temporal correlations without an explicit adjacency matrix (Jeon et al., 17 Sep 2025).
The stated significance of SE-Attention is that it places unified spatio-temporal encoding directly inside each attention head. In the reported analysis, this component learns smooth spatial correlations via the learned node-dependent rotations and excels at regular daily traffic flows (Jeon et al., 17 Sep 2025).
3. Memory Retrieval Feed-Forward Network
MRFFN replaces the vanilla FFN with a retrieval-augmented mixture-of-experts that dynamically recalls historical spatio-temporal patterns. The memory module maintains 4 key-value pairs 5. Given an input hidden vector 6, similarity scores are computed, the top-7 keys are selected, and the retrieved representation is formed as
8
To adapt memory to new patterns, each retrieved key is updated by a small step toward the current input:
9
The retrieved vector is then integrated through a mixture-of-experts mechanism. The concatenated vector 0 is passed through a gating network 1 to produce expert weights 2. With 3 small FFNs 4, the output is
5
This dynamic routing allows the network to specialize its feed-forward processing based on the retrieved historical patterns, and the framework describes this as bolstering long-term temporal stability (Jeon et al., 17 Sep 2025).
In the empirical discussion, MRFFN is associated with greater gains on long horizons, and its retrieval of similar historical contexts is reported to allow rapid adaptation to sudden traffic drops or surges, including Sensor-10 and Sensor-48 case studies on METR-LA (Jeon et al., 17 Sep 2025). A potential misconception would be to treat MRFFN as an external datastore mechanism; the reported characterization is instead end-to-end memory learning with no external datastore.
4. Training and inference procedure
The data preprocessing pipeline begins by aggregating raw sensor or station readings into graph signals 6. RevIN statistics are computed per node over the input window, and each nodeâs sequence is normalized. Learnable temporal embeddings for each of the 7 time steps and node embeddings for each of the 8 nodes are added, after which the inputs are projected into an initial hidden dimension 9 (Jeon et al., 17 Sep 2025).
Training minimizes mean squared error between predictions and ground truth for the next $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$0 steps:
$[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$1
The optimizer is AdamW with learning rate $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$2 for demand or $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$3 for speed, weight decay $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$4, and dropout $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$5â$[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$6. The batch size is $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$7. The train/validation/test split is $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$8 for demand or $[\text{Input features + temporal %%%%45%%%% spatial embeddings}] \rightarrow \text{RevIN normalization} \rightarrow \text{linear projections to }Q,K,V \rightarrow \text{SE-Attention} \rightarrow \text{revIN denormalization %%%%45%%%% residual} \rightarrow \text{MRFFN} \rightarrow \text{residual %%%%45%%%% layer norm} \rightarrow \text{next layer}.$9 for speed. Only the top 0 layers are updated; the bottom layers remain frozen except LayerNorm (Jeon et al., 17 Sep 2025).
During inference, a window of length 1 is slid over historical data. For each window, the data are normalized, embedded, and passed through the model with SE-Attention and MRFFN, with the memory module active. The system outputs 2-step predictions, and no further fine-tuning or prompt-engineering is required (Jeon et al., 17 Sep 2025).
The procedure is notable for combining parameter-efficient fine-tuning with a retrieval-enabled forward pass. A plausible implication is that this setup attempts to separate general sequence modeling retained from pre-training from spatio-temporal adaptation learned in the upper layers.
5. Benchmarks and empirical results
The reported experiments cover both demand forecasting and speed forecasting. Demand forecasting uses NYCTaxi with 266 stations and Citi Bike with 250 stations, aggregated every 5 minutes and predicting the next 12 steps. Speed forecasting uses METR-LA with 207 sensors and PEMS-BAY with 325 sensors, also at 5-minute intervals and with 12-step prediction (Jeon et al., 17 Sep 2025).
The baselines are grouped into graph-based, attention-based, and LLM-based methods. The graph-based baselines are DCRNN, STGCN, AGCRN, GWN, ASTGCN, ASTGNN, STG-NCDE, DGCRN, D3STGNN, and STIDGCN. The attention-based baselines are GMAN, PDFormer, STAEFormer, STTN, PM-MemNet, MegaCRN, and TESTAM. The LLM-based baselines are OFA, GATGPT, and ST-LLM (Jeon et al., 17 Sep 2025).
| Setting | Baseline comparison | ST-LINK |
|---|---|---|
| NYCTaxi Pickup | DCRNN: 5.40; ST-LLM: 5.29 | 5.09 |
| Citi Bike Drop-off | GWN: 1.95; ST-LLM: 1.89 | 1.87 |
| METR-LA@12 | OFA: 3.66; ST-LLM: 3.58 | 3.51 |
| PEMS-BAY@12 | ST-LLM: 1.95 | 1.90 |
On NYCTaxi Pickup, the result is described as a 4% relative improvement over ST-LLM. Across all tasks and horizons, ST-LINK is reported to consistently attain the lowest MAE and RMSE, particularly on long horizons where MRFFNâs retrieval gives greater gains (Jeon et al., 17 Sep 2025).
The reported efficiency results are also specific. ST-LINK is described as parameter-efficient fine-tuning with only approximately 2.2 M trainable parameters versus 42 M in ST-LLM, yielding approximately 38% faster training and approximately 37% faster inference (Jeon et al., 17 Sep 2025). These figures are presented as strengths of the design rather than as a claim that the full model is small in absolute terms.
6. Interpretation, limitations, and prospective extensions
The analysis distinguishes between regular and abrupt traffic patterns. SE-Attention is said to learn smooth spatial correlations via learned node-dependent rotations, which is associated with regular daily traffic flows. MRFFNâs retrieval of similar historical contexts is said to allow rapid adaptation to sudden traffic drops or surges, with Sensor-10 and Sensor-48 case studies on METR-LA used as examples (Jeon et al., 17 Sep 2025).
Several strengths are explicitly identified: unified spatio-temporal encoding directly inside each attention head, end-to-end memory learning with no external datastore, and parameter-efficient fine-tuning relative to ST-LLM (Jeon et al., 17 Sep 2025). These points clarify that the framework is neither a purely graph-neural architecture nor a minimally modified sequential LLM; instead, it inserts spatial encoding and memory retrieval into the standard transformer block structure.
The limitations are also explicit. Memory slots may suffer catastrophic forgetting if domain shifts occur too abruptly, and the EMA rate 4 requires careful tuning. In addition, scaling to much larger LLMs, such as billions of parameters, may incur heavy memory and compute unless sparsity or expert routing is further optimized (Jeon et al., 17 Sep 2025). This suggests that the frameworkâs retrieval mechanism is not presented as a universally stable memory solution under all nonstationary regimes.
The proposed future improvements are fourfold: adaptive memory-update schedules or regularization to balance plasticity versus stability; hierarchical multi-scale SE-Attention to capture both local and global spatial relations; integration of graph-topology learning, such as learnable adjacency, alongside SE-Attention to refine node embeddings; and evaluation in other spatio-temporal domains, including weather forecasting and disease spread (Jeon et al., 17 Sep 2025). These directions indicate that the current formulation emphasizes traffic forecasting, while the underlying mechanisms are intended to be extensible to broader spatio-temporal prediction settings.