Title: Lightweight and Interpretable Transformer via Unrolling of Mixed Graph Algorithms for Traffic Forecast
Primary area: optimization
Keywords: algorithm unrolling, graph signal processing, mixed graph, traffic forecasting

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.

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.