2025-11-20T09:28:14.240195

Lightweight and Interpretable Transformer via Mixed Graph Algorithm Unrolling for Traffic Forecast

Qi, Do, Liu et al.

Unlike conventional "black-box" transformers with classical self-attention mechanism, we build a lightweight and interpretable transformer-like neural net by unrolling a mixed-graph-based optimization algorithm to forecast traffic with spatial and temporal dimensions. We construct two graphs: an undirected graph $\mathcal{G}^u$ capturing spatial correlations across geography, and a directed graph $\mathcal{G}^d$ capturing sequential relationships over time. We predict future samples of signal $\mathbf{x}$, assuming it is "smooth" with respect to both $\mathcal{G}^u$ and $\mathcal{G}^d$, where we design new $\ell_2$ and $\ell_1$-norm variational terms to quantify and promote signal smoothness (low-frequency reconstruction) on a directed graph. We design an iterative algorithm based on alternating direction method of multipliers (ADMM), and unroll it into a feed-forward network for data-driven parameter learning. We insert graph learning modules for $\mathcal{G}^u$ and $\mathcal{G}^d$ that play the role of self-attention. Experiments show that our unrolled networks achieve competitive traffic forecast performance as state-of-the-art prediction schemes, while reducing parameter counts drastically. Our code is available in https://github.com/SingularityUndefined/Unrolling-GSP-STForecast .

academic

Lightweight and Interpretable Transformer via Mixed Graph Algorithm Unrolling for Traffic Forecast

基本信息

论文ID: 2505.13102
标题: Lightweight and Interpretable Transformer via Mixed Graph Algorithm Unrolling for Traffic Forecast
作者: Ji Qi, Mingxiao Liu, Tam Thuc Do, Yuzhe Li, Zhuoshi Pan, Gene Cheung, H. Vicky Zhao
分类: cs.LG cs.AI eess.SP
发表时间: 2025年10月12日 (arXiv v2)
论文链接: https://arxiv.org/abs/2505.13102

摘要

本文提出了一种基于混合图算法展开的轻量级可解释Transformer模型用于交通预测。与传统的"黑盒"Transformer不同，该方法通过展开混合图优化算法构建了一个可解释的类Transformer神经网络。模型构建了两个图：无向图 $\mathcal{G}^u$ 捕捉地理空间相关性，有向图 $\mathcal{G}^d$ 捕捉时序关系。通过设计新的 $\ell_2$ 和 $\ell_1$ 范数变分项来量化和促进有向图上的信号平滑性，并基于交替方向乘子法(ADMM)设计迭代算法，将其展开为前馈网络进行数据驱动的参数学习。实验表明，该模型在保持竞争性交通预测性能的同时，大幅减少了参数数量。

研究背景与动机

问题定义

交通预测是一个重要的时空数据建模问题，需要同时捕捉：

空间相关性：地理位置相近的监测站点之间的相关性
时间依赖性：历史观测对未来的影响关系

现有方法的局限性

传统Transformer：参数量巨大，缺乏可解释性，在实际部署中面临计算和内存约束
基于模型的方法：往往独立处理空间和时间维度，未能充分利用时空关系
现有深度学习方法：虽然性能优秀但仍是"黑盒"模型，参数量大

研究动机

工业应用对轻量级模型的迫切需求
算法展开(Algorithm Unrolling)提供了模型驱动与数据驱动相结合的新范式
现有工作仅使用正无向图，无法有效建模复杂的时空关系

核心贡献

首次提出混合图算法展开：结合无向图(空间)和有向图(时间)来建模复杂的时空关系
创新的有向图正则化项：设计了有向图拉普拉斯正则化器(DGLR)和有向图总变分(DGTV)
轻量级可解释Transformer：通过ADMM算法展开实现了参数大幅减少(仅为PDFormer的6.4%)
理论贡献：证明了有向图频率定义在无权有向线图情况下退化为经典傅里叶频率

方法详解

任务定义

给定N个监测站点在过去T+1个时刻的观测值，预测未来S个时刻的交通状态。输入为部分观测的时空信号 $y \in \mathbb{R}^M$ ，输出为完整的时空信号 $x \in \mathbb{R}^{N(T+S+1)}$ 。

混合图构建

无向图 $\mathcal{G}^u$

连接相同时刻地理位置相近的节点
捕捉空间相关性
使用对称邻接矩阵 $W^u$

有向图 $\mathcal{G}^d$

从时刻 $\tau$ 的节点连接到 $\tau+1, ..., \tau+W$ 时刻的相同节点
捕捉时间因果关系
使用非对称邻接矩阵 $W^d$

有向图变分项设计

$\ell_2$ 范数项：有向图拉普拉斯正则化器(DGLR)

$x^T\mathcal{L}_r^d x = x^T(L_r^d)^T L_r^d x = \|x - W_r^d x\|_2^2$

其中 $L_r^d = I - W_r^d$ 为随机游走拉普拉斯矩阵， $W_r^d = (D^d)^{-1}W^d$ 为行随机邻接矩阵。

$\ell_1$ 范数项：有向图总变分(DGTV)

$\|L_r^d x\|_1 = \sum_{j \in \bar{S}} |x_j - \sum_i w_{j,i} x_i|$

优化目标函数

$\min_x \|y - Hx\|_2^2 + \mu_u x^T L^u x + \mu_{d,2} x^T \mathcal{L}_r^d x + \mu_{d,1} \|L_r^d x\|_1$

其中 $H$ 为采样矩阵， $\mu_u, \mu_{d,2}, \mu_{d,1}$ 为权重参数。

ADMM算法设计

通过引入辅助变量 $\phi$ ，将优化问题转化为： $\min_{x,\phi} \|y - Hx\|_2^2 + \mu_u x^T L^u x + \mu_{d,2} x^T \mathcal{L}_r^d x + \mu_{d,1} \|\phi\|_1$ $\text{s.t. } \phi = L_r^d x$

子问题求解

$x$ 子问题：通过共轭梯度法求解线性系统
$\phi$ 子问题：软阈值运算 $\phi_i^{\tau+1} = \text{sign}(\delta) \cdot \max(|\delta| - \rho^{-1}\mu_{d,1}, 0)$ 其中 $\delta = (L_r^d)_i x^{\tau+1} - \rho^{-1}\gamma_i^\tau$

图学习模块

无向图学习(UGL)

使用Mahalanobis距离计算节点相似性： $d^u(i,j) = (f_i^u - f_j^u)^T M (f_i^u - f_j^u)$

边权重通过归一化指数函数计算： $w_{i,j}^u = \frac{\exp(-d^u(i,j))}{\sqrt{\sum_{l \in \mathcal{N}_i} \exp(-d^u(i,l))} \sqrt{\sum_{k \in \mathcal{N}_j} \exp(-d^u(k,j))}}$

有向图学习(DGL)

类似地使用度量矩阵 $P$ 计算有向边权重。

网络架构

将ADMM的每次迭代实现为神经层：

5个ADMM块，每块25层
每块前插入图学习模块
使用多头注意力机制(4个并行图学习模块)

实验设置

数据集

METR-LA：洛杉矶交通速度数据，207个节点，1315条边
PEMS03：交通流量数据，358个节点，547条边
采样间隔：5分钟
数据划分：6:2:2(训练:验证:测试)

评价指标

RMSE：均方根误差
MAE：平均绝对误差
MAPE：平均绝对百分比误差

对比方法

包括6类基线方法：

基于模型：VAR
GNN方法：STGCN, STSGCN
GAT方法：GMAN, ST-Wave
Transformer方法：PDFormer, STAEformer
自适应图方法：Graph WaveNet, AGCRN
简单线性模型：STID, SimpleTM

实现细节

预测时长：30/60/120分钟(6/12/24步)
历史窗口：60分钟(12步)
优化器：Adam，学习率5×10⁻⁴
损失函数：Huber损失(δ=1)
硬件：NVIDIA GeForce RTX 3090

实验结果

主要结果

数据集	时长	本文方法	最佳基线	参数量对比
PEMS03	30min	26.10/17.03/18.85	23.71/15.05/18.16	34K vs 531K
PEMS03	60min	27.67/17.46/17.72	25.56/15.97/15.49	(6.4%参数量)
METR-LA	60min	12.34/5.18/11.80	11.96/5.49/9.65

关键发现

参数效率：仅使用PDFormer 6.4%的参数量达到竞争性性能
长期预测优势：预测时长越长，与最佳方法的性能差距越小
数据效率：在数据稀缺情况下性能更稳定

消融实验

变体	PEMS03 (RMSE/MAE/MAPE)	METR-LA (RMSE/MAE/MAPE)
完整模型	27.67/17.46/17.72	12.34/5.18/11.80
无DGTV	27.78/17.85/17.90	12.36/5.40/12.31
无DGLR	30.89/20.02/21.10	12.41/5.35/12.20
无向时间图	27.52/17.87/18.82	12.51/5.42/12.11