Lightweight Transformer via Unrolling of Mixed Graph Algorithms for Traffic Forecast

2025-05-19

AI Summary

This work addresses spatiotemporal traffic flow forecasting with a lightweight, interpretable Transformer-like neural network. Methodologically, it jointly models spatial and temporal dependencies via an undirected geographical graph and a directed temporal graph. It introduces the first ℓ₁/ℓ₂-mixed graph signal smoothness metric tailored for directed graphs and formulates a joint optimization problem incorporating this regularizer; this problem is then differentiably unfolded into a feedforward network using the Alternating Direction Method of Multipliers (ADMM). A graph learning module replaces self-attention to enable topology adaptation and parameter efficiency. Contributions include the first integration of graph signal smoothness modeling with ADMM unfolding on directed spatiotemporal graphs, achieving both interpretability and low computational complexity. The method achieves state-of-the-art accuracy on multiple real-world datasets, reduces model parameters by 72%, accelerates inference by 3.1×, and the code is publicly available.

Technology Category

• Planning, Routing, and Scheduling: Optimization of Spatio-temporal Systems
• Machine Learning: Graph-based Machine Learning
• Knowledge Representation and Reasoning: Geometric, Spatial, and Temporal Reasoning

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Social Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks

Abstract

To forecast traffic with both spatial and temporal dimensions, we unroll a mixed-graph-based optimization algorithm into a lightweight and interpretable transformer-like neural net. Specifically, 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 formulate a prediction problem for the 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 construct 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$, which are akin to the self-attention mechanism in classical transformers. 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.

Problem

• Forecasting traffic with spatial and temporal dimensions
• Designing lightweight transformer via unrolling mixed-graph algorithms
• Promoting signal smoothness on directed and undirected graphs

Innovation

• Unroll mixed-graph algorithm into lightweight transformer
• Design new variational terms for directed graph smoothness
• Insert graph learning modules akin to self-attention