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

TL;DR

This work tackles traffic forecasting by replacing a conventional black-box transformer with a lightweight, interpretable transformer-like network built via mixed-graph algorithm unrolling. It learns an undirected spatial graph ${\mathcal G}^u$ and a directed temporal graph ${\mathcal G}^d$, and unrolls a convex ADMM objective that includes a graph Laplacian regularizer ${\mathbf x}^T {\mathbf L}^u {\mathbf x}$, a directed-graph regularizer ${\mathcal L}_r^d$, and a directed total variation term ${\mathbf L}^d_r {\mathbf x}$, into a trainable neural architecture. The two graph-learning modules act as self-attention mechanisms, enabling data-driven edge weights with significantly fewer parameters compared to standard transformers; the resulting model achieves competitive traffic forecast performance with only about 6.4% of PDFormer’s parameters. Experiments on METR-LA and PEMS03 demonstrate robust performance across horizons and data-efficiency, highlighting practical value for resource-constrained deployments. Future work includes extending to signed distances and richer graph structures to capture more complex directed interactions.

Abstract

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 .

Figures (4)

• Figure 1: Example of a mixed graph with undirected edges (blue) connecting nodes of the same time instants, and directed edges (red) connecting nodes at $t$ to nodes in window $\{t+1, t+2\}$.
• Figure 2: Example of a 4-node DAG ${\mathcal{G}}^d$, with an added self-loop at node $1$, specified by ${\mathbf L}^d_r$ (a), and corresponding undirected graph ${\mathcal{G}}^u$, specified by ${\mathcal{L}}^d_r = ({\mathbf L}^d_r)^\top {\mathbf L}^d_r = {\mathbf L}^u$ (b). Example of a 3-node DAG (c), and a 3-node DAG with opposite directional edges (d).
• Figure 3: Unrolling of proposed iterative ADMM algorithm into blocks and neural layers.
• Figure 4: Framework of our feature extractor. (a) We first propose a non-parametric input embedding layer to combine signal with spatial/temporal embeddings. (b) Then aggregate embedded inputs of spatial neighbors and history information in a time window to generate a layer of feature extractor.

Theorems & Definitions (1)

• Theorem 3.1