$\texttt{lrnnx}$: A library for Linear RNNs

TL;DR

The paper tackles the fragmentation of linear recurrent neural network (LRNN) implementations by introducing lrnnx, a unified PyTorch-based library that standardizes SOTA LRNN architectures under a single interface. It delivers core recurrences for both Linear Time Invariant (LTI) and Linear Time Varying (LTV) layers, supports multiple discretization schemes, and provides high-performance custom CUDA kernels to enable scalable inference. Key contributions include a uniform API across LRNN variants, rigorous validation of numerical and gradient correctness, and tutorials and infrastructure to facilitate end-to-end applications and benchmarking. By reducing engineering overhead and enabling reproducible long-range sequence modeling across domains like audio, vision, and event streams, lrnnx bridges research and deployment for LRNNs while highlighting practical limitations and avenues for future ecosystem integrations.

Abstract

Linear recurrent neural networks (LRNNs) provide a structured approach to sequence modeling that bridges classical linear dynamical systems and modern deep learning, offering both expressive power and theoretical guarantees on stability and trainability. In recent years, multiple LRNN-based architectures have been proposed, each introducing distinct parameterizations, discretization schemes, and implementation constraints. However, existing implementations are fragmented across different software frameworks, often rely on framework-specific optimizations, and in some cases require custom CUDA kernels or lack publicly available code altogether. As a result, using, comparing, or extending LRNNs requires substantial implementation effort. To address this, we introduce $\texttt{lrnnx}$, a unified software library that implements several modern LRNN architectures under a common interface. The library exposes multiple levels of control, allowing users to work directly with core components or higher-level model abstractions. $\texttt{lrnnx}$ aims to improve accessibility, reproducibility, and extensibility of LRNN research and applications. We make our code available under a permissive MIT license.

Figures (8)

• Figure 1: Class diagram describing lrnnx.
• Figure 2: Training Time Comparison
• Figure 3: LRU Training Benchmarks.
• Figure 4: S5 Training Benchmarks.
• Figure 5: Mamba Training Benchmarks.