Improved state mixing in higher-order and block diagonal linear recurrent networks

📄 arXiv: 2602.12021v1 📥 PDF

作者: Igor Dubinin, Antonio Orvieto, Felix Effenberger

分类: cs.LG

发布日期: 2026-02-12

💡 一句话要点

提出高阶和块对角线性循环网络，提升长序列建模的效率与表达能力。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)

关键词: 线性循环网络 状态空间模型 长序列建模 高阶递归 块对角结构 状态混合 序列模型

📋 核心要点

  1. 线性循环网络在长序列建模中效率高，但表达能力受限，无法媲美LSTM等非线性模型。
  2. 论文提出高阶线性循环单元(H-LRU)和块对角线性循环单元(BD-LRU)，增强状态混合，提升表达能力。
  3. 实验表明，BD-LRU在合成序列任务中性能优于现有线性模型，H-LRU在压缩任务中参数效率更高。

📝 摘要（中文）

线性循环网络(LRNNs)和线性状态空间模型(SSMs)在长序列建模任务中具有计算和内存效率的优势，但其对角状态转移限制了表达能力。另一方面，密集和非线性架构(如LSTMs)在表达能力上更强，但计算成本较高。本文探讨了如何在保持竞争效率的同时，通过更丰富的跨时间和通道的状态混合来提高LRNNs的表达能力。具体而言，我们提出了两种结构化的LRNN架构：(i)高阶线性循环单元(H-LRU)，它将一阶递归推广到高阶，混合多个过去的状态；(ii)块对角线性循环单元(BD-LRU)，它支持密集的块内通道混合。选择性门控的逐通道(H-LRU)或逐行(BD-LRU)L1-归一化稳定了训练，并允许缩放窗口/块大小。所提出的架构的并行扫描实现使得吞吐量与中等阶数(H-LRU)和块大小(BD-LRU)的对角LRNNs具有竞争力。在合成序列建模任务中，BD-LRU的性能与线性SSMs (Mamba)、低秩LRNNs (DeltaNet)和LSTM基线相匹配或超过它们，而H-LRU在压缩任务中被发现是最具参数效率的。在合成序列建模和语言建模中，我们的结果表明，状态混合的结构而非宽度本身决定了LRNNs的表达能力，为缩小线性序列模型中的效率-表达能力差距提供了一条实用途径。

🔬 方法详解

问题定义：线性循环网络（LRNNs）虽然在长序列建模中具有计算和内存效率优势，但由于其状态转移矩阵的对角结构，限制了模型的表达能力。现有的密集型或非线性模型（如LSTM）虽然表达能力更强，但计算成本也更高。因此，如何在保持LRNNs效率的同时，提升其表达能力是一个关键问题。

核心思路：论文的核心思路是通过引入更丰富的状态混合机制来增强LRNNs的表达能力。具体来说，通过设计特定的结构，使得模型能够混合多个过去的状态（H-LRU）或者进行块内的密集通道混合（BD-LRU），从而突破对角结构的限制。这种结构化的状态混合方式旨在在效率和表达能力之间取得平衡。

技术框架：论文提出了两种新的LRNN架构：H-LRU和BD-LRU。H-LRU将一阶递归推广到高阶，允许模型混合多个过去的状态信息。BD-LRU则采用块对角结构，在每个块内进行密集的通道混合。为了稳定训练和允许缩放窗口/块大小，论文还采用了选择性门控的L1归一化。此外，论文还针对提出的架构设计了并行扫描实现，以保持较高的吞吐量。

关键创新：论文的关键创新在于提出了两种结构化的状态混合方式：高阶递归和块对角混合。这两种方法都旨在突破传统LRNNs对角结构的限制，从而提升模型的表达能力。与直接增加模型宽度相比，论文强调了状态混合的结构对于提升表达能力的重要性。

关键设计：H-LRU的关键设计在于高阶递归的实现方式，需要仔细设计如何混合多个过去的状态。BD-LRU的关键设计在于块大小的选择和块内混合的实现方式。此外，选择性门控的L1归一化对于稳定训练至关重要，其归一化方式（逐通道或逐行）也需要根据具体架构进行调整。并行扫描实现也是一个重要的技术细节，它保证了模型在具有更高表达能力的同时，仍然能够保持较高的计算效率。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，BD-LRU在合成序列建模任务中，性能与线性SSMs (Mamba)、低秩LRNNs (DeltaNet)和LSTM基线相匹配或超过它们。H-LRU在压缩任务中表现出最高的参数效率。这些结果表明，所提出的结构化状态混合方法能够有效提升LRNNs的表达能力，并在某些任务中达到甚至超过现有模型的性能。

🎯 应用场景

该研究成果可应用于各种长序列建模任务，如自然语言处理、语音识别、时间序列预测等。通过提升线性循环网络的表达能力，可以更有效地处理长程依赖关系，提高模型性能。该研究对于开发更高效、更强大的序列模型具有重要意义。

📄 摘要（原文）

Linear recurrent networks (LRNNs) and linear state space models (SSMs) promise computational and memory efficiency on long-sequence modeling tasks, yet their diagonal state transitions limit expressivity. Dense and 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 first-order recurrence to higher order, mixing multiple past states, and (ii) Block-Diagonal LRUs (BD-LRU), which enable dense intra-block channel mixing. Per-channel (H-LRU) or per-row (BD-LRU) L1-normalization of selective gates stabilizes training and allows for scaling window/block sizes. A parallel-scan implementation of the proposed architectures keeps the throughput competitive with diagonal LRNNs for moderate orders (H-LRU) and block sizes (BD-LRU). In synthetic sequence modeling tasks, the performance of BD-LRU matches or exceeds those of linear SSMs (Mamba), low-rank LRNNs (DeltaNet) and LSTM baselines, while H-LRU is found to be the most parameter-efficient in compression task. In both synthetic sequence modeling and language modeling, our results indicate that the structure of state mixing rather than width alone shapes expressivity of LRNNs, offering a practical route to closing the efficiency-expressivity gap in linear sequence models.