Pre-training Enhanced Spatial-temporal Graph Neural Network for Multivariate Time Series Forecasting

In plain words

Traffic sensors across a city send readings every few minutes, and the sensors influence each other: a jam on one road shows up on the roads that feed it. Good forecasters use both the time pattern of each sensor and the links between sensors. The usual models look only at a short recent window, because the models that link sensors get expensive with long inputs.

This paper splits the job in two. First, a separate network learns from a very long history by hiding most of it and trying to fill the gaps back in. Its short summaries of each stretch of history then tell a standard forecaster what the wider context looks like. Second, the same summaries help the model guess which sensors are related, even when no road map is given. The forecaster itself is an off-the-shelf graph model that takes these extra inputs.

The first network is trained once per dataset and then frozen. The forecaster is trained on top of it, so the cost of the long history is paid in advance.

TL;DR

• STEP adds a frozen, pre-trained time-series Transformer (TSFormer) to a spatial-temporal graph forecaster; the paper's backend is Graph WaveNet §3.2.
• TSFormer is trained by masked patch reconstruction (75% of patches hidden) on a long history, and it supplies a representation of the last short patch and a kNN-based regularizer for a learned sparse graph §3.1 §3.2.
• Author-reported: best MAE, RMSE and MAPE in nearly all cells on METR-LA, PEMS-BAY and PEMS04 at horizons 3, 6 and 12, with a t-test at p < 0.05 §4.2 tab. 2.
• The code differs from the paper in several places that change the learned graph; see Paper vs code and Pitfalls before reimplementing.

Problem

Multivariate traffic series have several periodicities (daily and weekly), and neighbouring sensors are correlated §1. Most spatial-temporal GNNs only read a short window, such as one hour of 5-minute readings, because their cost grows with the input length. That short window cannot show the wider context, and it is a noisy basis for estimating how sensors depend on each other §1. Many datasets also lack a reliable predefined graph, and learning one jointly with the forecaster is hard because there is no direct supervision for the graph §3.2.

Method

Overview. STEP has two stages §3. In the pre-training stage, TSFormer reads a long history (one week for METR-LA and PEMS-BAY, two weeks for PEMS04 in the experiments) and is trained to reconstruct hidden patches §3.1 §4.1. In the forecasting stage, the TSFormer encoder is frozen. Its output is used in three ways: (a) the representation of the last short patch is added to the backend's latent features through a small projector eq. 5; (b) a discrete sparse graph is sampled from learned Bernoulli parameters eq. 2 eq. 4; (c) a kNN graph built from TSFormer states regularises those parameters eq. 3. The backend and graph learner are trained end to end with a regression loss plus a weighted graph loss eq. 6 eq. 7 §3.2.

Pre-training (TSFormer).

1. Each series is cut into P non-overlapping patches of length L = 12 §3.1. Patches are embedded linearly into width d = 96 eq. 1.
2. Learnable positional embeddings are added to all patches, before masking §3.1 tsformer.py.
3. A random subset of 75% of the patches is hidden. The paper states the ratio, not whether the subset is drawn per sample §3.1. The code draws one subset per forward pass and reuses it for every series and every sample in the batch mask.py tsformer.py.
4. The encoder sees only the visible patches. It is a stack of 4 Transformer blocks followed by a LayerNorm §3.1 §4.1 tsformer.py.
5. The decoder maps the encoder output to width d, inserts a shared mask token at every hidden position, and runs a single Transformer block §3.1 §4.1 tsformer.py.
6. A linear layer maps each decoded patch back to L values §3.1 tsformer.py.
7. The loss is the mean absolute error over the hidden patches only §3.1 eq. 1 TSFormer_METR-LA.py.

# TSFormer pre-training step  [§3.1]
# Code: step/step_arch/tsformer/tsformer.py#L179-L188
# Input:  series (B, N, P*L) normalised values; P patches of length L = 12
# Output: scalar loss (MAE over hidden patches)

def pretrain_loss(series):
    # Split each series into P non-overlapping patches  [§3.1]
    patches = split_patches(series)              # (B, N, P, L)
    # Linear patch embedding to width d = 96  [eq. 1]
    tokens = embed(patches)                      # (B, N, P, d)
    # Add learnable positional embeddings to all patches  [§3.1]
    tokens = tokens + position_embedding         # (B, N, P, d)
    # Hide 75 percent of patches; one draw for the whole batch
    # [§3.1] [code: step/step_arch/tsformer/mask.py#L15-L24]
    visible, hidden = random_split(P, ratio=0.75)
    # Encode only the visible patches, 4 Transformer blocks  [§3.1]
    latent = encoder(tokens[:, :, visible])      # (B, N, P_vis, d)
    # Project to decoder width; insert shared mask tokens at hidden slots
    # Code also adds positional embeddings to the mask tokens here
    # [code: step/step_arch/tsformer/tsformer.py#L122-L126]
    full = insert_mask_tokens(project(latent), hidden)  # (B, N, P, d)
    # One decoder block, then a linear map to L values per patch  [§3.1]
    recon = output_linear(decoder(full))         # (B, N, P, L)
    # MAE over hidden patches only, zero targets ignored  [§3.1]
    # [code: basicts/metrics/mae.py#L17-L23]
    return masked_mae(recon[:, :, hidden], patches[:, :, hidden])

Forecasting (STEP).

1. The TSFormer weights are loaded from the pre-trained checkpoint and frozen §3.2 step.py. The encoder runs without masking on the long history tsformer.py.
2. The graph learner builds a static global feature per series from the training split. It uses a small 1-D CNN and a linear layer §3.2 eq. 2. The paper's text also adds a dynamic term from the TSFormer states; the code removes that term discrete_graph_learning.py.
3. Edge logits Θ_ij come from two linear layers over the concatenated node features of i and j, giving two numbers per edge eq. 2.
4. A discrete graph A is sampled per edge with the Gumbel-Softmax trick eq. 4. Self-loops are removed discrete_graph_learning.py.
5. A kNN graph is built from cosine similarities of the full TSFormer output, concatenated over patches §3.2 discrete_graph_learning.py.
6. The backend receives the short 12-step input and the sampled graph, and returns its latent features §3.2 model.py.
7. The last patch's TSFormer state goes through a 2-layer MLP projector, and the result is added to the backend's latent features eq. 5 §3.2 step.py model.py.
8. The regression head predicts the next 12 steps. The loss is MAE on the forecast plus a graph term weighted by lambda eq. 6 eq. 7 §3.2 step_loss.py.

# STEP forecasting step  [§3.2] [code: step/step_arch/step.py#L57-L72]
# Input:  short (B, 12, N, C); long (B, T_long, N, C); epoch
# Output: forecast (B, N, 12); theta (B, N, N); knn (B, N, N); lambda

def step_forward(short, long, epoch):
    # Sample a graph from learned Bernoulli logits; kNN target  [eq. 2, 4]
    theta, states, knn, sampled = graph_learner(long, frozen_tsformer)
    # Keep the representation of the last short patch only  [§3.2]
    last_state = states[:, :, -1, :]             # (B, N, d)
    # Backend on the short window and sampled graph; the projected state
    # is added before its output layers  [eq. 5] [code: model.py#L215-L217]
    forecast = backend(short, sampled, project(last_state))  # (B, N, 12)
    # Graph-loss weight: a step function of epoch  [app. A.2]
    lam = 1 / (int(epoch / 6) + 1)               # code: step.py#L69
    return forecast, theta, knn, lam

Key equations (as written in the paper; the code differs in places, see Paper vs code).

\mathbf{U}_j^i = \mathbf{W}\cdot\mathbf{S}^i_j + \mathbf{b} eq. 1

\begin{aligned} \mathbf{\Theta}_{ij} &= \text{FC}(\text{relu}(\text{FC}(\mathbf{Z}^i \parallel \mathbf{Z}^j))) \\ \mathbf{Z}^i &= \text{relu}(\text{FC}(\mathbf{H}^i)) + \mathbf{G}^i \end{aligned} eq. 2

\mathcal{L}_{graph} = \sum_{ij} -\mathbf{A}^a_{ij}\log\mathbf{\Theta}'_{ij} - (1-\mathbf{A}^a_{ij})\log(1-\mathbf{\Theta}'_{ij}) eq. 3

with \mathbf{\Theta}'_{ij} = \text{softmax}(\mathbf{\Theta}_{ij}).

\mathbf{A}_{ij} = \text{softmax}((\mathbf{\Theta}_{ij} + \mathbf{g}) / \tau) eq. 4

with \mathbf{g} i.i.d. Gumbel(0, 1) samples and \tau the temperature.

\mathbf{H}_{final} = \text{SP}(\mathbf{H}_P) + \mathbf{H}_{gw} eq. 5

\mathcal{L}_{regression} = \frac{1}{T_f N C} \sum_{j}\sum_{i}\sum_{k} \left|\hat{\mathcal{Y}}_{ijk} - \mathcal{Y}_{ijk}\right| eq. 6

\mathcal{L} = \mathcal{L}_{regression} + \lambda\,\mathcal{L}_{graph} eq. 7

Hyperparameters

Values from the paper unless the "Where" column names the code. Code paths are relative to the repository root.

Results

Author-reported. Setup: three public traffic datasets (METR-LA, 207 sensors; PEMS-BAY, 325 sensors; PEMS04, 307 sensors; all 5-minute readings) §4.1 tab. 1. Splits are 70/10/20 for METR-LA and PEMS-BAY and 60/20/20 for PEMS04 §4.1 generate_training_data.py generate_training_data.py. Baselines: HA, VAR, SVR, FC-LSTM, DCRNN, STGCN, Graph WaveNet, ASTGCN, STSGCN, GMAN, MTGNN, GTS §4.1 app. A.1. STEP is pre-trained separately for each dataset §4.1. Table below copies every row of the paper's main table, with the significance marker on STEP where the paper gives it tab. 2.

* marks the paper's significance (t-test, p < 0.05) against the best baseline. Bold marks STEP's cells and, in the PEMS04 horizon-12 RMSE column, GMAN's 30.21, which beats STEP's 31.33; STEP is not marked significant in that cell.

Ablations and studies are shown in figures only; the paper gives no table for them fig. 4. In words: removing the graph learner hurts; replacing the TSFormer-based kNN regulariser with a kNN built from raw training series also hurts; STEP also improves DCRNN as a second backend; and both masking ratio and kNN k have an interior optimum on METR-LA fig. 4 §4.4 §4.5.

Inspection of TSFormer (PEMS04): cosine similarity of patch states shows daily and weekly structure, learned positional embeddings reproduce the periodicity, and sinusoidal embeddings are reported to fail fig. 3 §4.3. Reconstructions of hidden patches from visible ones are visualised in fig. 7, and forecasts on six METR-LA series in fig. 8, both in the appendix app. C.

Efficiency (METR-LA, one RTX 3090): the pre-training cost falls as the masking ratio grows fig. 5. In the forecasting stage the paper reports the graph-structure module at about 55 seconds per epoch on average, and reports that pre-computing TSFormer outputs removes repeated work app. B fig. 6.

Paper vs code

Official code is zezhishao/STEP (Apache-2.0), linked from the paper §4.1. Code at commit 566e2738da2d83f055718d8edb609ad8dc325204. Checked against the method text, equations, appendix and hyperparameter tables. Each entry gives both sides.

1. Decoder positional embeddings. The paper says the decoder adds no positional embeddings, because the encoder already did §3.1. The code adds them to the mask tokens at hidden positions tsformer.py.
2. Decoder output. The paper says MLPs predict the patch values §3.1. The code uses one linear layer tsformer.py.
3. Dynamic graph feature. Eq. 2 includes a term from TSFormer states eq. 2. The code removes it and says it was unstable discrete_graph_learning.py. The graph logits then depend only on the static global feature discrete_graph_learning.py.
4. Graph loss reduction. Eq. 3 is a sum over all pairs eq. 3. The code uses PyTorch's mean BCE step_loss.py. The sum over one sample's N² pairs is N² times the mean; for N = 207 that is 42,849. Lambda then scales whichever version is used.
5. kNN construction. The paper uses k = 10 neighbours per node §4.1. The code keeps the top k × N entries of the whole N × N similarity matrix, so it is a global budget of 10 N edges, not per node discrete_graph_learning.py discrete_graph_learning.py.
6. Lambda schedule. The appendix gives lambda = 1 / ceil(epoch / 6) app. A.2. The code uses 1 / (floor(epoch / 6) + 1) step.py. The text calls it a gradual decay §3.2; the code is a step function. The runner's curriculum code subtracts 1 from the epoch base_tsf_runner.py, which suggests epochs are 1-based. If so, the code drops lambda one epoch earlier than the appendix formula.
7. Fusion point. Eq. 5 adds the projector output to the backend latent eq. 5. The code adds it to the skip accumulator just before the output ReLU and convolutions model.py.
8. Backend adaptive adjacency. The paper uses only the learned graph and says Graph WaveNet's details are omitted §3.2. The code also uses an adaptive adjacency learned from a 10-dimensional node embedding model.py model.py.
9. Pre-train optimizer. The appendix table says AdamW tab. 3. The config uses Adam with weight decay 0 TSFormer_METR-LA.py. These are equivalent here.
10. Forecast batch size. The table gives 64 for METR-LA and PEMS-BAY, and 32 for PEMS04 tab. 4. Code: 32 for METR-LA STEP_METR-LA.py, 32 for PEMS-BAY STEP_PEMS-BAY.py, 8 for PEMS04 STEP_PEMS04.py.
11. Forecast gradient clip. The table gives 5 tab. 4. The METR-LA and PEMS04 configs use 3 STEP_METR-LA.py STEP_PEMS04.py.
12. Curriculum. The table gives cl_num 3 and warm_num 30 for all datasets tab. 4. The METR-LA config uses 6 and 0 STEP_METR-LA.py. The PEMS-BAY config matches the table STEP_PEMS-BAY.py.
13. Pre-train learning rate and batch. The appendix says lr = base × batch / 8 with base 5.0e-4 app. A.2. The METR-LA (batch 8, lr 5e-4) and PEMS-BAY (batch 16, lr 1e-3) configs follow this rule TSFormer_METR-LA.py TSFormer_PEMS-BAY.py. The PEMS04 config uses batch 6 with lr 1e-3, which the rule does not give TSFormer_PEMS04.py TSFormer_PEMS04.py. PEMS04 pre-training also runs 200 epochs TSFormer_PEMS04.py.
14. Unstated architecture detail. The encoder and decoder inputs are multiplied by sqrt(d) before the blocks transformer_layers.py. The blocks are PyTorch's default encoder layers transformer_layers.py. The paper states neither.
15. Repository log. The shipped METR-LA training log is dated 2022-10-02, after the paper's v2. Its final horizon-12 test line reads MAE 3.3626, RMSE 7.0157, MAPE 0.0982 [log: training_logs/STEP_METR-LA.log#L1954]. That is close to the paper's 3.37 / 6.99 / 9.61% tab. 2, but it is one run, not the reported setup.

Pitfalls

• The learned graph is almost static. With the dynamic term removed, the Bernoulli logits are the same for every sample, and only the Gumbel noise differs discrete_graph_learning.py. The paper's claim that the graph follows the TSFormer states does not hold in this code. Reimplement with eq. 2 if you want that behaviour, and expect the code's numbers to differ.
• Test-time predictions are random. The sampler always draws Gumbel noise, with no check for evaluation mode discrete_graph_learning.py discrete_graph_learning.py. Predictions vary across runs unless you seed the sampler, and the paper does not say how it handled this.
• Frozen encoder may still apply dropout. requires_grad is set to False on the encoder step.py, but nothing in the STEP code sets it to eval mode during training. The encoder's dropout (0.1) may therefore stay active while the forecaster trains. This is inferred from the code and was not run.
• Missing values in the loss are unclear. The masked MAE ignores targets equal to 0.0 mae.py. The data are standard-normalised generate_training_data.py, so a raw zero reading may not reach the loss as an exact 0.0. This was not checked, and the card does not claim the training loss excludes missing readings. Metrics are computed after rescaling.
• Zero-filled long history. Samples without enough history before them get an all-zero long input, with no mask flag forecasting_dataset.py forecasting_dataset.py. The model cannot tell it from real data equal to zero in normalised space. The paper is silent.
• Positional table limit. The learnable positional embedding has 1000 rows positional_encoding.py. The paper's longest history (P = 336) fits; a longer history needs a code change.
• Hard-coded sizes. The graph learner hard-codes node counts, training lengths and the flattened CNN output size per dataset discrete_graph_learning.py discrete_graph_learning.py. Changing the patch count, patch length or node count needs code edits.
• Input channels. The forecast config passes three features, but Graph WaveNet keeps only the first two and left-pads one zero time step STEP_METR-LA.py model.py.
• Pre-training mask shared across the batch. One mask is drawn per forward pass with Python's random module, not the torch RNG, and it is shared by all series in the batch mask.py. Seed both RNGs for reproducibility.
• Compute. The forecaster calls TSFormer on every batch unless you pre-compute its outputs step.py. The paper says pre-computing significantly reduces repeated work fig. 6 app. B.
• Significance details. The paper does not give the number of runs behind its t-test. The authors ran GTS five times and reported its best run §4.2.
• Evaluation. All numbers are author-reported, on the paper's splits and preprocessing. Compare against them only with the same splits and metric scale.

Lineage

• Builds on: masked autoencoders for images, which set the masking design and the asymmetric encoder and decoder §3.1; the Transformer §3.1; Graph WaveNet as the backend, with an arXiv id given in the code docstring model.py; the Gumbel-Softmax reparameterisation §3.2; and kNN-based graph regularisation in the style of GTS §3.2.
• Compared against: HA, VAR, SVR, FC-LSTM, DCRNN, STGCN, Graph WaveNet, ASTGCN, STSGCN, GMAN, MTGNN and GTS §4.1.
• Superseded by: not stated in content on this wiki.