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

AI summary

Overview

Research area: Deep learning software infrastructure for sequence modeling — specifically linear recurrent neural networks (LRNNs) and state-space models (SSMs).

Technical level: Intermediate. The paper is a software-library paper rather than a new-architecture paper; it assumes familiarity with recurrent models, state-space parameterizations, and PyTorch, but the design and benchmark discussion is readable without deep signal-processing background.

One-sentence scope: The paper introduces lrnnx, a unified PyTorch library that implements several modern linear RNN architectures behind a common interface, with custom CUDA kernels and benchmarks against original releases.

What This Paper Is About

LRNN research has produced many architectures (S4, S5, LRU, S6, S7, RG-LRU, STREAM, Centaurus, Event-SSM), but their implementations are scattered across different frameworks, sometimes depend on custom CUDA kernels, and in several cases are not publicly released at all. This fragmentation means that using, comparing, or extending LRNNs requires substantial re-engineering effort. The paper's goal is to consolidate those architectures into a single library with a shared interface, control over discretization and parameterization, and performance competitive with the original implementations.

Key Contributions

1. A unified interface for LRNN layers. lrnnx standardizes multiple LRNN architectures under a consistent interface derived from a single state-space equation (eq. 2), so switching between state-space parameterizations or discretization schemes amounts to instantiating a different class without changing surrounding training or evaluation code.

2. A three-tier inheritance hierarchy with decoupled discretization. A base LRNN class defines the forward interface and selects discretization; LTI_LRNN and LTV_LRNN subclasses organize layers by time-invariance, and each subclass defines its own parameterization of (A, B, C). Supported discretization schemes include ZOH, bilinear, dirac, and asynchronous (event-driven) discretization, with some models restricting the allowed set (for example, Centaurus uses only ZOH).

3. High-performance custom CUDA kernels for time-varying layers. Derived from the selective scan implementation in Mamba, these kernels integrate multiple discretization methods (ZOH, bilinear, dirac) and support asynchronous inputs within a fused scan and output projection. The paper positions this as an advantage over some JAX implementations, which it says suffer memory bottlenecks from materializing the hidden state.

4. Multi-level abstractions plus end-to-end model heads and tutorials. The library exposes both low-level building blocks (core recurrences) and higher-level modules with activations and skip connections, along with components and tutorials for language modeling, classification, and autoencoders. An LRNNLMHeadModel wraps an LRNN backbone with embeddings, stacked residual blocks, and a language modeling head, and its mixer_types argument allows mixing LRNN backends and attention layers.

Main Findings

• Fragmented implementations are the motivating problem: Table 1 catalogues nine contemporary SSM architectures and their availability. S4 (PyTorch), S5 (JAX), Event-SSM (JAX), S6 (PyTorch), STREAM (PyTorch), and Centaurus (PyTorch) are marked as public; LRU, RG-LRU, and S7 are marked as not public (N/A implementation).
• SISO/LTI classification: Table 1 distinguishes architectures by whether they are single-input single-output (S4, S6, STREAM) or not, and whether they are linear time invariant (S4, S5, LRU, Event-SSM) or not (S6, STREAM, RG-LRU, S7, Centaurus). Centaurus is described as being in-between SISO and MIMO models.
• Benchmark setup: All GPU benchmarks ran on an NVIDIA A100 40GB GPU with Python 3.12 and CUDA 12.9.
• Comparison targets: The library's LRU implementation (PyTorch) was compared against a popular public repository (JAX); the S5 implementation against the original release; and the Mamba implementation against the official repository. Benchmarks used random tensors.
• Timing protocol: For each configuration, the authors ran 10 warm-up passes, then timed 90 forward passes; this was repeated for 5 experiments, reporting the mean and standard deviation across those 5 experiment means. The same sweep settings (batch sizes, sequence lengths, model dimensions) were mirrored across all three models, and the state dimension was set to 16 wherever required.
• Training performance: The paper states that its implementations are competitive with the public baselines (Figure 2, Training Time Comparison), with full benchmark plots in Appendix A (Figures 3–8 covering LRU, S5, and Mamba training and inference benchmarks).
• Inference performance: The authors report that their models match other public implementations on training speed but are slightly slower for inference, which they attribute to known CPU overheads in PyTorch inference execution rather than model-specific design choices. They expect inference performance to be very similar for production workloads in high batch size and long sequence length regimes.
• Correctness validation: Numerical equivalence is validated between parallel, recurrent, and step-wise execution modes for every layer via a test suite spanning sequence lengths, batch sizes, model dimensions, initializations, and discretizations; gradient consistency is also verified between custom CUDA kernels and reference PyTorch implementations.
• Inference acceleration approach: Because PyTorch lacks a native analogue of jax.lax.scan, the library uses CUDA Graphs to avoid CPU synchronization after each step. The paper describes this as competitive at large sequence lengths and adding only a few ms at small ones.

Methodology in Plain English

The authors did not propose a new recurrent architecture. Instead, they surveyed the existing LRNN and SSM literature, identified which models had public implementations and which did not, and then re-implemented a set of them inside one PyTorch codebase.

All layers were anchored to the same state-space recurrence, with learning concentrated in how the matrices A, B, and C are parameterized. Layers were sorted by whether their matrices change over time: time-invariant layers (S4, S5, LRU, Event-SSM) inherit from an LTI_LRNN class and use optimal einsum contractions for efficiency gains, while time-varying layers (S6, S7, RG-LRU, STREAM, Centaurus) inherit from LTV_LRNN.

Discretization was deliberately separated from layer definition, so a user can pick ZOH, bilinear, dirac, or asynchronous event-driven discretization without rewriting the layer. Every layer gets the same constructor signature and a step method for autoregressive generation.

For the time-varying models, the team wrote custom CUDA kernels adapted from Mamba's selective scan, extended to handle multiple discretization schemes and asynchronous inputs inside a fused scan-plus-output-projection operation. To build confidence that parallel, recurrent, and step-by-step paths agree numerically, they ran a test suite that varies sequence length, batch size, model dimension, initialization, and discretization, and checked that gradients from the CUDA kernels match reference PyTorch gradients.

Finally, they benchmarked three implementations — LRU, S5, and Mamba — against the corresponding public reference implementations under an identical sweep of batch size, sequence length, and model dimension, reporting averaged execution times for training (forward plus backward) and autoregressive inference.

Why This Matters

Impact on research. The paper argues that inconsistent frameworks and missing code make it hard to reproduce results, benchmark models under consistent conditions, or extend beyond a single architecture. A shared interface turns architecture-swapping into a constructor change, which lowers the barrier to controlled comparison studies and to building hybrids (the mixer_types argument already allows mixing LRNN backends with attention layers in one model).

Real-world applications (as listed in the paper's applications section):

• Audio: text-to-speech, automatic speech recognition, and speech enhancement.
• RNA modeling.
• Vision.
• Event streams and point clouds.

Industry relevance. LRNNs offer O(1) inference-time complexity, in contrast to the O(n) inference cost of Transformer attention with global attention over sequence length n, while training in a parallelizable way. The paper notes that Transformers are typically hard to train for very long sequences (≥ 2^10), which is precisely where these layers prove extremely useful. That combination — cheap inference plus long-context handling — is what makes the library relevant to deployment. The library also integrates with standard PyTorch workflows including checkpointing, gradient checkpointing, mixed-precision training, and fused operations, and is released under a permissive MIT license (the arXiv listing shows a CC BY-NC-SA 4.0 license).

Future Directions

• Framework coverage. The library is restricted to PyTorch, precluding direct use by researchers in the JAX or TensorFlow communities; extending to other backends is an open direction.
• Hardware portability. Optimal performance for several LTV layers relies on custom CUDA kernels, limiting optimal execution to NVIDIA hardware and hindering accessibility on alternative backends.
• Ecosystem integrations. There are no native wrappers for established tools such as Hugging Face, DeepSpeed, or FSDP, so large-scale distributed workflows currently require manually developed adapter layers.
• Unimplemented architectural features. Bidirectional variants of LRNN layers are not yet provided, although the base interface is designed to support them. Non-linear RNNs such as xLSTM are described as orthogonal to the library's focus and have not been implemented. Closing the reported inference-speed gap versus other public implementations is also an implicit target.

Target Audience

This paper is most useful to practitioners and researchers who want to use, compare, or extend linear RNN and state-space models without re-implementing them: sequence-modeling researchers benchmarking SSM variants, engineers integrating long-sequence or low-latency inference models into audio, vision, event-stream, or RNA pipelines, and PyTorch developers who want drop-in LRNN layers. Readers evaluating library choices will find the architecture availability table and the benchmark protocol most valuable; readers looking for new model architectures or new theory will not find them here.

Authors’ 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.