关于自监督学习时空纠缠Transformer交通流预测的论文。

论文基本信息

项目	内容
标题	SSL-STMFormer: Self-Supervised Learning Spatio-Temporal Entanglement Transformer for Traffic Flow Prediction
作者	李泽涛、胡正、韩鹏（通讯作者）、顾宇、蔡世民
单位	电子科技大学，成都
会议	AAAI-25 (The Thirty-Ninth AAAI Conference on Artificial Intelligence)
页码	12130-12138

---

研究背景与核心问题

交通流预测的重要性

快速城市化和人口增长导致城市交通管理面临严峻挑战：交通拥堵、出行时间增加、空气污染。准确的交通流预测对缓解这些问题至关重要。

现有方法的三大核心局限

1. 长距离长时间纠缠捕获不足

• 交通系统存在长距离长时间纠缠：相距遥远的节点可能在特定时段表现出相似交通模式
• 案例（图1）：旧金山国际机场(Node A)、奥克兰国际机场(Node B)、Bernal Heights(Node C)
  • 7/10-7/13：三处交通流趋势高度相似
  • 7/14-7/17：相似性急剧下降
• 现有静态建模方法无法有效捕捉这种动态变化

2. 时空异质性建模不足

• 城市区域功能分割导致不同区域交通模式差异显著
• 即使相距遥远的两个位置也可能因功能相似（如商业区、住宅区）而表现出相似交通模式
• 这种相似性同样随时间动态变化

3. 时空纠缠的静态建模缺陷

• 时空依赖性和异质性因出行模式变化而动态演变
• 现有方法以预定义/静态方式建模，无法捕捉真实交通环境的动态特性
• GNN的过平滑问题进一步阻碍动态长期依赖的捕获

---

核心创新：SSL-STMFormer架构

整体框架（图2）

┌─────────────────────────────────────────────────────────────────┐
│                        输入数据嵌入层                              │
│  X_emb = X_data + X_spe(图拉普拉斯) + X_w(周周期) + X_d(日周期) + X_tpe │
├─────────────────────────────────────────────────────────────────┤
│                    时空编码器层 (ST Encoder Layer)                  │
│  ┌─────────────────────────┐  ┌─────────────────────────┐      │
│  │    空间注意力模块 (SA)     │  │    时间注意力模块 (TA)     │      │
│  │  • 短程掩码 M_short       │  │  • 随机掩码 M_rand        │      │
│  │  • 长程掩码 M_long (DTW)  │  │  • 管道掩码 M_tube        │      │
│  │  • 随机掩码 M_rand        │  │  • 块掩码 M_block         │      │
│  │  • 管道掩码 M_tube        │  │  • 时间掩码 M_temp        │      │
│  │  • 块掩码 M_block         │  │                         │      │
│  └─────────────────────────┘  └─────────────────────────┘      │
│  ↓                           ↓                                  │
│  空间纠缠感知模块              时间纠缠感知模块                        │
│  (Spatial Entanglement-Aware)  (Temporal Entanglement-Aware)      │
├─────────────────────────────────────────────────────────────────┤
│                    多头注意力融合 (Concat Multi-Heads)              │
│  STA = cat(Z^SRand, Z^STube, Z^SBlock, Z^Short, Z^Long,           │
│           Z^TRand, Z^TTube, Z^TBlock, Z^Temp, Z^T) · W^l        │
├─────────────────────────────────────────────────────────────────┤
│  逐位置前馈网络 (Position-Wise FFN) + 输出投影 (Output Projection)   │
├─────────────────────────────────────────────────────────────────┤
│                    自监督学习分支 (SSL Branch)                       │
│  ┌─────────────────────────┐  ┌─────────────────────────┐      │
│  │  空间异质性建模 (SHM)      │  │  时间异质性建模 (THM)      │      │
│  │  • 软聚类自监督任务        │  │  • 正/负样本对比学习        │      │
│  │  • 区域功能识别            │  │  • 时间步判别               │      │
│  └─────────────────────────┘  └─────────────────────────┘      │
│  L_joint = L_p(预测损失) + L_s(SHM损失) + L_t(THM损失)             │
└─────────────────────────────────────────────────────────────────┘

---

核心模块详解

模块一：数据嵌入层

多维度嵌入融合：
$${\mathbf{X}}_{emb}={\mathbf{X}}_{data}+{\mathbf{X}}_{spe}+{\mathbf{X}}_w+{\mathbf{X}}_d+{\mathbf{X}}_{tpe}$$

嵌入类型	维度	功能
数据嵌入 ${\mathbf{X}}_{data}$	${\mathbb{R}}^{T\times N\times d}$	原始交通流数据经全连接层转换
空间图拉普拉斯嵌入 ${\mathbf{X}}_{spe}$	${\mathbb{R}}^{N\times d}$	基于图拉普拉斯特征向量，描述节点间距离信息
周周期嵌入 ${\mathbf{X}}_w$	${\mathbb{R}}^{T\times d}$	捕捉每周周期性（周一到周日）
日周期嵌入 ${\mathbf{X}}_d$	${\mathbb{R}}^{T\times d}$	捕捉每日周期性（1-1440分钟）
时间位置编码 ${\mathbf{X}}_{tpe}$	${\mathbb{R}}^{T\times d}$	标准Transformer位置编码

矩阵加法通过维度复制实现。

---

模块二：时空编码器层

2.1 空间注意力模块 (Spatial Attention, SA)

标准自注意力：
$${\mathbf{Q}}_t^{(S)}={\mathbf{X}}_{t::}{\mathbf{W}}_Q^S,{\mathbf{K}}_t^{(S)}={\mathbf{X}}_{t::}{\mathbf{W}}_K^S,{\mathbf{V}}_t^{(S)}={\mathbf{X}}_{t::}{\mathbf{W}}_V^S$$

$${\mathbf{A}}_t^{(S)}=\frac{({\mathbf{Q}}_t^{(S)})({\mathbf{K}}_t^{(S)}{)}^{\top }}{\sqrt{d^{\prime }}}$$

$$\text{SA}({\mathbf{Q}}_t^{(S)},{\mathbf{K}}_t^{(S)},{\mathbf{V}}_t^{(S)})=F({\mathbf{A}}_t^{(S)}){\mathbf{V}}_t^{(S)}$$

关键洞察：空间依赖在不同时间片动态变化，SA模块专门捕捉这些动态空间依赖。

---

2.2 空间纠缠感知模块 (Spatial Entanglement-Aware)

核心思想：并非所有节点对都重要，需要识别关键节点对（邻近节点对 + 远距离相似节点对）

五种掩码策略：

掩码策略	数学表达	目标
随机掩码 ${\mathbf{M}}_{rand}$	$\text{RandSA}=F({\mathbf{A}}_t^{(S)}\odot {\mathbf{M}}_{rand}^{(S)}){\mathbf{V}}_t^{(S)}$	捕捉细粒度时空关系（类似MAE）
管道掩码 ${\mathbf{M}}_{tube}$	$\text{TubeSA}=F({\mathbf{A}}_t^{(S)}\odot {\mathbf{M}}_{tube}^{(S)}){\mathbf{V}}_t^{(S)}$	增强空间外推能力（模拟传感器故障）
块掩码 ${\mathbf{M}}_{block}$	$\text{BlockSA}=F({\mathbf{A}}_t^{(S)}\odot {\mathbf{M}}_{block}^{(S)}){\mathbf{V}}_t^{(S)}$	增强空间可迁移性
短程掩码 ${\mathbf{M}}_{short}$	$\text{ShortSA}=F({\mathbf{A}}_t^{(S)}\odot {\mathbf{M}}_{short}){\mathbf{V}}_t^{(S)}$	距离<λ跳的节点权重为1
长程掩码 ${\mathbf{M}}_{long}$	$\text{LongSA}=F({\mathbf{A}}_t^{(S)}\odot {\mathbf{M}}_{long}){\mathbf{V}}_t^{(S)}$	DTW相似度Top-K节点权重为1

短程掩码构建：

• 若两节点图距离（跳数）< 阈值λ，则权重为1，否则为0

长程掩码构建（Dynamic Time Warping, DTW）：

1. 用DTW算法计算节点间历史交通流相似度
2. 对每个节点，识别K个最相似节点作为长程邻居
3. 当前节点与长程邻居间权重设为1，其余为0

Hadamard积 $\odot$ 实现掩码与注意力分数的逐元素相乘。

---

2.3 时间注意力模块 (Temporal Attention, TA)

节点级时间自注意力：
$${\mathbf{Q}}_n^{(T)}={\mathbf{X}}_{:n:}{\mathbf{W}}_Q^T,{\mathbf{K}}_n^{(T)}={\mathbf{X}}_{:n:}{\mathbf{W}}_K^T,{\mathbf{V}}_n^{(T)}={\mathbf{X}}_{:n:}{\mathbf{W}}_V^T$$

$${\mathbf{A}}_n^{(T)}=\frac{({\mathbf{Q}}_n^{(T)})({\mathbf{K}}_n^{(T)}{)}^{\top }}{\sqrt{d^{\prime }}}$$

$$\text{TA}({\mathbf{Q}}_n^{(T)},{\mathbf{K}}_n^{(T)},{\mathbf{V}}_n^{(T)})=F({\mathbf{A}}_n^{(T)}){\mathbf{V}}_n^{(T)}$$

---

2.4 时间纠缠感知模块 (Temporal Entanglement-Aware)

四种掩码策略：

• 随机掩码、管道掩码、块掩码（同空间模块）
• 时间掩码 ${\mathbf{M}}_{temp}$：专门设计，遮蔽未来数据，强制模型仅基于历史信息重建未来交通状况

$$\text{TempTA}=F({\mathbf{A}}_t^{(S)}\odot {\mathbf{M}}_{temp}){\mathbf{V}}_t^{(S)}$$

目标：增强模型从历史到未来的时间依赖捕获能力，感知时间纠缠。

---

2.5 多头注意力融合 (Multiple Attention Fusion)

$$\text{STA}=\text{cat}\left(Z_{1\cdots h_{SRand}}^{SRand},Z_{1\cdots h_{STube}}^{STube},Z_{1\cdots h_{SBlock}}^{SBlock},Z_{1\cdots h_{short}}^{short},Z_{1\cdots h_{long}}^{long},Z_{1\cdots h_{TRand}}^{TRand},Z_{1\cdots h_{TTube}}^{TTube},Z_{1\cdots h_{TBlock}}^{TBlock},Z_{1\cdots h_{temp}}^{temp},Z_{1\cdots h_t}^t\right){\mathbf{W}}^l$$

• cat：拼接操作
• $h$：各注意力头的数量
• ${\mathbf{W}}^l$：可学习投影矩阵

后续处理：位置级前馈网络 + 层归一化 + 残差连接

---

模块三：自适应数据增强

3.1 交通流增强 (Traffic Flow Augmentation)

核心思想：基于学习的时间感知交通模式依赖，遮蔽低相关性的交通流量

掩码概率：
$${\rho }_{\tau ,n}\sim \text{Ber}(1-p_{\tau ,n})$$

• ${\rho }_{\tau ,n}$ 越高 → $x_{\tau ,n}$ 越可能被遮蔽（该时间步与该区域整体交通规律相关性低）
• 降低噪声扰动影响

3.2 道路网络增强 (Road Network Augmentation)

• 缓解低互相关交通模式区域连接的偏差
• 捕获全局城市背景下的长程区域依赖

---

模块四：自监督学习——时空异质性建模

4.1 空间异质性建模 (SHM, Spatial Heterogeneity Modeling)

核心思想：通过软聚类自监督任务，使区域嵌入捕捉空间异质性

步骤：

1. 区域投影到多个潜在表示空间（对应不同城市功能：住宅区、商业区等）

2. 生成K个聚类嵌入 {c1,…,cK}\{c_1, \ldots, c_K\}{c1​,…,cK​}

3. 聚类关联分数计算：
   $${\tilde{z}}_{i,n,k}={\mathbf{c}}_k^{\top }{\tilde{\mathbf{h}}}_{i,n}$$

其中 ${\tilde{\mathbf{h}}}_{i,n}$ 为增强数据编码的区域嵌入

4. 自监督预测任务：基于原始网络编码的区域嵌入 ${\mathbf{h}}_{i,n}$ 预测聚类分配

$${\hat{z}}_{i,n,k}={\mathbf{c}}_k^{\top }{\mathbf{h}}_{i,n}$$

5. 损失函数（带温度参数$\gamma$的软聚类）：
   $$\ell ({\mathbf{h}}_{i,n},{\tilde{\mathbf{z}}}_{i,n})=-\sum _k{\tilde{z}}_{i,n,k}\log \frac{\exp ({\hat{z}}_{i,n,k}/\gamma )}{{\sum }_j\exp ({\hat{z}}_{i,n,j}/\gamma )}$$

$${\mathcal{L}}_S=\sum _{i=1}^{t-T+1}\sum _{n=1}^N\ell ({\mathbf{h}}_{i,n},{\tilde{\mathbf{z}}}_{i,n})$$

分布正则化策略：确保聚类结果反映城市空间区域特征的真实分布。

---

4.2 时间异质性建模 (THM, Temporal Heterogeneity Modeling)

核心思想：通过对比学习区分不同时间步的异质性

步骤：

1. 融合原始与增强数据的时间维度嵌入：
   $${\mathbf{v}}_{t,n}={\mathbf{w}}_1\odot {\mathbf{h}}_{t,n}+{\mathbf{w}}_2\odot {\tilde{\mathbf{h}}}_{t,n}$$

2. 城市级表示聚合：
   $${\mathbf{s}}_t=\sigma \left(\frac{1}{N}\sum _{n=1}^N{\mathbf{v}}_{t,n}\right)$$

3. 对比学习：

   • 正样本对：同一时间步的区级嵌入 ${\mathbf{v}}_{t,n}$ 与城市级嵌入 ${\mathbf{s}}_t$
   • 负样本对：不同时间步的嵌入

4. 损失函数（交叉熵形式）：
   $${\mathcal{L}}_T=-\sum _{n=1}^N\log g({\mathbf{v}}_{t,n},{\mathbf{s}}_t)-\sum _{n=1}^N\log (1-g({\mathbf{v}}_{t^{\prime },n},{\mathbf{s}}_t))$$

其中 $g({\mathbf{v}}_{t,n},{\mathbf{s}}_t)=\sigma ({\mathbf{v}}_{t,n}^{\top }{\mathbf{W}}_3{\mathbf{s}}_t)$，${\mathbf{W}}_3$ 为可学习变换矩阵。

---

联合损失函数

$${\mathcal{L}}_{joint}={\mathcal{L}}_p+{\mathcal{L}}_s+{\mathcal{L}}_t$$

损失项	公式	功能
预测损失 ${\mathcal{L}}_p$	$\text{MSE}({\mathbf{x}}_{t+1,t+T^{\prime }},{\hat{\mathbf{x}}}_{t+1,t+T^{\prime }})$	主任务：交通流预测
空间异质性损失 ${\mathcal{L}}_s$	软聚类交叉熵	自监督：区域功能识别
时间异质性损失 ${\mathcal{L}}_t$	对比学习交叉熵	自监督：时间步判别

---

实验验证

数据集（6个真实世界数据集）

数据集	类型	节点数	边数	时间步	间隔	时间范围
PeMS04	图-高速公路	307	340	16,992	5min	2018/01-02
PeMS07	图-高速公路	883	866	28,224	5min	2017/05-08
PeMS08	图-高速公路	170	295	17,856	5min	2016/07-08
NYCTaxi	网格-城市	75 (15×5)	484	17,520	30min	2014全年
CHIBike	网格-城市	270 (15×18)	1,966	4,416	30min	2020/07-09
T-Drive	网格-城市	1,024 (32×32)	7,812	3,600	60min	2015/02-06

基准模型（19个）

类别	模型
网格模型	STResNet, DMVSTNet, DSAN
图神经网络	DCRNN, STGCN, GWNET, MTGNN, STSGCN, STFGNN, STGODE, STGNCDE
自注意力模型	STTN, GMAN, TFormer, PDFormer, ASTGNN
自监督/预训练	ST-SSL, GraphST, STD-MAE

核心实验结果

图数据集结果（表2）

模型	PeMS04 (MAE/MAPE/RMSE)	PeMS07 (MAE/MAPE/RMSE)	PeMS08 (MAE/MAPE/RMSE)
DCRNN	22.73/14.75/36.57	23.63/12.28/36.51	18.18/11.23/28.17
STGCN	21.75/13.87/34.76	22.89/11.98/35.44	17.83/11.21/27.12
GWNET	19.35/13.30/31.71	21.22/9.07/34.11	15.06/9.51/24.85
MTGNN	19.07/12.96/31.56	20.82/9.03/34.08	15.39/10.17/24.93
PDFormer	18.32/12.10/29.96	19.83/8.52/32.87	13.58/9.04/23.50
ST-SSL	18.56/12.74/31.01	20.36/9.64/33.98	14.73/10.05/24.49
STD-MAE*	17.80/12.07/29.25	18.65/8.47/31.71	13.44/8.98/22.47
SSL-STMFormer	17.06/11.32/28.11	18.53/8.41/31.44	12.05/8.11/20.15
提升幅度	4.15%/6.21%/3.89%	0.64%/0.70%/0.85%	10.34%/9.68%/10.32%

平均提升：MAE 5.04%，MAPE 5.53%，RMSE 5.02%

---

网格数据集结果（表3）

模型	NYCTaxi (In/Out)	T-Drive (In/Out)	CHIBike (In/Out)
STResNet	14.49/12.79	19.63/19.61	4.76/4.62
PDFormer	13.15/11.57	17.83/17.74	3.95/3.83
ST-SSL	12.97/9.78	16.21/16.69	4.03/3.94
GraphST	11.67/10.79	15.97/15.84	4.01/3.86
STD-MAE*	10.71/9.53	15.42/14.93	3.91/3.89
SSL-STMFormer	9.49/8.98	12.85/12.79	3.49/3.44
提升幅度	11.39%/8.50%	10.34%/14.33%	10.74%/10.18%

平均提升：MAE 10.91%，MAPE 9.86%，RMSE 7.18%

---

消融实验（表4）

变体	PeMS08 (MAE/MAPE/RMSE)	T-Drive In (MAE/MAPE/RMSE)	T-Drive Out (MAE/MAPE/RMSE)
w/o SSL（无自监督）	13.70/9.16/23.30	14.14/10.76/30.90	14.15/10.69/30.92
w/o EA（无纠缠感知）	13.59/9.14/23.52	14.20/11.04/30.79	14.21/11.00/30.78
w/o SEA（无空间纠缠）	13.46/9.00/23.22	13.51/10.60/29.83	13.47/10.64/29.80
w/o TEA（无时间纠缠）	13.40/9.12/23.32	13.47/10.16/29.61	13.38/10.06/29.50
w/ STEA（混合纠缠）	13.86/9.31/23.61	13.65/10.89/29.90	13.60/10.82/29.83
SSL-STMFormer	12.05/8.11/20.15	12.85/9.68/29.10	12.79/9.60/29.00

关键发现：

结论	证据
混合时空掩码策略劣于分离策略	w/STEA性能低于完整模型
移除任一纠缠模块均导致性能下降	w/o SEA和w/o TEA均劣于完整模型
SSL组件对捕捉时空异质性至关重要	w/o SSL性能显著下降
掩码策略对识别关键节点对不可或缺	各掩码变体均验证其必要性

---

超参数敏感性分析（图3）

参数	搜索范围	最优值	观察
隐藏维度 d	{8, 16, 32, 64, 128}	64	维度增加提升表达能力，但过大增加计算成本
编码器深度	{2, 4, 6, 8, 10}	6-8	深度增加提升复杂依赖建模，但过深导致过拟合
掩码比例	{0.05, 0.1, 0.15, 0.2, 0.25}	0.1-0.15	过低无法有效正则化，过高丢失关键信息

---

四篇文献的对比与关联

维度	第一篇（风电预测）	第二篇（调度优化）	第三篇（医疗-交通）	第四篇（交通预测）
会议/期刊	IEEE TSTE	IEEE TSTE	CAAI Trans	AAAI-25
核心问题	风电功率预测精度	风电不确定性量化	医疗-交通融合预测	交通流动态预测
技术范式	数据驱动+物理知识嵌入	深度学习+DRO优化	生成式AI+多尺度Transformer	自监督学习+时空纠缠Transformer
注意力机制	多尺度空间+时间	多教师知识蒸馏	医疗区域+应急路径	多掩码空间+时间注意力融合
自监督/对比学习	无	无	无	有（SHM+THM）
掩码策略	无	无	无	随机/管道/块/短程/长程/时间
异质性建模	误差分布形状损失	JSD模糊集	多尺度门控融合	时空软聚类+对比学习
核心创新	物理边界约束	增强模糊集	医疗感知嵌入	时空纠缠感知
损失函数	MSE+边界+分布	预测+DRO	MAE+紧急损失	MSE+聚类+对比

---

论文核心贡献总结

贡献	技术实现	效果
1. 时空注意力机制	动态捕获不同区域和时间步的依赖	增强复杂交通模式理解
2. 自监督学习任务	辅助模型捕捉交通数据底层结构	提升泛化能力
3. 时空纠缠感知模块	多种掩码策略感知动态交互	准确检测演变交通模式
4. 自适应数据增强	基于学习的时间感知依赖掩码	降低噪声，捕获长程依赖
5. 时空异质性建模	软聚类(SHM) + 对比学习(THM)	区分区域功能，识别时间模式

---

关键方法论启示

1. 纠缠感知的核心思想

• 交通系统的时空依赖不是静态预定义的
• 通过多种掩码策略强制模型学习"哪些节点对/时间步是真正重要的"
• 短程掩码捕捉局部连通性，长程掩码（DTW）捕捉功能相似性

2. 自监督学习的双重作用

• SHM：通过区域功能聚类，使嵌入具有语义解释性（住宅区、商业区等）
• THM：通过时间步对比，强化同时间段内城市级模式的一致性

3. 从静态到动态的范式转变

• 传统GNN：预定义邻接矩阵，静态图结构
• SSL-STMFormer：动态注意力 + 自适应掩码，图结构随时间和数据演变