Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Parallelization of Non-linear State-Space Models: Scaling Up Liquid-Resistance Liquid-Capacitance Networks for Efficient Sequence Modeling

Mónika Farsang, Ramin Hasani, Daniela Rus, Radu Grosu

TL;DR

The paper introduces LrcSSM, a non-linear state-space model with a diagonal Jacobian that enables exact parallel updates for long sequences, achieving O(TD) compute and O(log T) depth while providing a formal gradient-stability guarantee. By designing the model so that its state-transition Jacobian is inherently diagonal, the authors realize a scalable parallelization approach that matches or surpasses leading linear SSMs and Transformer variants on long-horizon tasks. Empirically, LrcSSM demonstrates strong performance on EthanolConcentration and competitive results across six datasets, and ablations show the importance of both state- and input-dependent terms as well as the robustness of the diagonal design. The work further shows that the diagonal-model design can be generalized to other non-linear RNNs, offering a broad, biologically inspired pathway to efficient sequence modeling with stable training dynamics.

Abstract

We present LrcSSM, a $\textit{non-linear}$ recurrent model that processes long sequences as fast as today's linear state-space layers. By forcing its Jacobian matrix to be diagonal, the full sequence can be solved in parallel, giving $\mathcal{O}(TD)$ computational work and memory and only $\mathcal{O}(\log T)$ sequential depth, for input-sequence length $T$ and a state dimension $D$. Moreover, LrcSSM offers a formal gradient-stability guarantee that other input-varying systems such as Liquid-S4 and Mamba do not provide. Importantly, the diagonal Jacobian structure of our model results in no performance loss compared to the original model with dense Jacobian, and the approach can be generalized to other non-linear recurrent models, demonstrating broader applicability. On a suite of long-range forecasting tasks, we demonstrate that LrcSSM outperforms Transformers, LRU, S5, and Mamba.

Parallelization of Non-linear State-Space Models: Scaling Up Liquid-Resistance Liquid-Capacitance Networks for Efficient Sequence Modeling

TL;DR

The paper introduces LrcSSM, a non-linear state-space model with a diagonal Jacobian that enables exact parallel updates for long sequences, achieving O(TD) compute and O(log T) depth while providing a formal gradient-stability guarantee. By designing the model so that its state-transition Jacobian is inherently diagonal, the authors realize a scalable parallelization approach that matches or surpasses leading linear SSMs and Transformer variants on long-horizon tasks. Empirically, LrcSSM demonstrates strong performance on EthanolConcentration and competitive results across six datasets, and ablations show the importance of both state- and input-dependent terms as well as the robustness of the diagonal design. The work further shows that the diagonal-model design can be generalized to other non-linear RNNs, offering a broad, biologically inspired pathway to efficient sequence modeling with stable training dynamics.

Abstract

We present LrcSSM, a recurrent model that processes long sequences as fast as today's linear state-space layers. By forcing its Jacobian matrix to be diagonal, the full sequence can be solved in parallel, giving computational work and memory and only sequential depth, for input-sequence length and a state dimension . Moreover, LrcSSM offers a formal gradient-stability guarantee that other input-varying systems such as Liquid-S4 and Mamba do not provide. Importantly, the diagonal Jacobian structure of our model results in no performance loss compared to the original model with dense Jacobian, and the approach can be generalized to other non-linear recurrent models, demonstrating broader applicability. On a suite of long-range forecasting tasks, we demonstrate that LrcSSM outperforms Transformers, LRU, S5, and Mamba.

Paper Structure

This paper contains 54 sections, 2 theorems, 14 equations, 2 figures, 11 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Background
  3. Bio-inspired Liquid Neural Networks
  4. Parallelization Techniques
  5. Scaling Up Non-linear LRCs
  6. Our Parallelization.
  7. Proposed Model
  8. Comparison to Linear State Space Models
  9. Comparison to Liquid-Resistance Liquid-Capacitance Networks (LRCs)
  10. Theoretical Insights
  11. Related Work
  12. Linear Structural State-Space Models.
  13. Parallelizing Non-linear Recurrent Neural Networks.
  14. Positioning LrcSSM in Recent Advances.
  15. Experiments
  16. ...and 39 more sections

Key Result

Lemma 1

For any $x,y\in\mathbb{R}^{D}$ we have $\|x_{t+1}-y_{t+1}\|_2 \;=\; \|\lambda_t(x-y)\|_2 \;\le\; \rho\,\|x-y\|_2$.

Figures (2)

  • Figure 1: Liquid-Resistance Liquid-Capacitance SSM (LrcSSM) architecture. The input sequence of length $T$ and input dimension $p$ is first passed through an input encoder, followed by a normalization layer. The core component is a non-linear, state-and-input dependent LRC with hidden dimension $D$ and sequence length $T$. This SSM is computed by a parallelizable iterative linearization method. The final state values are then processed by an MLP, with a skip connection added to preserve information flow. The LrcSSM block can be stacked and repeated an arbitrary number of times (we use ${2, 4, 6}$ layers in our experiments). A post-normalization layer is applied before the output is passed to the decoder, which produces the final output.
  • Figure 2: Iterations needed for convergence for LrcSSM on average per dataset.

Theorems & Definitions (4)

  • Lemma 1: $\rho$‚Äìcontraction
  • proof
  • Theorem 1: Gradient stability
  • proof