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

Ji Qi, Mingxiao Liu, VIET HO TAM THUC DO, Yuzhe Li, Zhuoshi Pan, Gene Cheung, H. Vicky Zhao

Keywords: algorithm unrolling, graph signal processing, mixed graph, traffic forecasting

TL;DR: We unroll a mixed-graph-based optimization algorithm into a lightweight and interpretable transformer for traffic forecasting, which achieves comparable performance to SOTA schemes with drastically reduced parameters.

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.

Primary Area: optimization

Submission Number: 12372