Source-linked AI summary

Spatial-Temporal Large Language Model for Traffic Prediction

Chenxi Liu, Sun Yang, Qianxiong Xu, Zhishuai Li, Cheng Long, Ziyue Li, Rui Zhao

arXiv:2401.10134v4cs.LGcs.CL

TL;DR

Traffic prediction requires forecasting future traffic features from historical data despite complex spatial-temporal dependencies and limitations in existing models. ST-LLM represents each location’s timesteps as tokens, learns unified spatial-temporal representations, and adapts an LLM with partially frozen attention. Experiments report that it outperforms state-of-the-art models and other LLMs, including strong few-shot and zero-shot performance.

  • Problem

    Traffic prediction must forecast future features from historical data, but existing models struggle with complex spatial-temporal dependencies and LLM-based methods often overlook spatial information.

  • Method

    ST-LLM defines timesteps at each location as tokens, embeds spatial-temporal information, fuses token representations, and uses partially frozen attention to capture global dependencies.

  • Results

    ST-LLM outperforms other LLMs and state-of-the-art traffic prediction models across traffic prediction scenarios, with a 22.5% average MAE improvement over OFA.

  • Takeaways & Limitations

    ST-LLM shows promise as a spatial-temporal learner for traffic prediction, including few-shot and zero-shot settings.

Abstract

from arXiv · show

Traffic prediction, an essential component for intelligent transportation systems, endeavours to use historical data to foresee future traffic features at specific locations. Although existing traffic prediction models often emphasize developing complex neural network structures, their accuracy has not improved. Recently, large language models have shown outstanding capabilities in time series analysis. Differing from existing models, LLMs progress mainly through parameter expansion and extensive pretraining while maintaining their fundamental structures. Motivated by these developments, we propose a Spatial-Temporal Large Language Model (ST-LLM) for traffic prediction. In the ST-LLM, we define timesteps at each location as tokens and design a spatial-temporal embedding to learn the spatial location and global temporal patterns of these tokens. Additionally, we integrate these embeddings by a fusion convolution to each token for a unified spatial-temporal representation. Furthermore, we innovate a partially frozen attention strategy to adapt the LLM to capture global spatial-temporal dependencies for traffic prediction. Comprehensive experiments on real traffic datasets offer evidence that ST-LLM is a powerful spatial-temporal learner that outperforms state-of-the-art models. Notably, the ST-LLM also exhibits robust performance in both few-shot and zero-shot prediction scenarios. The code is publicly available at https://github.com/ChenxiLiu-HNU/ST-LLM.

I. INTRODUCTION

Traffic prediction remains difficult because traffic data combines irregular spatial structure with complex temporal dependencies, while existing LLM-based methods often emphasize time and overlook space. ST-LLM addresses this gap by representing locations’ timesteps as tokens, learning spatial-temporal embeddings, and adapting attention for global dependencies.

  • Motivation: Traffic prediction supports intelligent transportation tasks such as traffic management and public-transportation scheduling.Historical data are used to predict future traffic features at specific locations, including flow.
  • Limitations of prior methods: CNNs and RNNs struggle with traffic’s non-Euclidean spatial structure and complex periodicity when modeling spatial-temporal dependencies.These limitations motivated later graph- and attention-based approaches.
  • LLM motivation: Existing LLM-based traffic methods often focus on temporal data despite strong spatial correlations and potentially larger spatial than temporal dimensions.A representative setting uses twelve historical and twelve future timesteps across hundreds of locations.
  • ST-LLM: ST-LLM defines each location’s timesteps as a token and uses spatial-temporal embeddings with fusion convolution to form unified representations.The representation emphasizes spatial locations and global temporal patterns.
  • ST-LLM: Partially frozen attention adapts the LLM by selectively freezing multi-head attention layers to capture global spatial-temporal dependencies between tokens.The strategy is designed specifically for traffic prediction.
  • Evaluation: Experiments on real traffic datasets report superior ST-LLM performance across settings, including few-shot and zero-shot prediction.The paper also describes these settings as evidence of intra-domain and inter-domain knowledge transfer.

B. Traffic Prediction

Traffic prediction models must forecast future features from historical data while handling spatial-temporal dependencies across irregular urban locations. ST-LLM represents locations as tokens, embeds their spatial-temporal information, processes the representations with an adapted LLM, and regresses future traffic features.

  • Problem setting: Traffic prediction uses historical traffic data to forecast future traffic features, but spatial-temporal dependencies make classic time-series models inadequate.Early approaches included ARIMA and Kalman filters.
  • Related challenges: CNNs face difficulty modeling traffic because urban spatial structure is non-Euclidean rather than grid-like.This irregularity limits accurate capture of spatial dependencies.
  • Related challenges: GCN-based models provide local spatial modeling but suffer from over-smoothing, limiting their ability to capture broader spatial patterns.Their use in traffic prediction reflects the importance of graph-based spatial structure.
  • Problem definition: Traffic data are represented as X ∈ R^(T×N×C), where T is timesteps, N is spatial stations, and C is the feature dimension.When C = 1, the feature can represent traffic pick-up or drop-off flow.
  • Problem definition: Given P historical timesteps, the prediction objective is to learn f(·) with parameters θ that forecasts the following S timesteps.The historical input is X^P and the target is Y^S.
  • ST-LLM framework: ST-LLM processes location tokens through spatial-temporal embedding, fusion convolution, a partially frozen LLM, and regression convolution to predict future traffic.Frozen early layers preserve pretrained knowledge while unfrozen later attention layers focus on spatial-temporal dependencies.

B. Spatial-Temporal Embedding and Fusion

ST-LLM converts traffic timesteps at each location into tokens and builds unified spatial-temporal representations by combining token, spatial, and temporal embeddings.

  • Traffic timesteps at each location are treated as tokens and transformed into spatial-temporal representations containing spatial correlations, hour-of-day, day-of-week, and token information.
  • A pointwise convolution maps input traffic data into token embeddings with hidden dimension D.The pointwise convolution uses a 1 × 1 kernel and learnable parameters.
  • Separate learnable embeddings encode hour-of-day and day-of-week patterns, whose sum forms the temporal representation ET.
  • An adaptive embedding represents spatial correlations among token pairs through learnable parameters and an activation function.
  • Fusion convolution concatenates token, spatial, and temporal embeddings and projects them into the dimensions required by the LLM, producing a uniform token representation.The concatenated tensor is denoted HF ∈ R^N×3D before fusion.

C. Partially Frozen Attention (PFA) LLM

The PFA LLM adapts a pretrained GPT2-based model to traffic prediction by freezing early attention and feed-forward layers while unfreezing later attention layers for spatial-temporal dependencies.

  • PFA addresses the limited performance of fully frozen pretrained transformers on short-term and long-term traffic prediction.
  • Unlike FPT, PFA freezes the first F layers but unfreezes multi-head attention in the last U layers to adapt to spatial-temporal dependencies.
  • PFA reverses the traditional calculation dimension from temporal to spatial to capture global dependencies in traffic data.The paper states that this shift reflects the importance of spatial dynamics in determining flow patterns.
  • The PFA LLM uses a Transformer-based GPT2 architecture with layer-normalization modifications and adapts it to the fused tensor HF.
  • After the PFA LLM, a regression convolution predicts traffic features for the following S timesteps, using a loss with L2 regularization.

V. EXPERIMENTS

The experiments evaluate ST-LLM against baselines through extensive experimental evaluations.

  • The study validates ST-LLM through a series of extensive experimental evaluations.

A. Datasets

The evaluation uses real-world NYC taxi and Citi Bike traffic datasets collected over the same three-month period and represented at 30-minute intervals.

  • The experiments use real-world traffic data from the NYCTaxi and CHBike datasets.
  • NYCTaxi contains over 35 million trips categorized into 266 virtual stations across 4,368 half-hour timesteps from April 1 to June 30, 2016.
  • CHBike contains approximately 2.6 million Citi Bike orders from 250 frequently used stations across 4,368 30-minute timesteps during the same period.

B. Baselines

The paper compares ST-LLM with GNN-based, attention-based, and LLM baselines spanning established traffic-prediction architectures and pretrained-transformer adaptations.

  • The baselines comprise six GNN-based models, three attention-based models, and four LLMs, including OFA, GATGPT, GCNGPT, and LLAMA2.
  • GNN-based models: DCRNN models traffic as a directed graph using diffusion convolutional recurrent networks.
  • GNN-based models: STGCN combines graph convolution with one-dimensional convolution for traffic time-series prediction.
  • GNN-based models: GWN employs graph convolution with an adaptive adjacency matrix, while AGCRN learns nodes and interdependencies among traffic series.
  • GNN-based models: STG-NCDE uses graph neural controlled differential equations, and DGCRN uses dynamic graph convolutional recurrent networks.
  • Attention-based and LLM models: ASTGCN, GMAN, and ASTGNN represent attention-based alternatives, while GATGPT and GCNGPT combine graph modules with frozen GPT2.

C. Implementations

Experiments use standardized temporal splits, multi-step prediction settings, four error metrics, and NVIDIA A100 hardware for evaluation and training.

  • NYCTaxi and CHBike are divided into training, validation, and test sets using a 6:2:2 ratio.
  • The experiments use 12 historical timesteps and predict the following 12 timesteps, with weekly and daily periods set to 7 and 48.
  • Training runs on NVIDIA A100 GPUs with 40GB of memory per GPU.
  • MAE and RMSE measure absolute errors, while MAPE and WAPE measure relative errors; lower values indicate better performance.

E. Main Results

ST-LLM achieves the strongest reported traffic-prediction performance among the compared models, and ablations show that its LLM, spatial-temporal embeddings, and partially frozen attention are all important.

  • Main comparison: ST-LLM outperforms other LLMs in four traffic-prediction scenarios across the evaluated datasets.
  • Main comparison: 22.5% average MAE improvement over OFA and 20.8% over LLAMA2 are reported for ST-LLM.
  • Main comparison: LLM-based methods rank highest overall, followed by attention-based models and then GCN-based models.
  • Component ablation: Removing the LLM increases error across all metrics on the NYCTaxi ablation, indicating degraded prediction performance without the language model.
  • Component ablation: Removing spatial-temporal, temporal, or spatial embeddings worsens performance, while the full ST-LLM records the lowest error rates across all metrics.
  • Attention ablation: PFA outperforms FPT, No Pretrain, Full Layer, and Full Tuning across all metrics on all datasets.

G. Parameter Analysis

The number of unfrozen final attention layers affects performance non-monotonically and varies across datasets and metrics.

  • On NYCTaxi Pick-up, WAPE improves as U increases to 2, then degrades when U exceeds 2.
  • On CHBike Pick-up, U=1 yields the lowest WAPE and MAE, while increasing U raises both errors.

H. Inference Time Analysis

ST-LLM achieves the lowest MAE while maintaining competitive inference times across the evaluated NYCTaxi and CHBike settings. LLAMA2 is excluded because its inference time is substantially longer than the other LLMs.

  • LLAMA2 is omitted because its inference time is significantly longer than that of the other evaluated LLMs.
  • NYCTaxi: ST-LLM achieves the lowest MAE on NYCTaxi while maintaining competitive inference times.OFA has similar inference time but slightly higher MAE, whereas GATGPT and GCNGPT are slower and less accurate.
  • CHBike: ST-LLM achieves the lowest MAE on both CHBike pick-up and drop-off datasets while remaining close to OFA in inference speed.GATGPT and GCNGPT consistently require longer inference times and produce higher MAEs.
  • ST-LLM provides the best balance between inference speed and predictive accuracy across the NYCTaxi and CHBike datasets.
Loading 2401.10134v4…