Anchors as Tokens: Anchor-Conditioned Spatio-Temporal Forecasting with Pretrained Language Models

Overview

ST-AaT (Anchors as Tokens) is a novel anchor-conditioned PLM framework for spatio-temporal forecasting that explicitly integrates historical baseline awareness.

Historical mean statistics are promoted to semantic anchor tokens to cleanly relate current observations to expected baselines, while a multi-factor tokenizer extracts complementary temporal dynamics. Reinforced by an Anchor Deviation Loss, this conditioning mechanism enforces consistency between historical and predicted variations, allowing the model to overcome deviation blindness and generate highly accurate, contextually-informed forecasts.

This repository provides the official PyTorch implementation for our paper "Anchors as Tokens: Anchor-Conditioned Spatio-Temporal Forecasting with Pretrained Language Models".

Setup

The project uses Conda for environment management. We provide a requirements.yaml file to help you set up the environment with all necessary dependencies easily.

1. Clone the repository:

   git clone https://github.com/your-username/Spatio-LLM.git
   cd Spatio-LLM

2. Create and activate the Conda environment:

   conda env create -f requirements.yaml
   conda activate stllm_test

3. Data Preparation: Make sure you download the traffic datasets (e.g., PEMS04, PEMS07, PEMS08) and put them in the corresponding data directories such as data/traffic/PEMS04/ containing .npz and .csv files.

Usage

To train and evaluate the model, you can run the provided shell scripts locate in the code/scripts directory.

Running via Scripts

For example, to run experiments on the PEMS04 dataset:

cd code/scripts
bash pems04.sh

Other available scripts include pems07.sh and pems08.sh.

Running Manually

Alternatively, you can run the primary training script main.py directly from the code/src directory. Here is an example of the command to run the model:

cd code/src
python main.py \
    --data_path '../../data/traffic/PEMS04/PEMS04.npz' \
    --adj_filename ../../data/traffic/PEMS04/PEMS04.csv \
    --dataset PEMS04FLOW \
    --desc ft_${loss}_${day}-${week}_PEMS04_pre \
    --sample_len 12 \
    --predict_len 12 \
    --train_ratio 0.6 \
    --val_ratio 0.2 \
    --epoch 500 \
    --val_epoch 1 \
    --test_epoch 5 \
    --batch_size 64 \
    --lr 0.001 \
    --causal 0 \
    --model gpt2 \
    --lora \
    --ln_grad \
    --patience 50 \
    --t_dim 64 \
    --node_emb_dim 64 \
    --node_embedding \
    --llm_layers 3 \
    --time_token \
    --dropout 0.05 \
    --trunc_k 64 \
    --weight_decay 0 \
    --task prediction \
    --sandglassAttn 1 \
    --sag_dim 128 \
    --sag_tokens 128 \
    --input_dim 1 \
    --output_dim 1 \
    --anchor_day_loss_weight $day \
    --anchor_week_loss_weight $week \
    --anchor_loss_type $loss \
    --use_anchor_day \
    --use_anchor_week

Citation

If you find this code or our paper useful, please consider citing it:

@article{anchors_as_tokens_2026,
  title={Anchors as Tokens: Anchor-Conditioned Spatio-Temporal Forecasting with Pretrained Language Models},
  author={Trung Kien Le, Thi Thu Nguyen, Nhat Hai Nguyen},
  journal={},
  year={2026}
}