Position-Aware Unified Embedding with Linear Attention for Distinguishable Flow Modeling

Tao Cui, Yudong Lu, Di Dong, Chongguang Ren, Zhijian Qu, Panjing Li. Transportation Research Part C: Emerging Technologies, 2026.

This is a PyTorch implementation of STLAformer, as described in our paper:
https://doi.org/10.1016/j.trc.2025.105423

---

---

Performance on Traffic Forecasting Benchmarks

---

Citation

@article{cui2026linear,
  title={Position-Aware Unified Embedding with Linear Attention for Distinguishable Flow Modeling},
  author={Cui, Tao and Lu, Yudong and Dong, Di and Ren, Chongguang and Qu, Zhijian and Li, Panjing},
  journal={Transportation Research Part C: Emerging Technologies},
  pages={105423},
  year={2026},
  publisher={Elsevier}
}

---

Required Packages

pytorch>=1.11
numpy
pandas
matplotlib
pyyaml
pickle
torchinfo

---

PeMSD3 & PeMSD4 & PeMSD7 & PeMSD8

• Download the data 百度网盘 or Google Drive.
• Unzip them to corresponding folders.

---

Training Commands

cd model/  
python train.py --d {dataset} --g {gpu_id}  

---

Testing Commands

Select TestPreTrained.py, choose the dataset to be tested, and import the trained model parameters from {DATASET_NAME}.pth in the pre-trained folder.