[论文解读] Theoretical Foundations of Deep Selective State-Space Models

Nicola Muca Cirone, Antonio Orvieto+3 more|arXiv (Cornell University)|Feb 29, 2024

Complex Systems and Decision Making被引用 4

一句话总结

本文利用粗糙路径理论为深度选择性状态空间模型（SSMs）建立了理论基础，证明了输入控制的线性递推关系（如 Mamba 和 GLA 中的）可通过将输入投影到序列的签名上，显式捕捉跨时间尺度的非线性相互作用。关键结果表明，广泛且随机初始化的选择性 SSM 在无需额外 MLP 的情况下具备完全表达能力，而对角变体则通过深度高效计算高阶统计量。

ABSTRACT

Structured state-space models (SSMs) such as S4, stemming from the seminal work of Gu et al., are gaining popularity as effective approaches for modeling sequential data. Deep SSMs demonstrate outstanding performance across a diverse set of domains, at a reduced training and inference cost compared to attention-based transformers. Recent developments show that if the linear recurrence powering SSMs allows for multiplicative interactions between inputs and hidden states (e.g. GateLoop, Mamba, GLA), then the resulting architecture can surpass in both in accuracy and efficiency attention-powered foundation models trained on text, at scales of billion parameters. In this paper, we give theoretical grounding to this recent finding using tools from Rough Path Theory: we show that when random linear recurrences are equipped with simple input-controlled transitions (selectivity mechanism), then the hidden state is provably a low-dimensional projection of a powerful mathematical object called the signature of the input -- capturing non-linear interactions between tokens at distinct timescales. Our theory not only motivates the success of modern selective state-space models such as Mamba but also provides a solid framework to understand the expressive power of future SSM variants.

研究动机与目标

• 为利用粗糙路径理论工具分析现代选择性 SSM 提供理论框架。
• 证明 SSM 中的输入控制线性递推关系可完全表达逼近输入序列的连续函数的能力。
• 证明对角 SSM（如 Mamba）通过深度高效计算高阶统计量，其表达能力在深度极限下可与密集变体相匹配。
• 基于受控微分方程（CDEs）将 S4、Mamba、GLA 及相关模型统一于同一数学框架下进行分析。
• 为理解与设计未来具备更强表达能力与更高效率的 SSM 变体提供可推广的理论基础。

提出的方法

• 将 SSM 建模为由输入序列驱动的受控微分方程（CDEs），其中隐藏状态通过依赖输入的线性变换演化。
• 应用粗糙路径理论，将隐藏状态形式化为输入签名的低维投影，以捕捉非线性时间相互作用。
• 将状态转移矩阵 $ W_{s,t} $ 定义为输入驱动向量场的迭代积分之和，推广了 S4 和 Mamba 的递推关系。
• 利用 Wronskian 矩阵与路径积分表示，推导含加法输入项的 CDE 的显式解。
• 证明随机初始化的密集输入控制递推关系可通过隐藏状态捕捉足够的输入统计量，从而实现完全表达能力。
• 证明将对角 SSM 模块级联可计算高阶全局统计量，在深度极限下逼近密集模型的表达能力。

Figure 1 : Example of the dynamics of the S4 and Mamba CDEs derived in Sec. 3.1 . The continuous-time approximation of these algorithms can be written as a Linear CDE. In S4, $\omega^{\text{X}}_{t}=t$ , while $\xi^{\text{X}}_{t}=\int_{0}^{t}X_{s}ds$ . In Mamba, $\omega^{\text{X}}_{t}=\int_{0}^{t}\si

实验结果

研究问题

• RQ1 如何从理论上为 Mamba 等现代选择性 SSM 的表达能力提供依据？
• RQ2 选择性 SSM 能够跨长序列建模非线性相互作用的数学结构是什么？
• RQ3 在不引入额外非线性组件（如 MLP）的情况下，输入控制的线性递推关系能否实现完全表达能力？
• RQ4 对角 SSM 与密集变体在捕捉高阶输入统计量方面有何差异？
• RQ5 能否基于受控微分方程与粗糙路径理论，构建统一框架以分析和比较多种 SSM 架构？

主要发现

• 选择性 SSM 中的输入控制线性递推关系可显式将隐藏状态表示为输入签名的低维投影，从而捕捉非线性时间相互作用。
• 广泛且随机初始化的密集输入控制 SSM 具备完全表达能力：无需额外 MLP 即可逼近任意从输入序列到输出的连续函数。
• 如 Mamba 等对角 SSM 通过深度高效计算高阶统计量，在深度极限下其表达能力可与密集模型相媲美。
• 选择性 SSM 的隐藏状态在数学上等价于求解受控微分方程（CDE），其显式解可通过路径积分表示与 Wronskian 矩阵推导得出。
• 该框架将 S4、Mamba、GLA 及相关模型统一于同一理论基础之上，可实现对新型 SSM 变体的严谨比较与分析。
• 该理论解释了为何选择性 SSM 在长程建模中优于标准 S4 和注意力机制：其通过结构化、输入相关的转移机制捕捉了更丰富的输入统计量。

Figure 3 : Comparison of Linear CDE, Mamba and S5 on anti-symmetric signature prediction tasks. As suggested by the theory Mamba struggles to immediately capture these high order statistics, but performance improves with chaining . On the other hand S5 does not learn such non-linear features even af