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 utilize a parallel-scan implementation that keeps 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 in both synthetic sequence modeling and language modeling that our architectures significantly benefit from the increased expressivity of structured state mixing. 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.