重温被Mamba带火的SSM：线性系统和HiPPO矩阵 2024-06-21 作者 C3P00 近年来，深度学习领域涌现出许多新的模型架构，其中状态空间模型（SSM，State Space Model）因其优雅的数学性质和强大的表达能力，逐渐成为Transformer的强劲对手。而Mamba，作为SSM最新的变体，更是凭借其在长序列建模上的优异表现，吸引了众多研究者的关注。 本文将带您深入了解SSM的核心概念，并重温其重要奠基之作——HiPPO（High-order Polynomial Projection Operators）。通过HiPPO的推导，您可以理解线性系统在SSM中的重要性，以及它如何通过有限维的向量来储存无限维的函数信息。 线性系统：简单而强大的表达能力 SSM的核心思想是利用线性微分方程（ODE）系统来建模序列数据。一个典型的线性 ODE 系统可以表示为： $$\frac{dh}{dt} = Ah + Bu$$ 其中，$h$ 代表系统的状态，$u$ 代表输入，$A$ 和 $B$ 是模型参数。 那么，为什么SSM会选择线性系统呢？答案是：线性系统既足够简单，也足够复杂。 简单是指，线性化通常是复杂系统的一个最基本近似。而复杂是指，即使是如此简单的系统，也可以拟合异常复杂的函数。例如，一个简单的线性系统： $$\frac{dh}{dt} = h$$ 其解为 $h(t) = h(0)e^t$。这意味着，只要时间足够长，该线性系统就可以通过指数函数来拟合足够复杂的函数。 HiPPO：从正交基逼近到线性系统 HiPPO 为我们提供了一种更本质的理解：当我们试图用正交基去逼近一个动态更新的函数时，其结果就是如上的线性系统。 假设我们要用一个有限维的向量来储存一段信号 $x(t)$ 的信息。如果我们假设 $x(t)$ 在某点 $t_0$ 阶可导，那么它对应的 $t_0$ 阶泰勒展开式往往是 $x(t)$ 的良好近似。我们可以只储存展开式的系数，从而将 $x(t)$ 压缩为一个有限维向量。 然而，实际遇到的数据通常无法满足“阶可导”这种苛刻的条件。因此，我们更倾向于使用正交函数基展开，比如傅里叶级数。其系数计算公式为： $$c_k = \int_{-\infty}^{\infty} x(t)e^{-2\pi ikt} dt$$ 通过只保留有限个系数，我们可以将 $x(t)$ 压缩为一个有限维向量。 接下来，问题难度升级。实际中，$x(t)$ 代表的是持续采集的信号，所以它是不断有新数据进入的。我们需要更新逼近结果来记忆整个信号的历史。 为了解决这个问题，我们可以将 $x(t)$ 映射到一个有限区间 $[0, T]$，然后计算其在该区间上的傅里叶系数。 $$c_k(T. = \int_{0}^{T} x(t)e^{-2\pi ikt} dt✅$$ 当新的数据进入时，我们可以重新计算系数，从而更新对 $x(t)$ 的逼近。 通过对系数的导数进行推导，我们可以发现，系数的变化满足一个线性 ODE 系统。这意味着，当我们试图用傅里叶级数去记忆一个实时函数的最邻近窗口内的状态时，结果自然而言地导致了一个线性 ODE 系统。 HiPPO 矩阵：勒让德多项式的应用 HiPPO 的核心是选取多项式为基函数。其中，勒让德多项式因其在实数空间中的定义和简化推导过程的优势，成为了 HiPPO 的首选。 勒让德多项式 $P_n(x)$ 是关于 $x$ 的 $n$ 次函数，定义域为 $[-1, 1]$，满足： $$\int_{-1}^{1} P_m(x)P_n(x) dx = \frac{2}{2n + 1}\delta_{mn}$$ 通过将勒让德多项式作为基函数，并利用其递归公式，我们可以推导出一个恒等式： $$\frac{d}{dT}P_n(2T/T – 1) = \frac{2n}{T}P_{n-1}(2T/T – 1)$$ 利用该恒等式，我们可以得到 HiPPO 矩阵，它描述了系数随时间的变化规律。 HiPPO 的应用：SSM 的基石 HiPPO 的结论被后来诸多 SSM 模型所使用，例如 S4 和 Mamba。HiPPO 为我们提供了一种理解 SSM 的全新视角，它揭示了线性系统在 SSM 中的本质意义，以及它如何通过有限维的向量来储存无限维的函数信息。 总结 本文以尽可能简单的方式重复了 HiPPO 的主要推导。HiPPO 通过适当的记忆假设，自下而上地导出了线性 ODE 系统，并针对勒让德多项式的情形求出了相应的解析解（HiPPO 矩阵）。其结果被后来诸多 SSM 模型使用，可谓是 SSM 的重要奠基之作。 参考文献 [1] https://papers.cool/arxiv/2312.00752[2] https://papers.cool/arxiv/2305.13048[3] https://papers.cool/arxiv/2307.08621[4] https://papers.cool/arxiv/2008.07669[5] https://dblp.org/pid/130/0612.html[6] https://en.wikipedia.org/wiki/Kronecker_delta[7] https://en.wikipedia.org/wiki/Legendre_polynomials[8] https://en.wikipedia.org/wiki/Gram–Schmidt_process[9] https://en.wikipedia.org/wiki/Chebyshev_polynomials[10] https://en.wikipedia.org/wiki/Laguerre_polynomials[11] https://proceedings.neurips.cc/paper/2019/file/952285b9b7e7a1be5aa7849f32ffff05-Paper.pdf
近年来，深度学习领域涌现出许多新的模型架构，其中状态空间模型（SSM，State Space Model）因其优雅的数学性质和强大的表达能力，逐渐成为Transformer的强劲对手。而Mamba，作为SSM最新的变体，更是凭借其在长序列建模上的优异表现，吸引了众多研究者的关注。
本文将带您深入了解SSM的核心概念，并重温其重要奠基之作——HiPPO（High-order Polynomial Projection Operators）。通过HiPPO的推导，您可以理解线性系统在SSM中的重要性，以及它如何通过有限维的向量来储存无限维的函数信息。
线性系统：简单而强大的表达能力
SSM的核心思想是利用线性微分方程（ODE）系统来建模序列数据。一个典型的线性 ODE 系统可以表示为：
$$
\frac{dh}{dt} = Ah + Bu
$$
其中，$h$ 代表系统的状态，$u$ 代表输入，$A$ 和 $B$ 是模型参数。
那么，为什么SSM会选择线性系统呢？答案是：线性系统既足够简单，也足够复杂。
简单是指，线性化通常是复杂系统的一个最基本近似。而复杂是指，即使是如此简单的系统，也可以拟合异常复杂的函数。例如，一个简单的线性系统：
$$
\frac{dh}{dt} = h
$$
其解为 $h(t) = h(0)e^t$。这意味着，只要时间足够长，该线性系统就可以通过指数函数来拟合足够复杂的函数。
HiPPO：从正交基逼近到线性系统
HiPPO 为我们提供了一种更本质的理解：当我们试图用正交基去逼近一个动态更新的函数时，其结果就是如上的线性系统。
假设我们要用一个有限维的向量来储存一段信号 $x(t)$ 的信息。如果我们假设 $x(t)$ 在某点 $t_0$ 阶可导，那么它对应的 $t_0$ 阶泰勒展开式往往是 $x(t)$ 的良好近似。我们可以只储存展开式的系数，从而将 $x(t)$ 压缩为一个有限维向量。
然而，实际遇到的数据通常无法满足“阶可导”这种苛刻的条件。因此，我们更倾向于使用正交函数基展开，比如傅里叶级数。其系数计算公式为：
$$
c_k = \int_{-\infty}^{\infty} x(t)e^{-2\pi ikt} dt
$$
通过只保留有限个系数，我们可以将 $x(t)$ 压缩为一个有限维向量。
接下来，问题难度升级。实际中，$x(t)$ 代表的是持续采集的信号，所以它是不断有新数据进入的。我们需要更新逼近结果来记忆整个信号的历史。
为了解决这个问题，我们可以将 $x(t)$ 映射到一个有限区间 $[0, T]$，然后计算其在该区间上的傅里叶系数。
$$
c_k(T. = \int_{0}^{T} x(t)e^{-2\pi ikt} dt✅
$$
当新的数据进入时，我们可以重新计算系数，从而更新对 $x(t)$ 的逼近。
通过对系数的导数进行推导，我们可以发现，系数的变化满足一个线性 ODE 系统。这意味着，当我们试图用傅里叶级数去记忆一个实时函数的最邻近窗口内的状态时，结果自然而言地导致了一个线性 ODE 系统。
HiPPO 矩阵：勒让德多项式的应用
HiPPO 的核心是选取多项式为基函数。其中，勒让德多项式因其在实数空间中的定义和简化推导过程的优势，成为了 HiPPO 的首选。
勒让德多项式 $P_n(x)$ 是关于 $x$ 的 $n$ 次函数，定义域为 $[-1, 1]$，满足：
$$
\int_{-1}^{1} P_m(x)P_n(x) dx = \frac{2}{2n + 1}\delta_{mn}
$$
通过将勒让德多项式作为基函数，并利用其递归公式，我们可以推导出一个恒等式：
$$
\frac{d}{dT}P_n(2T/T – 1) = \frac{2n}{T}P_{n-1}(2T/T – 1)
$$
利用该恒等式，我们可以得到 HiPPO 矩阵，它描述了系数随时间的变化规律。
HiPPO 的应用：SSM 的基石
HiPPO 的结论被后来诸多 SSM 模型所使用，例如 S4 和 Mamba。HiPPO 为我们提供了一种理解 SSM 的全新视角，它揭示了线性系统在 SSM 中的本质意义，以及它如何通过有限维的向量来储存无限维的函数信息。
总结
本文以尽可能简单的方式重复了 HiPPO 的主要推导。HiPPO 通过适当的记忆假设，自下而上地导出了线性 ODE 系统，并针对勒让德多项式的情形求出了相应的解析解（HiPPO 矩阵）。其结果被后来诸多 SSM 模型使用，可谓是 SSM 的重要奠基之作。
参考文献
[1] https://papers.cool/arxiv/2312.00752
[2] https://papers.cool/arxiv/2305.13048
[3] https://papers.cool/arxiv/2307.08621
[4] https://papers.cool/arxiv/2008.07669
[5] https://dblp.org/pid/130/0612.html
[6] https://en.wikipedia.org/wiki/Kronecker_delta
[7] https://en.wikipedia.org/wiki/Legendre_polynomials
[8] https://en.wikipedia.org/wiki/Gram–Schmidt_process
[9] https://en.wikipedia.org/wiki/Chebyshev_polynomials
[10] https://en.wikipedia.org/wiki/Laguerre_polynomials
[11] https://proceedings.neurips.cc/paper/2019/file/952285b9b7e7a1be5aa7849f32ffff05-Paper.pdf