Spatial-Temporal Graph ODE Networks for Traffic Flow Forecasting

In plain words

Road sensors report how busy a road is every few minutes. Forecasting the next hour is hard because what happens on one road depends on roads that may be far away on the map, and on roads whose traffic behaves the same way even though they are not neighbours.

Many earlier models looked only at nearby roads and only a few layers deep. Stacking many layers of that kind makes every sensor's numbers blur into the same average. This paper treats the stacking as a continuous process, described by a differential equation, and adds a restart term that keeps pulling each sensor back toward what it measured at the start. The sensors stay distinct even when the process runs for many steps, so information from further away can reach each one.

The model uses two kinds of road graph, one from distances and one from how similar the traffic patterns are. It handles time with dilated convolutions, which look far back in the history without a recurrent loop. The authors report that it beats six earlier forecasting models on six California datasets.

TL;DR

STGODE predicts the next 60 minutes of traffic flow from the previous 60 minutes §5.3. It runs graph propagation over a spatial graph and a semantic graph built from DTW similarity, written as a tensor ODE with a restart term §4.2. Dilated temporal convolutions wrap each ODE block §4.3. Author-reported: the lowest RMSE, MAE and MAPE on all six PeMS datasets tab. 2. The official code is Apache-2.0 but departs from the paper in ways that change the model (see Paper vs code), so implement from the equations and the code together.

Problem

Traffic flow depends on road layout, and on distant roads with similar patterns, at the same time §1 fig. 1. Trips vary in length, so near and far dependencies both matter §1. Two gaps are named. Most models treat spatial and temporal patterns with separate modules, which misses their interaction §1 §5.4. Graph networks also gain little from depth: deep stacks over-smooth, so node representations converge. The best results in prior work come from two cascaded graph layers §1 §2.

The method assumes a fixed sensor network with distances between sensors, and traffic series aggregated to 5-minute steps §4.1 §5.1.

Method

Overview: two spatial-temporal graph ODE layers are cascaded. Each layer has several blocks in parallel, one group per adjacency matrix. Their outputs are max-pooled and a two-layer MLP produces the forecast §4.4 §4.5 fig. 3. Each block is a sandwich of a temporal convolution, a tensor ODE solver and a second temporal convolution §4.4 fig. 3 fig. 5.

Adjacency matrices §4.1.

• Spatial: a Gaussian of the road distance, sparsified by a threshold eq. 4. The text calls the Gaussian width σ² and the threshold ε, and §5.3 gives σ = 10 and ε = 0.5.
• Semantic: the DTW distance between two node series eq. 5. The pair is linked when that distance is below ε eq. 6. The text gives ε = 0.6 without a unit §5.3.

Regularised adjacency §3.2. Â = (α/2)(I + D^{-1/2} A D^{-1/2}) eq. 1, with α in (0, 1). Its eigenvalues lie in [0, α]. The paper sets α = 0.8 §5.3.

Tensor residual layer §4.2. A graph convolution is H_{l+1} = σ( H_l W) eq. 7. The tensor version multiplies the node mode by Â, the time mode by U and the feature mode by W, and adds the initial state H_0 from the upstream network eq. 8. Expanding the recursion gives a sum over all powers of  eq. 9. Without the H_0 term the powers of  collapse onto the top eigenvector, which loses information eq. 10.

Continuous form §4.2. Letting the number of layers grow turns the sum into an integral and then a differential equation eq. 12 eq. 13. Corollary 1 states the form with logarithms eq. 14. The logarithms are approximated by their first-order Taylor term, ln M ≈ M − I. This gives the simplified ODE with (Â − I), (U − I), (W − I) and the restart term H_0 §4.2. U and W are parameterised as P diag(λ) P^T and Q diag(μ) Q^T, where P and Q are learnable orthogonal matrices and λ, μ are clamped to (0, 1) §4.2. The ODE is solved with an Euler solver, and the solver's step count and size are not stated §4.2. The paper also gives an analytic solution (Corollary 2) §4.2.

Temporal convolution §4.3. Each layer is a dilated 1-D convolution over time, with dilation 2^(l−1) and zero padding that keeps the length. A residual structure is added §4.3 fig. 5. The text does not give the kernel size, the number of layers, or the activation.

Output and loss §4.5. The outputs of the parallel blocks are max-pooled element-wise. A two-layer MLP produces the 60-minute forecast. The loss is Huber with a threshold δ that the paper does not give §4.5.

# STGODE block: Euler solve of the tensor ODE  [§4.2] [eq. 14]
# Input:  H0 (N, T, F) from the upstream network; A_hat (N, N)
#         U (T, T) and W (F, F), eigenvalues in (0, 1)
# Output: state (N, T, F)

def stgode_block(H0, A_hat, U, W, num_steps, step_size):
    # The initial state is the upstream output  [§4.2]
    state = H0
    # Euler steps; the step count and size are not stated  [§4.2]
    for step in range(num_steps):
        # Simplified derivative, ln M replaced by M - I  [§4.2]
        spatial = mode_product(state, A_hat - I, mode=1)
        temporal = mode_product(state, U - I, mode=2)
        feature = mode_product(state, W - I, mode=3)
        # Restart term added at every step  [§4.2]
        derivative = spatial + temporal + feature + H0
        state = state + step_size * derivative
    return state

# Spatial adjacency  [§4.1] [eq. 4]
# Input:  distances (N, N); bandwidth (the paper's sigma^2); eps
# Output: weights (N, N), continuous, zero below eps

def spatial_adjacency(distances, bandwidth, eps):
    # Gaussian weight of each sensor pair  [§4.1] [eq. 4]
    weights = exp(-distances ** 2 / bandwidth)
    # Drop links whose weight falls below eps  [§4.1] [eq. 4]
    weights[weights < eps] = 0
    return weights

# Semantic adjacency  [§4.1] [eq. 5] [eq. 6]
# Input:  series (N, M), one traffic series per node; eps
# Output: binary adjacency (N, N)

def semantic_adjacency(series, eps):
    adjacency = zeros(N, N)
    for i in range(N):
        for j in range(N):
            # DTW distance between two node series  [§4.1] [eq. 5]
            distance = dtw(series[i], series[j])
            # Link the pair when the distance is below eps  [§4.1] [eq. 6]
            if distance < eps:
                adjacency[i, j] = 1
    return adjacency

# STGODE network  [§4.4] [§4.5] [fig. 3]
# Input:  X (N, T, F); A_spatial and A_semantic (N, N), regularised
# Output: forecast (N, T_out)
# temporal_conv: dilated convolution with residual  [§4.3]
# U and W are learned per block  [§4.2]

def stgode_network(X, A_spatial, A_semantic):
    branch_outputs = []
    # One group of parallel branches per adjacency  [§4.4] [fig. 3]
    for A_hat in [A_spatial, A_semantic]:
        # Three blocks per layer  [§5.3]; the split by adjacency is inferred
        for branch in range(3):
            H = X
            # Two cascaded layers  [§4.4] [fig. 3]
            for layer in range(2):
                # Sandwich: temporal conv, ODE block, temporal conv  [fig. 5]
                H = temporal_conv(H)
                H = stgode_block(H, A_hat, U, W, num_steps, step_size)
                H = temporal_conv(H)
            branch_outputs.append(H)
    # Element-wise max over all branches  [§4.5]
    H = elementwise_max(branch_outputs)
    # Two-layer MLP maps the features to the forecast  [§4.5]
    return mlp(flatten(H))

Hyperparameters

Results

Author-reported. Setup: six PeMS datasets tab. 1; one hour of history predicts the next 60 minutes §5.3; baselines ARIMA, STGCN, DCRNN, ASTGCN(r), GraphWaveNet and STSGCN §5.2; hardware: Intel Xeon E5-2682 v4 and an NVIDIA Tesla V100 16 GB §5.3. Model size and run time are not stated.

Datasets tab. 1:

Comparison tab. 2. Best value per row is in bold in the paper. STGODE is labelled STODE in the table. STSGCN is the best baseline in 12 of 18 rows and ASTGCN(r) in 6.

Margins over the best baseline are small in some rows. Examples: PeMS04 RMSE 32.82 against 33.65 for STSGCN, and PeMS08 RMSE 25.97 against 26.22 for ASTGCN(r) tab. 2. No variance, repeated runs or significance tests are reported §5.4.

Ablation on PeMS04 §5.6.1 fig. 7. The authors compare four variants. The text says that the graph-convolution variant (STGCN*) beats the original STGCN, that the spatial-only variant is poor, that the matrix-based version is inferior to the tensor one, and that removing the H_0 restart hurts. The numbers are only in the figure.

Depth §5.6.2 fig. 8. The authors plot error against "depth", which here means the input length of the solver or the number of convolution layers in the STGCN variant. The text says the ODE model stays stable as this grows and the STGCN variant gets worse. Values are only in the figure.

Case study §5.5 fig. 6. Two nodes are compared with STGCN. The text says the ODE model follows abrupt changes more quickly and is smoother under noise.

Paper vs code

Official code: https://github.com/square-coder/STGODE, linked in a footnote on the abstract, Apache-2.0 LICENSE, commit e5e7bfc2194f8253063bcc6ba1d2903a48deb1a4. The code was read at that commit. Discrepancies that change the model:

1. TCN residual bypass. In TemporalConvNet.forward, the residual is a single conditional expression: the sum of the dilated stack and the 1×1 projection when a projection exists, and otherwise the raw input, wrapped in a ReLU model.py. The 1×1 projection exists only when the input channels differ from the last output channels model.py. When they match, the layer returns relu(input) and never calls its convolutions. That covers temporal2 in both blocks and temporal1 of the second block, where the input and output are both 64 channels model.py model.py model.py. Only the first temporal1 of each branch runs the dilated stack, so 6 of 24 TCN instances do. The paper describes a residual dilated stack for every TCN §4.3.

2. Spatial term scaling in the ODE. The paper uses Â, with α fixed at 0.8 eq. 1 §3.2. The code's adjacency is already 0.4 (I + Ã) utils.py. The ODE then multiplies the spatial term by alpha / 2, where alpha = sigmoid(a) and a is a learnable per-node vector initialised at 0.8 odegcn.py odegcn.py odegcn.py. At initialisation this gives 0.345 × 0.4 (I + Ã) = 0.138 (I + Ã), not 0.4 (I + Ã). The learnable factor is not in the paper.

3. U and W are not orthogonal. The paper parameterises them as P diag(λ) P^T with P orthogonal §4.2. The code uses unconstrained square matrices initialised to the identity and forms W diag(d) W^T with d clamped to [0, 1] odegcn.py odegcn.py odegcn.py. Without orthogonality the eigenvalues are not bounded by the clamp, so the stability argument in §4.2 does not carry over.

4. Solver step. The paper says Euler, with no length or step count §4.2. The code integrates on the grid [0, 6] and takes the state at index 1 odegcn.py model.py. I read torchdiffeq's fixed-grid solvers as taking one step of size 6 in this case. I did not verify this by running it.

5. Spatial graph. The paper's Eq. 4 uses raw distances eq. 4. The code z-scores all finite distances, then applies exp(−z²/σ²) with σ = 10 [code: utils.py#L108-L114] args.py. The effective bandwidth is σ² = 100 in z units. The kernel weights are kept as continuous values.

6. Semantic graph. The paper links pairs whose DTW distance is below 0.6 eq. 6. The code runs fastdtw with radius 6 on the average z-scored flow for each time of day, 288 steps averaged over days utils.py. That is an approximation on a daily profile, not the full series. It z-scores the matrix, applies exp(−z²/0.1²), and keeps entries above 0.6 utils.py args.py args.py. exp(−z²/0.01) > 0.6 holds only when |z| < 0.0715. So the code links pairs whose DTW distance lies within about 0.07 standard deviations of the mean, and drops both the most and the least similar pairs. This is derived from the code and was not run, because the pems04.npz file that the loader needs is not in the repository.

7. Normalisation and leakage. The z-score mean and standard deviation, and the daily DTW profile, are computed over the full series, which includes the validation and test periods utils.py utils.py. The paper does not say which period is used.

8. Optimiser and schedule. The paper says Adam, learning rate 0.01, batch 32 §5.3. The code uses AdamW run_stode.py, learning rate 2e-3 by default args.py, batch 16 args.py, and halves the rate every 50 epochs run_stode.py. No config file for the paper's numbers is included.

9. Loss. The paper's Huber threshold is not stated §4.5. The code uses SmoothL1Loss, which is Huber with δ = 1 run_stode.py.

10. Parallel branches. The paper says three blocks per layer §5.3. The code builds three spatial chains and three semantic chains, each with two cascaded blocks model.py. The mapping is inferred from the paper.

11. Evaluation. The paper does not describe these choices. The code masks targets equal to zero for RMSE, MAE and MAPE eval.py eval.py. It averages the per-batch metrics rather than computing them over the whole set run_stode.py, and the test loader is shuffled utils.py, so the batch composition varies. The checkpoint is chosen by best validation RMSE run_stode.py, and the test score comes from that checkpoint run_stode.py.

12. Dataset files. The loader expects PEMS04/pems04.npz utils.py. The repository ships only the PEMS04 CSV, its distance file and two cached .npy matrices. The other datasets are not included.

13. Naming and typos in the paper. The text uses STGODE, the table STODE §5.4 tab. 2. The text says PeMSD3 where the table says PeMS03 §5.1 tab. 1. The text gives 288 sensors for PeMSD7(M), but the table and the code give 228 §5.1 tab. 1 utils.py.

Pitfalls

• The corollary's constant depends on a limit the text does not justify. The step that fixes the constant to H_0 comes from "letting t go to −1" §4.2. The text gives no argument for it.
• ln  may be undefined. The eigenvalues of  lie in [0, α] §3.2. The logarithm is undefined at zero, and the first-order Taylor step ln M ≈ M − I avoids that. The error of this step is not analysed §4.2.
• Step count changes the model. The paper does not give the number of Euler steps or the integration length §4.2. Different choices give different models. The code's choice is noted under Paper vs code.
• "Depth" is not model depth. The depth experiment varies the solver's input window or the number of convolution layers in STGCN* §5.6.2 fig. 8. It does not vary the number of ODE steps.
• σ² and σ are ambiguous. The text calls σ² a threshold §4.1 and §5.3 gives σ = 10. The code uses σ = 10 on z-scored distances args.py.
• The DTW threshold has no unit. ε = 0.6 is given for DTW distances §4.1 §5.3, but it is not stated whether the distances are raw or standardised.
• Missing details for reimplementation. The TCN kernel size, layer count, activation and residual form; the Huber threshold; the MLP width; dropout; the learning-rate schedule; the normalisation scope; the upstream network that produces H_0; and the step count are all unstated or only given in code §4.2 §4.3 §4.5 §5.3.
• Cost is not reported. The paper gives no parameter count, run time or memory §5.3. The code runs six branches of two blocks each. Each adjacency is a dense N × N matrix utils.py, and PeMSD7(L) has 1026 nodes tab. 1.
• Single horizon, single run. Only the 60-minute result is reported §5.3. It is not stated whether the baselines were re-run or copied from their papers §5.2. No variance or significance is given §5.4.
• Ablation and depth numbers are only in figures. The text states the order of the results §5.6.1 §5.6.2, but the values are in fig. 7 and fig. 8.

Lineage

• Builds on: neural ODEs, cited as chen2018neural §2 §3.3; the continuous GNN cited as xhonneux2020continuous, which uses a restart distribution against over-smoothing §2 §4.2; residual networks §4.2 §4.3; STGCN, for the spatial adjacency construction §4.1; dilated temporal convolutions §4.3.
• Compared against: ARIMA, STGCN, DCRNN, ASTGCN(r), GraphWaveNet and STSGCN §5.2 tab. 2.
• Superseded by: not known from content in this wiki.