Improved state mixing in higher-order and block diagonal linear recurrent networks

원문 초록

Linear recurrent networks (LRNNs) and linear state space models (SSMs) promise computational and memory efficiency on sequence modeling tasks, yet their diagonal state transitions limit expressivity. Dense and/or nonlinear architectures (e.g., LSTMs) on the other hand are provably more expressive, but computationally costly. Here, we explore how expressivity in LRNNs can be increased via richer state mixing across time and channels while maintaining competitive efficiency. Specifically, we introduce two structured LRNN architectures: (i) Higher-order Linear Recurrent Units (H-LRU), which generalize recurrences to arbitrary order, mixing multiple past states, and (ii) Block-Diagonal LRUs (BD-LRU), which enable dense intra-block channel mixing. To ensure stable training, we introduce a selective gate normalization scheme that allows for scalable window and block sizes. To maintain efficiency, we implement block-diagonal parallel scans in PyTorch and Triton that keep the throughput competitive with diagonal LRNNs for moderate orders (H-LRU) and block sizes (BD-LRU). Consistent with prior theoretical studies on the limitations of diagonal models, we empirically demonstrate that our architectures benefit from the increased expressivity of structured state mixing, on synthetic sequence modeling and on language modeling tasks that require tracking structure. Our results show that the structure of state mixing is a critical driver of performance in LRNNs, offering a practical pathway to closing the efficiency-expressivity gap in linear sequence models.