Lightweight and Interpretable Transformer via Unrolling of Mixed Graph Algorithms for Traffic Forecast

Unlike conventional "black-box" transformers with classical self-attention mechanisms, we build a lightweight and interpretable transformer-like neural network 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 periodically 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. Code: https://github.com/SingularityUndefined/Unrolling-GSP-STForecast.

We challenge the parameter efficiency of modern transformer architectures in this work. Conventional transformers rely on dense self-attention mechanisms and large black-box feed-forward networks (FFNs), which often result in high computational cost and limited interpretability. Using city-scale traffic forecasting as a case study, we develop a lightweight and interpretable transformer-like neural network by unrolling a mixed-graph smoothing algorithm. Despite using only 7.2% of the parameters of the transformer-based PDFormer model, our approach achieves competitive performance against a wide range of strong baselines. Our method models traffic data using a mixed graph composed of two parts: an undirected graph that captures spatial relationships between road sensors, and a directed graph that models how traffic evolves over time. Building on ideas from graph signal processing, we design smoothing priors for both graph types — including a new prior for directed graphs — and formulate traffic forecasting as a graph signal recovery problem. We solve this optimization problem using an iterative ADMM algorithm, whose operations can be interpreted as a sequence of low-pass graph filters over the learned mixed graph. By unrolling these iterative steps into neural network layers and learning only a small number of filter parameters, we obtain an interpretable and parameter-efficient alternative to conventional transformer FFNs. We further show that transformer attention learning is closely related to Gaussian-kernel graph learning, allowing us to implement local attention structures at dramatically lower parameter cost. Our results suggest that many parameter-heavy transformer components can be replaced by structured graph-based operations with strong mathematical interpretations. Although we evaluate our approach on traffic forecasting, the broader goal of this work is to study how directed relationships — analogous to directional attention scores — can be modeled using graph learning and graph signal processing techniques. We believe these lightweight and interpretable transformer designs may generalize to a much broader range of applications beyond traffic systems.