Journal of Computer Applications

    Next Articles

Traffic prediction based on spatio-temporal bottleneck attention enhanced by pre-trained language model

  

  • Received:2026-02-02 Revised:2026-04-03 Online:2026-05-13 Published:2026-05-13

基于预训练语言模型增强时空瓶颈注意力的交通预测

柴艺函1,王语雁1,杜圣东2,胡节3   

  1. 1. 西南交通大学
    2. 西南交通大学(信息科学与技术学院)
    3. 西南交通大学计算机与人工智能学院
  • 通讯作者: 胡节

Abstract: To address high computational complexity and difficulty of modeling complex spatio-temporal dependencies in large-scale traffic flow prediction, a spatio-temporal bottleneck attention model enhanced by Pre-trained Language Model (PLM), termed STBA-PLM, was proposed. First, traffic flow data were mapped into spatio-temporally aware node representations through spatio-temporal tokenization. Then, a bottleneck attention encoding module integrating spatial regions and temporal patterns was constructed to compress large-scale node-level representations into constant-size regional representations, thereby decoupling PLM input sequence length from number of nodes. Based on this, the compressed regional representations were fed into PLM for deep feature extraction, and a symmetric decoder was employed to accurately reconstruct node-level representations. Meanwhile, a temporal difference iteration module was designed to explicitly model differences between historical and future temporal patterns, enhancing model’s capability in temporal pattern learning. Experimental results on PEMS04, PEMS07, and PEMS08 datasets show that STBA-PLM outperforms models such as PDFormer, SFADNet, and ST-LLM in both prediction accuracy and computational efficiency. Compared with ST-LLM, a PLM-based model that focuses more on spatial relationships, the proposed model reduces RMSE, MAE, and MAPE by at least 2.6%, 4.1%, and 9.6%, respectively. When number of nodes reaches 1000, training time is reduced by 93.3%. Results demonstrate that the proposed model effectively reduces computational complexity while capturing spatio-temporal dependencies and temporal patterns in traffic data, achieving a good balance between prediction accuracy and computational efficiency, making it suitable for large-scale road network traffic flow prediction tasks.

摘要: 针对大规模路网交通流量预测任务中计算复杂度高与复杂时空特征建模困难的问题,提出一种基于预训练语言模型(PLM)增强时空瓶颈注意力的交通流预测模型STBA-PLM。该模型首先通过时空分词器将流量映射为具有时空感知的节点表征;随后构建融合空间区域与时间模式的瓶颈注意力编码模块,将大规模的节点级表征压缩为常数级时空区域表征,使PLM输入序列长度与节点规模解耦;在此基础上,利用PLM对压缩后的区域表征进行深层特征提取,并通过对称的解码器实现节点空间的精确还原;同时,设计时间差分迭代模块,显式建模历史与预测之间的时间模式差异,增强模型对时间模式的建模能力。在PEMS04、PEMS07和PEMS08数据集上的实验结果表明,STBA-PLM在预测精度和运算效率上均优于PDFormer、SFADNet及ST-LLM等模型。相较于更侧重空间关系的PLM模型ST-LLM,本文模型的RMSE、MAE和MAPE指标分别至少降低2.6%、4.1%和9.6%,在节点规模为1000时训练时间降低93.3%。所提模型在有效降低计算复杂度的同时,能够充分建模交通数据中的时空依赖与时间模式,适用于大规模路网交通流量预测任务。

CLC Number: