作者：苏剑林
单位：科学空间
研究方向：NLP、神经网络
前几天，笔者看了几篇介绍 SSM（State Space Model）的文章，才发现自己从未认真了解过 SSM，于是决定深入学习这一领域，并记录下学习所得。SSM 的概念由来已久，但这里我们特别指深度学习中的 SSM。可以说，2021年的 S4（Structured State Space for Sequence Modeling）是 SSM 的开篇之作，而最近最火的变体大概是去年的 Mamba。
SSM 的基础与发展
SSM 的概念并不新鲜，但其在深度学习领域的应用却是近几年的事。当我们谈到 SSM 时，通常指的是一切线性 RNN 模型，如 RWKV、RetNet 以及 Google 推出的 LRU 等。这些模型的目标之一是成为 Transformer 的竞争者，尽管完全替代 Transformer 的可能性不大，但 SSM 本身优雅的数学性质值得深入研究。
HiPPO 的引入
在 S4 之前，SSM 的奠基之作是《HiPPO: Recurrent Memory with Optimal Polynomial Projections》（简称 HiPPO）。HiPPO 提出了用正交基逼近动态更新的函数，其结果是一个线性系统。这不仅告诉我们线性系统可以逼近复杂的函数，还提供了具体的逼近方法和近似程度。
SSM 的基本形式
对于已经了解 SSM 的读者，SSM 建模所用的线性 ODE 系统公式如下：
$$
\dot{x}(t) = A x(t) + B u(t)
$$
其中，$x(t)$ 是系统状态，$A$ 和 $B$ 是矩阵，$u(t)$ 是输入信号。
即使将其离散化，线性 RNN 模型的实质依然是线性系统。这引出了一个自然的问题：为什么选择线性系统？线性系统够用吗？
答案是肯定的。线性系统既足够简单，也足够复杂。简单是指线性化常常是复杂系统的基本近似；复杂是指即便是简单的线性系统，也可以拟合复杂的函数。
举个例子，考虑一个简单的线性系统：
$$
\dot{x}(t) = \lambda x(t)
$$
其解为 $x(t) = x(0) e^{\lambda t}$，这意味着通过组合指数函数和三角函数，线性系统可以拟合复杂的函数。
有限压缩和函数逼近
我们考虑区间 $[0, t]$ 内的输入信号 $u(\tau)$，HiPPO 的目标是将其压缩为一个有限维的向量。通过傅里叶级数展开，我们可以将 $u(\tau)$ 压缩为一个有限维向量，但这只适用于静态区间。
而实际中，我们需要处理的是动态信号 $u(t)$，这就需要将傅里叶级数公式推广到动态情况。具体的方法是通过映射 $t$ 到 $[0, 1]$ 区间，然后计算正交基展开系数，并对其进行动态更新。
线性系统的本质
通过傅里叶级数和勒让德多项式等正交多项式展开，我们可以得出线性 ODE 系统的系数动力学。具体来说，当我们用正交基去逼近动态函数时，其结果自然导出一个线性 ODE 系统。
以勒让德多项式为例，其递归公式和边界值满足条件，能够在实数空间中简化推导过程。通过对函数基的正交化处理，我们得到的系数动力学满足线性 ODE 系统。
HiPPO 的一般框架
HiPPO 是一个自下而上的框架，通过正交基逼近函数，导出线性 ODE 系统。它并没有假设系统必须是线性的，而是通过逼近的角度反推出线性系统的能力。
HiPPO 的主要结论基于勒让德多项式，但也可以扩展到其他正交多项式，如切比雪夫多项式、拉盖尔多项式等。
结论
本文通过简单易懂的方式重温了 HiPPO 的主要推导。HiPPO 通过适当的记忆假设，自下而上地导出了线性 ODE 系统，并为 SSM 的发展奠定了基础。尽管 SSM 还在不断发展，其优雅的数学性质和强大的函数逼近能力使其成为深度学习中的重要工具。
参考文献
[1] Efficiently Modeling Long Sequences with Structured State Spaces
[2] Recurrent Memory with Optimal Polynomial Projections
[3] Legendre Polynomials
[4] Gram-Schmidt Process
[5] Kronecker Delta