让非线性序列模型并行计算,速度媲美线性模型。
Parallelization of Non-linear State-Space Models: Scaling Up Liquid-Resistance Liquid-Capacitance Networks for Efficient Sequence Modeling
- 通过强制雅可比矩阵对角化,实现长序列并行求解。
- 在长序列预测任务中优于Transformer、Mamba等模型。
- 兼具稳定梯度和可推广性,适合高效序列建模场景。
我们提出LrcSSM,一种非线性循环模型,能以与当前线性状态空间层相当的速度处理长序列。通过强制雅可比矩阵为对角阵,整个序列可并行求解,计算复杂度为𝒪(TD),内存占用为𝒪(TD),仅需𝒪(log T)的串行深度,其中T为输入序列长度,D为状态维度。此外,LrcSSM提供形式化的梯度稳定性保障,这是液态-阻抗-电容系统(如Liquid-S4和Mamba)等输入依赖型系统所不具备的。重要的是,该模型的对角雅可比结构在性能上不损失于原始的密集雅可比模型,且可推广至其他非线性循环模型,展现出更广适用性。在一系列长程预测任务中,LrcSSM的表现优于Transformer、LRU、S5和Mamba。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。