概述
HiPPO(High-order Polynomial Projection Operators)是目前大热的structured state space model (S4)及其后续工作的backbone. State space mode主要是控制学科里的内容,最近被引入深度学习领域来解决长距离依赖问题。长距离依赖建模的核心问题是如何通过有限的memory来尽可能记住之前所有的历史信息。当前的主流序列建模模型(即Transformer和RNN) 存在着普遍的遗忘问题
- fixed-size context windows: Transformer的window size通常是有限的,一般来说quadratic的attention最多建模到大约10k的token就到计算极限了
- vanishing gradient: RNN通过hidden state来存储历史信息,理论上能记住之前所有内容,但实际上的effective memory大概是<1k个token的level,可能的原因是gradient vanishing
HiPPO 通过数学方法分析来得到closed-form solution,并用数学证明HiPPO不会遗忘历史信息。
快速理解
- HiPPO的核心思想是用有限维向量储存连续函数的信息。
- 当我们试图用正交基去逼近一个动态更新的函数时,其结果就是线性ODE系统
- HiPPO不仅告诉我们线性系统可以逼近复杂函数,还告诉我们如何逼近以及近似程度。
HiPPO 框架利用了逼近理论中的经典工具,如正交多项式。最终,解决方案采取了一种简单的线性微分方程的形式,被称为 HiPPO 算子
基本形式
对于事先已经对SSM有所了解的读者,想必知道SSM建模所用的是线性ODE系统:
\[\begin{equation}\begin{aligned}
x'(t) =&\, A x(t) + B u(t) \\
y(t) =&\, C x(t) + D u(t)
\end{aligned}\tag{1}
\end{equation}\]
其中\(u(t)\in\mathbb{R}^{d_i}, x(t)\in\mathbb{R}^{d}, y(t)\in\mathbb{R}^{d_o}, A\in\mathbb{R}^{d\times d}, B\in\mathbb{R}^{d\times d_i}, C\in\mathbb{R}^{d_o\times d}, D\in\mathbb{R}^{d_o\times d_i}\)。当然我们也可以将它离散化,那么就变成一个线性RNN模型,这部分我们在后面的文章再展开。不管离散化与否,其关键词都是“线性”,那么马上就有一个很自然的问题:为什么是线性系统?线性系统够了吗?
我们可以从两个角度回答这个问题:线性系统既足够简单,也足够复杂。
- 简单是指从理论上来说,线性化往往是复杂系统的一个最基本近似,所以线性系统通常都是无法绕开的一个基本点;
- 复杂是指即便如此简单的系统,也可以拟合异常复杂的函数,为了理解这一点,我们只需要考虑一个 \(\mathbb{R}^4\) 的简单例子:
\[x'(t) =\begin{pmatrix} 1 & 0 & 0 & 0 \\
0 & -1 & 0 & 0 \\
0 & 0 & 0 & 1 \\
0 & 0 & -1 & 0
\end{pmatrix}x(t)\]
这个例子的基本解是 \(x(t) = (e^t, e^{-t}, \sin t, \cos t)\)。这意味着什么呢?意味着只要d足够大,该线性系统就可以通过指数函数和三角函数的组合来拟合足够复杂的函数,而我们知道拟合能力很强的傅里叶级数也只不过是三角函数的组合,如果在加上指数函数显然就更强了,因此可以想象线性系统也有足够复杂的拟合能力。
当然,这些解释某种意义上都是“马后炮”。HiPPO给出的结果更加本质:当我们试图用正交基去逼近一个动态更新的函数时,其结果就是如上的线性系统。这意味着,HiPPO不仅告诉我们线性系统可以逼近足够复杂的函数,还告诉我们怎么去逼近,甚至近似程度如何。
有限压缩
接下来,我们都只考虑 \(d_i=1\) 的特殊情形,\(d_i > 1\) 只不过是 \(d_i=1\) 时的平行推广。此时,\(u(t)\) 的输出是一个标量,进一步地,作为开头我们先假设 \(t\in[0, 1]\),HiPPO的目标是:用一个有限维的向量来储存这一段 \(u(t)\) 的信息。
看上去这是一个不大可能的需求,因为 \(t\in[0,1]\) 意味着 \(u(t)\) 可能相当于无限个点组成的向量,压缩到一个有限维的向量可能严重失真。不过,如果我们对 \(u(t)\) 做一些假设,并且允许一些损失,那么这个压缩是有可能做到的,并且大多数读者都已经尝试过。比如,当 \(u(t)\) 在某点 \(n+1\) 阶可导的,它对应的 \(n\) 阶泰勒展开式往往是 \(u(t)\) 的良好近似,于是我们可以只储存展开式的 \(n+1\) 个系数来作为 \(u(t)\) 的近似表征,这就成功将 \(u(t)\) 压缩为一个 \(n+1\) 维向量。
当然,对于实际遇到的数据来说,“\(n+1\) 阶可导”这种条件可谓极其苛刻,我们通常更愿意使用在平方可积条件下的正交函数基展开,比如傅里叶(Fourier)级数,它的系数计算公式为
\[\begin{equation}c_n = \int_0^1 u(t) e^{-2i\pi n t}dt \tag{2}
\end{equation}\]
这时候取一个足够大的整数 \(N\),只保留 \(|n|\leq N\) 的系数,那么就将 \(u(t)\) 压缩为一个 \(2N + 1\) 维的向量了。
接下来,问题难度就要升级了。刚才我们说 \(t\in[0,1]\),这是一个静态的区间,而实际中 \(u(t)\) 代表的是持续采集的信号,所以它是不断有新数据进入的,比如现在我们近似了 \([0,1]\) 区间的数据,马上就有 \([1,2]\) 的数据进来,你需要更新逼近结果来试图记忆整个\([0,2]\) 区间,接下来是 \([0,3]\)、\([0,4]\) 等等,这我们称为“在线函数逼近”。而上面的傅里叶系数(2),只适用于区间 \([0,1]\),因此需要将它进行推广。
为此,我们设 \(t\in[0,T]\),\(s\mapsto t_{\leq T}(s)\) 是 \([0,1]\) 到 \([0,T]\) 的一个映射,那么 \(u(t_{\leq T}(s))\) 作为 \(s\) 的函数时,它的定义区间就是 \([0,1]\),于是就可以复用(2):
\[\begin{equation}c_n(T) = \int_0^1 u(t_{\leq T}(s)) e^{-2i\pi n s}ds \tag{3}
\end{equation}\]
这里我们已经给系数加了标记\((T)\),以表明此时的系数会随着T的变化而变化。
线性初现
能将\([0,1]\) 映射到\([0,T]\)的函数有无穷多,而最终结果也因 \(t_{\leq T}(s)\) 而异,一些比较直观且相对简单的选择如下:
- \(t_{\leq T}(s) = sT\),即将 \([0,1]\) 均匀地映射到 \([0,T]\);
- 注意 \(t_{\leq T}(s)\) 并不必须是满射,所以像 \(t_{\leq T}(s)=s + T - 1\) 也是允许的,这意味着只保留了最邻近窗口 \([T-1,T]\) 的信息,丢掉了更早的部分,更一般地有 \(t_{\leq T}(s)=sw + T - w\),其中\(w\) 是一个常数,这意味着 \(T-w\) 前的信息被丢掉了;
- 也可以选择非均匀映射,比如 \(t_{\leq T}(s) = T\sqrt{s}\),它同样是 \([0,1]\) 到 \([0,T]\) 的满射,但 \(s=1/4\)时就映射到 \(T/2\) 了,这意味着我们虽然关注全局的历史,但同时更侧重于 \(T\) 时刻附近的信息。
现在我们以\(t_{\leq T}(s)=sw + T - w\) 为例,代入 (3) 得到
\[\begin{equation}c_n(T) = \int_0^1 u(sw + T - w) e^{-2i\pi n s}ds\tag{4}
\end{equation}\]
现在我们两边求关于\(T\)的导数:\(v(x) = \frac{e^{-2i\pi ns}}{-2i\pi n}\)
\[\begin{equation}\begin{aligned}
\frac{d}{dT}c_n(T) =&\, \int_0^1 u'(sw + T - w) e^{-2i\pi n s}ds \\
=&\, \left.\frac{1}{w} u(sw + T - w) e^{-2i\pi n s}\right|_{s=0}^{s=1} + \frac{2i\pi n}{w}\int_0^1 u(sw + T - w) e^{-2i\pi n s}ds \\
=&\, \frac{1}{w} u(T) - \frac{1}{w} u(T-w) + \frac{2i\pi n}{w} c_n(T) \\
\end{aligned}\ \ \ \ \ \tag{5}
\end{equation}\]
其中第二个等号我们用了分部积分公式。由于我们只保留了 \(|n|\leq N\) 的系数,所以根据傅立叶级数的公式,可以认为如下是 \(u(sw + T - w)\) 的一个良好近似:
\[\begin{equation}u(sw + T - w) \approx \sum_{k=-N}^{k=N} c_k(T) e^{2i\pi k s}\tag{6}
\end{equation}\]
那么 \(u(T - w) = u(sw + T - w)|_{s=0}\approx \sum\limits_{k=-N}^{k=N} c_k(T)\),代入 (5) 得:
\[\begin{equation}\frac{d}{dT}c_n(T) \approx \frac{1}{w} u(T) - \frac{1}{w} \sum_{k=-N}^{k=N} c_k(T) + \frac{2i\pi n}{w} c_n(T)\tag{7}
\end{equation}\]
将 \(T\) 换成 \(t\),然后所有的 \(c_n(t)\) 堆在一起记为 \(x(t) = (c_{-N},c_{-(N-1)},\cdots,c_0,\cdots,c_{N-1},c_N)\),并且不区分 \(\approx\) 和 \(=\),那么就可以写出
\[\begin{equation}x'(t) = Ax(t) + Bu(t),\quad A_{n,k} = \left\{\begin{array}{l}(2i\pi n - 1)/w, &k=n \\ -1/w,&k\neq n\end{array}\right.,\quad B_n = 1/w\tag{8}
\end{equation}\]
这就出现了如 (1) 所示的线性ODE系统。即当我们试图用傅里叶级数去记忆一个实时函数的最邻近窗口内的状态时,结果自然而言地导致了一个线性ODE系统。
HIPPO推导
一般框架
当然,目前只是选择了一个特殊的 \(t_{\leq T}(s)\),换一个 \(t_{\leq T}(s)\) 就不一定有这么简单的结果了。此外,傅里叶级数的结论是在复数范围内的,进一步实数化也可以,但形式会变得复杂起来。所以,我们要将上一节的过程推广成一个一般化的框架,从而得到更一般、更简单的纯实数结论。
设 \(t\in[a,b]\),并且有目标函数 \(u(t)\) 和函数基 \(\{g_n(t)\}_{n=0}^N\),我们希望有后者的线性组合来逼近前者,目标是最小化 \(L_2\) 距离:
\[\begin{equation}\mathop{\text{argmin}}_{c_1,\cdots,c_N}\int_a^b \left[u(t) - \sum_{n=0}^N c_n g_n(t)\right]^2 dt\tag{9}
\end{equation}\]
这里我们主要在实数范围内考虑,所以方括号直接平方就行,不用取模。更一般化的目标函数还可以再加个权重函数\(\rho(t)\),但我们这里就不考虑了,毕竟HiPPO的主要结论其实也没考虑这个权重函数。
对目标函数展开,得到
\[\begin{equation}\int_a^b u^2(t) dt - 2\sum_{n=0}^N c_n \int_a^b u(t) g_n(t)dt + \sum_{m=0}^N\sum_{n=0}^N c_m c_n \int_a^b g_m(t) g_n(t) dt\tag{10}
\end{equation}\]
这里我们只考虑标准正交函数基,其定义为 \(\int_a^b g_m(t) g_n(t) dt = \delta_{m,n}\),\(\delta_{m,n}\) 是克罗内克δ函数,此时上式可以简化成
\[\begin{equation}\int_a^b u^2(t) dt - 2\sum_{n=0}^N c_n \int_a^b u(t) g_n(t)dt + \sum_{n=0}^N c_n^2 \tag{11}
\end{equation}\]
这只是一个关于 \(c_n\) 的二次函数,它的最小值是有解析解的:
\[\begin{equation}c^*_n = \int_a^b u(t) g_n(t)dt\tag{12}
\end{equation}\]
这也被称为 \(u(t)\) 与 \(g_n(t)\) 的内积,它是有限维向量空间的内积到函数空间的平行推广。简单起见,在不至于混淆的情况下,我们默认 \(c_n\) 就是 \(c^*_n\)。
接下来的处理跟上一节是一样的,我们要对一般的 \(t\in[0, T]\) 考虑 \(u(t)\) 的近似,那么找一个 \([a,b]\) 到 \([0,T]\) 的映射 \(s\mapsto t_{\leq T}(s)\),然后计算系数
\[\begin{equation}c_n(T) = \int_a^b u(t_{\leq T}(s)) g_n(s) ds\tag{13}
\end{equation}\]
同样是两边求T的导数,然后用分部积分法
\[\begin{equation}\scriptsize\begin{aligned}
\frac{d}{dT}c_n(T) =&\, \int_a^b u'(t_{\leq T}(s)) \frac{\partial t_{\leq T}(s)}{\partial T} g_n(s) ds = \int_a^b \left(\frac{\partial t_{\leq T}(s)}{\partial T}\left/\frac{\partial t_{\leq T}(s)}{\partial s}\right.\right) g_n(s) d u(t_{\leq T}(s)) \\
=&\,\left.u(t_{\leq T}(s))\left(\frac{\partial t_{\leq T}(s)}{\partial T}\left/\frac{\partial t_{\leq T}(s)}{\partial s}\right.\right) g_n(s)\right|_{s=a}^{s=b} - \int_a^b u(t_{\leq T}(s)) \,d\left[\left(\frac{\partial t_{\leq T}(s)}{\partial T}\left/\frac{\partial t_{\leq T}(s)}{\partial s}\right.\right) g_n(s)\right]
\end{aligned}\ \ \ \ \ \ \ \tag{14}
\end{equation}\]
准备工作—勒让德多项式
接下来的计算,就依赖于 \(g_n(t)\) 和 \(t_{\leq T}(s)\) 的具体形式了。HiPPO的全称是High-order Polynomial Projection Operators,第一个P正是多项式(Polynomial)的首字母,所以HiPPO的关键是选取多项式为基。现在我们请出继傅里叶之后又一位大牛——勒让德(Legendre),接下来我们要选取的函数基正是以他命名的“勒让德多项式”。
勒让德多项式 \(p_n(t)\) 是关于 \(t\) 的 \(n\) 次函数,定义域为 \([-1,1]\),满足
\[\begin{equation}\int_{-1}^1 p_m(t) p_n(t) dt = \frac{2}{2n+1}\delta_{m,n}\tag{15}
\end{equation}\]
所以 \(p_n(t)\) 之间只是正交,还不是标准(平分积分为1),\(g_n(t)=\sqrt{\frac{2n+1}{2}} p_n(t)\)才是标准正交基。
当我们对函数基 \(\{1,t,t^2,\cdots, t^n\}\) 执行施密特正交化时,其结果正是勒让德多项式。相比傅里叶基,勒让德多项式的好处是它是纯粹定义在实数空间中的,并且多项式的形式能够有助于简化部分\(t_{\leq T}(s)\) 的推导过程,这一点我们后面就可以看到。勒让德多项式有很多不同的定义和性质,这里我们不一一展开,有兴趣的读者自行看链接中维基百科介绍即可。
接下来我们用到两个递归公式来推导一个恒等式,这两个递归公式是
\[\begin{align}
p_{n+1}'(t) - p_{n-1}'(t) = (2n+1)p_n(t)\tag{16}\\[5pt]
p_{n+1}'(t) = (n + 1)p_n(t) + t p_n'(t)\tag{17}
\end{align}\]
由第一个(16) 迭代得到:
\[\begin{equation}\begin{aligned}
p_{n+1}'(t) =&\, (2n+1)p_n(t) + (2n-3)p_{n-2}(t) + (2n-7)p_{n-4}(t) + \cdots \\
=&\, \sum_{k=0}^n (2k+1) \chi_{n-k} p_k(t)
\end{aligned}\tag{18}
\end{equation}\]
其中当 \(k\) 是偶数时 \(\chi_k=1\) 否则 \(\chi_k=0\)。代入第二个 (17) 得到
\[\begin{equation}t p_n'(t) = n p_n(t) + (2n-3)p_{n-2}(t) + (2n-7)p_{n-4}(t) + \cdots\tag{19}
\end{equation}\]
继而有
\[\begin{equation}\begin{aligned}
(t+1) p_n'(t) =&\, n p_n(t) + (2n-1)p_{n-1}(t) + (2n-3)p_{n-2}(t) + \cdots\\
=&\,-(n+1) p_n(t) + \sum_{k=0}^n (2k + 1) p_k(t)
\end{aligned}\tag{20}
\end{equation}\]
这些就是等会要用到的恒等式。此外,勒让德多项式满足 \(p_n(1)=1,p_n(-1)=(-1)^n\),这个边界值后面也会用到。
正如n维空间中不止有一组正交基也一样,正交多项式也不止有勒让德多项式一种,比如还有切比雪夫(Chebyshev)多项式,如果算上加权的目标函数(即\(\rho(t)\not\equiv 1\)),还有拉盖尔多项式等,这些在原论文中都有提及,但HiPPO的主要结论还是基于勒让德多项式展开的,所以剩余部分这里也不展开讨论了。
邻近窗口(LegT)
完成准备工作后,我们就可以代入具体的 \(t_{\leq T}(s)\) 进行计算了,计算过程跟傅里叶级数的例子大同小异,只不过基函数换成了勒让德多项式构造的标准正交基 \(g_n(t)=\sqrt{\frac{2n+1}{2}} p_n(t)\)。作为第一个例子,我们同样先考虑只保留最邻近窗口的信息,此时 \(t_{\leq T}(s) = (s + 1)w / 2 + T - w\) 将 \([-1,1]\) 映射到 \([T-w,T]\),原论文将这种情形称为“LegT(Translated Legendre)”。
直接代入 (14),马上得到
\[\small\frac{d}{dT}c_n(T) = \frac{\sqrt{2(2n+1)}}{w}\left[u(T) - (-1)^n u(T-w)\right] - \frac{2}{w}\int_{-1}^1 u((s + 1)w / 2 + T - w) g_n'(s) ds\]
我们首先处理 \(u(T-w)\) 项,跟傅里叶级数那里同样的思路,我们截断 \(n\leq N\) 作为 \(u((s + 1)w / 2 + T - w)\) 的一个近似:
\[\begin{equation}u((s + 1)w / 2 + T - w)\approx \sum_{k=0}^N c_k(T)g_k(s)\tag{21}
\end{equation}\]
从而有 \(u(T-w)\approx \sum\limits_{k=0}^N c_k(T)g_k(-1) = \sum\limits_{k=0}^N (-1)^k c_k(T) \sqrt{\frac{2k+1}{2}}\)。接着,利用 (16) 得到
\[\begin{equation}\begin{aligned}
&\,\int_{-1}^1 u((s + 1)w / 2 + T - w) g_n'(s) ds \\
=&\,\int_{-1}^1 u((s + 1)w / 2 + T - w) \sqrt{\frac{2n+1}{2}} p_n'(s) ds \\
=&\, \int_{-1}^1 u((s + 1)w / 2 + T - w)\sqrt{\frac{2n+1}{2}}\left[\sum_{k=0}^{n-1} (2k+1) \chi_{n-1-k} p_k(s)\right]ds \\
=&\, \int_{-1}^1 u((s + 1)w / 2 + T - w)\sqrt{\frac{2n+1}{2}}\left[\sum_{k=0}^{n-1} \sqrt{2(2k+1)} \chi_{n-1-k} g_k(s)\right]ds \\
=&\, \sqrt{2n+1}\sum_{k=0}^{n-1} \sqrt{2k+1} \chi_{n-1-k} c_k(T)
\end{aligned}\tag{22}
\end{equation}\]
将这些结果整合起来,就有
\[\begin{equation}\begin{aligned}
\frac{d}{dT}c_n(T) \approx &\, \frac{\sqrt{2(2n+1)}}{w}u(T) - \frac{\sqrt{2(2n+1)}}{w} (-1)^n \overbrace{\sum\limits_{k=0}^N (-1)^k c_k(T) \sqrt{\frac{2k+1}{2}}}^{u(T-w)} \\
&\quad- \frac{2}{w}\overbrace{\sqrt{2n+1}\sum_{k=0}^{n-1} \sqrt{2k+1} \chi_{n-1-k} c_k(T)}^{\int_{-1}^1 u((s + 1)w / 2 + T - w) g_n'(s) ds} \\[12pt]
= &\, \frac{\sqrt{2(2n+1)}}{w}u(T) - \frac{\sqrt{2n+1}}{w} \sum\limits_{k=0}^N (-1)^{n-k} c_k(T) \sqrt{2k+1} \\
&\quad- \frac{2}{w}\sqrt{2n+1}\sum_{k=0}^{n-1} \sqrt{2k+1} \chi_{n-1-k} c_k(T) \\[12pt]
= &\, \frac{\sqrt{2(2n+1)}}{w}u(T) - \frac{\sqrt{2n+1}}{w} \sum\limits_{k=n}^N (-1)^{n-k} c_k(T) \sqrt{2k+1} \\
&\quad- \frac{\sqrt{2n+1}}{w}\sum_{k=0}^{n-1} \sqrt{2k+1} \underbrace{\left(2\chi_{n-1-k} + (-1)^{n-k}\right)}_{\equiv 1}c_k(T) \\
\end{aligned}\tag{23}
\end{equation}\]
再次地,将 \(T\) 换回 \(t\),并将所有的 \(c_n(t)\) 堆在一起记为 \(x(t) = (c_0,c_1,\cdots,c_N)\),那么根据上式可以写出
\[\begin{equation}\begin{aligned}
x'(t) =&\, Ax(t) + Bu(t)\\[8pt]
\quad A_{n,k} =&\, -\frac{1}{w}\left\{\begin{array}{l}\sqrt{(2n+1)(2k+1)}, &k < n \\ (-1)^{n-k}\sqrt{(2n+1)(2k+1)}, &k \geq n\end{array}\right.\\[8pt]
B_n =&\, \frac{1}{w}\sqrt{2(2n+1)}
\end{aligned}\tag{24}
\end{equation}\]
我们还可以给每个 \(c_n(T)\) 都引入一个缩放因子,来使得上述结果更一般化。比如我们设 \(c_n(T) = \lambda_n \tilde{c}_n(T)\),代入(23) 整理得
\[\begin{equation}\begin{aligned}
\frac{d}{dt}\tilde{c}_n(T) \approx &\, \frac{\sqrt{2(2n+1)}}{w\lambda_n}u(T) - \frac{\sqrt{2n+1}}{w} \sum\limits_{k=n}^N (-1)^{n-k} \tilde{c}_k(T) \frac{\lambda_k\sqrt{2k+1}}{\lambda_n} \\
&\quad- \frac{\sqrt{2n+1}}{w}\sum_{k=0}^{n-1} \frac{\lambda_k\sqrt{2k+1}}{\lambda_n} \tilde{c}_k(T) \\
\end{aligned}\tag{25}
\end{equation}\]
如果取 \(\lambda_n = \sqrt{2}\),那么 \(A\) 不变,\(B_n = \frac{1}{w}\sqrt{2n+1}\),这就对齐了原论文的结果,如果取 \(\lambda_n = \frac{2}{\sqrt{2n+1}}\),那么就得到了 Legendre Memory Units 中的结果
\[\begin{equation}\begin{aligned}
x'(t) =&\, Ax(t) + Bu(t)\\[8pt]
\quad A_{n,k} =&\, -\frac{1}{w}\left\{\begin{array}{l}2n+1, &k < n \\ (-1)^{n-k}(2n+1), &k \geq n\end{array}\right.\\[8pt]
B_n =&\, \frac{1}{w}(2n+1)
\end{aligned}\tag{26}
\end{equation}\]
这些形式在理论上都是等价的,但可能存在不同的数值稳定性。比如一般来说当 \(u(t)\) 的性态不是特别糟糕时,我们可以预期 \(n\) 越大,\(|c_n|\) 的值就相对越小,这样直接用 \(c_n\) 的话 \(x(t)\) 向量的每个分量的尺度就不大对等,这样的系统在实际计算时容易出现数值稳定问题,而取 \(\lambda_n = \frac{2}{\sqrt{2n+1}}\) 改用 \(\tilde{c}_n\) 的话意味着数值小的分量会被适当放大,可能有助于缓解多尺度问题从而使得数值计算更稳定。
整个区间(LegS)
现在我们继续计算另一个例子:\(t_{\leq T}(s) = (s + 1)T / 2\),它将 \([-1,1]\) 均匀映射到 \([0,T]\),这意味着我们没有舍弃任何历史信息,并且平等地对待所有历史,原论文将这种情形称为“LegS(Scaled Legendre)”。
同样地,通过代入 (14) 得到
\[\frac{d}{dT}c_n(T) = \frac{\sqrt{2(2n+1)}}{T}u(T) - \frac{1}{T}\int_{-1}^1 u((s + 1)T / 2) \left[g_n(s) + (s+1) g_n'(s)\right] ds\]
利用(20) 得到
\[\begin{aligned}
&\,\int_{-1}^1 u((s + 1)T / 2) \left[g_n(s) + (s+1) g_n'(s)\right] ds \\
=&\,c_n(T) + \int_{-1}^1 u((s + 1)T / 2) (s+1) g_n'(s) ds \\
=&\, c_n(T) + \int_{-1}^1 u((s + 1)T / 2)(s+1) \sqrt{\frac{2n+1}{2}} p_n'(s) \\
=&\, c_n(T) + \int_{-1}^1 u((s + 1)T / 2)\sqrt{\frac{2n+1}{2}}\left[-(n+1) p_n(s) + \sum_{k=0}^n (2k + 1) p_k(s)\right] ds \\
=&\, c_n(T) + \int_{-1}^1 u((s + 1)T / 2)\left[-(n+1) g_n(s) + \sum_{k=0}^n \sqrt{(2n+1)(2k + 1)} g_k(s)\right] ds \\
=&\, -n c_n(T) + \sum_{k=0}^n \sqrt{(2n+1)(2k + 1)} c_k(T) \\
\end{aligned}\]
于是有
\[\begin{equation}\begin{aligned}\frac{d}{dT}c_n(T) &=\\ &\frac{\sqrt{2(2n+1)}}{T}u(T) - \frac{1}{T}\left(-n c_n(T) + \sum_{k=0}^n \sqrt{(2n+1)(2k + 1)} c_k(T)\right)\ \ \ \ \ \ \end{aligned} \tag{27}
\end{equation}\]
将 \(T\) 换回 \(t\),将所有的 \(c_n(t)\) 堆在一起记为 \(x(t) = (c_0,c_1,\cdots,c_N)\),那么根据上式可以写出
\[\begin{equation}\begin{aligned}
x'(t) =&\, \frac{A}{t}x(t) + \frac{B}{t}u(t)\\[8pt]
\quad A_{n,k} =&\, -\left\{\begin{array}{l}\sqrt{(2n+1)(2k+1)}, &k < n \\ n+1, &k = n \\
0, &k > n\end{array}\right.\\[8pt]
B_n =&\, \sqrt{2(2n+1)}
\end{aligned}\tag{28}
\end{equation}\]
引入缩放因子来一般化结果也是可行的:设 \(c_n(T) = \lambda_n \tilde{c}_n(T)\),代入(23) 整理得
\[\begin{equation}\begin{aligned}\frac{d}{dT}&\tilde{c}_n(T) \\= &\frac{\sqrt{2(2n+1)}}{T\lambda_n}u(T) - \frac{1}{T}\left(-n \tilde{c}n(T) + \sum{k=0}^n \frac{\sqrt{(2n+1)(2k + 1)}\lambda_k}{\lambda_n} \tilde{c}_k(T)\right)\end{aligned}\ \ \tag{29}
\end{equation}\]
取 \(\lambda_n=\sqrt{2}\) 就可以让 \(A\) 不变,\(B\) 变为 \(B_n = \sqrt{2n+1}\),就对齐了原论文的结果。如果取\(\lambda_n=\sqrt{\frac{2}{2n+1}}\) ,就可以像上一节LegT的结果一样去掉根号
\[\begin{equation}\begin{aligned}
x'(t) =&\, \frac{A}{t}x(t) + \frac{B}{t}u(t)\\[8pt]
\quad A_{n,k} =&\, -\left\{\begin{array}{l}2n+1, &k < n \\ n+1, &k = n \\
0, &k > n\end{array}\right.\\[8pt]
B_n =&\, 2n+1
\end{aligned}\tag{30}
\end{equation}\]
但原论文没有考虑这种情况,原因不详。
延伸思考
回顾Leg-S的整个推导,我们可以发现其中关键一步是将 \((s+1) g_n'(s)\) 拆成\(g_0(s),g_1(s),\cdots,g_n(s)\) 的线性组合,对于正交多项式来说,\((s+1) g_n'(s)\) 是一个 \(n\) 次多项式,所以这种拆分必然可以精确成立,但如果是傅立叶级数的情况,\(g_n(s)\) 是指数函数,此时类似的拆分做不到了,至少不能精确地做到,所以可以说选取正交多项式为基的根本目的是简化后面推导。
特别要指出的是,HiPPO是一个自下而上的框架,它并没有一开始就假设系统必须是线性的,而是从正交基逼近的角度反过来推出其系数的动力学满足一个线性ODE系统,这样一来我们就可以确信,只要认可所做的假设,那么线性ODE系统的能力就是足够的,而不用去担心线性系统的能力限制了你的发挥。
当然,HiPPO对于每一个解所做的假设及其物理含义也很清晰,所以对于重用了HiPPO矩阵的SSM,它怎么储存历史、能储存多少历史,从背后的HiPPO假设就一清二楚。比如LegT就是只保留 \(w\) 大小的最邻近窗口信息,如果你用了LegT的HiPPO矩阵,那么就类似于一个Sliding Window Attention;而LegS理论上可以捕捉全部历史,但这有个分辨率问题,因为 \(x(t)\) 的维度代表了拟合的阶数,它是一个固定值,用同阶的函数基去拟合另一个函数,肯定是区间越小越准确,区间越大误差也越大,这就好比为了一次性看完一幅大图,那么我们必须站得更远,从而看到的细节越少。
诸如RWKV、LRU等模型,并没有重用HiPPO矩阵,而是改为可训练的矩阵,原则上具有更多的可能性来突破瓶颈,但从前面的分析大致上可以感知到,不同矩阵的线性ODE只是函数基不同,但本质上可能都只是有限阶函数基逼近的系数动力学。既然如此,分辨率与记忆长度就依然不可兼得,想要记忆更长的输入并且保持效果不变,那就只能增加整个模型的体量(即相当于增加hidden_size),这大概是所有线性系统的特性。
离散格式
前面我们已经推导出了两类线性ODE系统,分别是:
\[\begin{align}
&\text{HiPPO-LegT:}\quad x'(t) = Ax(t) + Bu(t) \tag{31} \\[5pt]
&\text{HiPPO-LegS:}\quad x'(t) = \frac{A}{t}x(t) + \frac{B}{t}u(t)\tag{32} \end{align}\]
其中 \(A,B\) 是与时间 \(t\) 无关的常数矩阵,HiPPO矩阵主要指矩阵 \(A\)。在这一节中,我们讨论这两个ODE的离散化。
输入转换
在实际场景中,输入的数据点是离散的序列 \(u_0,u_1,u_2,\cdots,u_k,\cdots\),比如流式输入的音频信号、文本向量等,我们希望用如上的ODE系统来实时记忆这些离散点。为此,我们先定义
\[\begin{equation}u(t) = u_k,\quad \text{如果} t\in[k\epsilon, (k + 1)\epsilon)\tag{33}
\end{equation}\]
其中 \(\epsilon\) 就是离散化的步长。该定义也就是说在区间 \([k\epsilon, (k + 1)\epsilon)\) 内,\(u(t)\) 是一个常数函数,其值等于 \(u_k\)。很明显这样定义出来的 \(u(t)\) 无损原本 \(u_k\) 序列的信息,因此记忆 \(u(t)\) 就相当于记忆 \(u_k\) 序列。
从 \(u_k\) 变换到 \(u(t)\),可以使得输入信号重新变回连续区间上的函数,方便后面进行积分等运算,此外在离散化的区间内保持为常数,也能够简化离散化后的格式。
LegT版本
我们先以LegT型ODE (31) 为例,将它两端积分
\[\begin{equation}x(t+\epsilon) - x(t) = A\int_t^{t+\epsilon} x(s)ds + B\int_t^{t+\epsilon}u(s)ds\tag{34}
\end{equation}\]
其中 \(t=k\epsilon\)。根据 \(u(t)\) 的定义,它在 \([t, t + \epsilon)\) 区间内恒为 \(u_k\),于是 \(u(s)\) 的积分可以直接算出来:
\[\begin{equation}x(t+\epsilon) - x(t) = A\int_t^{t+\epsilon} x(s)ds + \epsilon B u_k\tag{35}
\end{equation}\]
接下来的结果,就取决于我们如何近似 \(x(s)\) 的积分了。假如我们认为在 \([t, t + \epsilon)\) 区间内 \(x(s)\) 近似恒等于 \(x(t)\),那么就得到前向欧拉格式
\[x(t+\epsilon) - x(t) = \epsilon A x(t) + \epsilon B u_k \quad\Rightarrow\quad x(t+\epsilon) = (I + \epsilon A)x(t) + \epsilon B u_k\]
我们认为在 \([t, t + \epsilon)\) 区间内 \(x(s)\) 近似恒等于 \(x(t+\epsilon)\),那么就得到后向欧拉格式
\[x(t+\epsilon) - x(t) = \epsilon A x(t+\epsilon) + \epsilon B u_k \quad\Rightarrow\quad x(t+\epsilon) = (I - \epsilon A)^{-1}(x(t) + \epsilon B u_k)\]
前后向欧拉都具有相同的理论精度,但后向通常会有更好的数值稳定性。如果要更准确一些,那么认为在 \([t, t + \epsilon)\) 区间内 \(x(s)\) 近似恒等于 \(\frac{1}{2}[x(t) + x(t+\epsilon)]\),那么得到双线性形式:
\[\begin{equation}\begin{gathered}
x(t+\epsilon) - x(t) = \frac{1}{2}\epsilon A [x(t) + x(t+\epsilon)] + \epsilon B u_k \\
\Downarrow \\
x(t+\epsilon) = (I - \epsilon A/2)^{-1}[(I + \epsilon A/2) x(t) + \epsilon B u_k]
\end{gathered}\tag{36}
\end{equation}\]
这也等价于先用前向欧拉走半步,再用后向欧拉走半步。更一般地,我们还可以认为在 \([t, t + \epsilon)\) 区间内 \(x(s)\) 近似恒等于 \(\alpha x(t) + (1 - \alpha) x(t+\epsilon)\),其中 \(\alpha\in[0,1]\),这就不进一步展开了。事实上,我们也可以完全不做近似,因为结合式 (31) 以及在区间 \([t,t+\epsilon)\) 中 \(u(s)\) 是常数 \(u_k\),我们完全可以用“常数变易法”来精确求解出来,结果是
\[\begin{equation}x(t+\epsilon) = e^{\epsilon A} x(t) + A^{-1} (e^{\epsilon A} - I) B u_k\tag{37}
\end{equation}\]
这里的矩阵指数按照级数来定义,可以参考《恒等式 det(exp(A)) = exp(Tr(A)) 赏析》。
LegS版本
现在轮到LegS型ODE了,它的思路跟LegT型基本一致,结果也大同小异。首先将式(32) 两端积分得到
\[\begin{equation}x(t+\epsilon) - x(t) = A\int_t^{t+\epsilon} \frac{x(s)}{s}ds + B\int_t^{t+\epsilon}\frac{u(s)}{s}ds\tag{38}
\end{equation}\]
根据 \(u(t)\) 定义,第二项积分的 \(u(s)\) 在 \([t,t+\epsilon)\) 恒为 \(u_k\),所以它相当于 \(1/s\) 的积分,可以直接积分出来得 \(\ln\frac{t+\epsilon}{t}\),当然直接换为一阶近似 \(\frac{\epsilon}{t}\) 也无妨,因为本身 \(u_k\) 到 \(u(t)\) 的变换有很大自由度,这点误差无所谓。至于第一项积分,我们直接采用精度更高的中点近似,得到
\[\begin{equation}\begin{gathered}
x(t+\epsilon) - x(t) = \frac{1}{2}\epsilon A\left(\frac{x(t)}{t}+\frac{x(t+\epsilon)}{t+\epsilon}\right) + \frac{\epsilon}{t} B u_k \\[5pt]
\Downarrow \\[5pt]
x(t+\epsilon) = \left(I - \frac{\epsilon A}{2(t+\epsilon)}\right)^{-1}\left[\left(I + \frac{\epsilon A}{2t}\right)x(t) + \frac{\epsilon}{t} B u_k\right]
\end{gathered}\tag{39}
\end{equation}\]
事实上,(32) 也可以精确求解,只需要留意到它等价于
\[\begin{equation}Ax(t) + Bu(t) = t x'(t) = \frac{d}{d\ln t} x(t)\tag{40}
\end{equation}\]
这意味着只需要做变量代换 \(\tau = \ln t\),那么LegS型ODE就可以转化为LegT型ODE:
\[\begin{equation}\frac{d}{d\tau} x(e^{\tau}) = Ax(e^{\tau}) + Bu(e^{\tau})\tag{41}
\end{equation}\]
利用(37)得到(由于变量代换,时间间隔由 \(\epsilon\) 变成 \(\ln(t+\epsilon) - \ln t\))
\[\begin{equation}x(t+\epsilon) = e^{(\ln(t+\epsilon) - \ln t) A} x(t) + A^{-1} \big(e^{(\ln(t+\epsilon) - \ln t) A} - I\big) B u_k\tag{42}
\end{equation}\]
然而,上式虽然是精确解,但不如同为精确解的(37)好用,因为(37)的指数矩阵部分是 \(e^{\epsilon A}\),跟时间\(t\) 无关,所以一次性计算完就可以了。但上式中 \(t\) 在矩阵指数里边,意味着在迭代过程中需要反复计算矩阵指数,对计算并不友好,所以LegS型ODE我们一般只会用 (39) 来离散化。
优良性质
接下来,LegS是我们的重点关注对象。重点关注LegS的原因并不难猜,因为从推导的假设来看,它是目前求解出来的唯一一个能够记忆整个历史的ODE系统,这对于很多场景如多轮对话来说至关重要。此外,它还有其他的一些比较良好且实用的性质。
尺度等变
比如,LegS的离散化格 (39) 是步长无关的,我们只需要将 \(t=k\epsilon\) 代入里边,并记 \(x(k\epsilon)=x_k\),就可以发现
\[\begin{equation}
x_{k+1} = \left(I - \frac{A}{2(k + 1)}\right)^{-1}\left[\left(I + \frac{A}{2k}\right)x_k + \frac{1}{k} B u_k\right]\tag{43}
\end{equation}\]
步长 \(\epsilon\) 被自动地消去了,从而自然地减少了一个需要调的超参数,这对于炼丹人士显然是一个好消息。注意步长无关是LegS型ODE的一个固有性质,它跟具体的离散化方式并无直接关系,比如精确解(42) 同样是步长无关的:
\[\begin{equation}x_{k+1} = e^{(\ln(k+1) - \ln k) A} x_k + A^{-1} \big(e^{(\ln(k+1) - \ln k) A} - I\big) B u_k\tag{44}
\end{equation}\]
其背后的原因,在于LegS型ODE满足“时间尺度等变性(Timescale equivariance)”——如果我们设\(t=\lambda\tau\) 代入LegS型ODE,将得到
\[\begin{equation}Ax(\alpha\tau) + Bu(\alpha\tau) = (\alpha\tau)\times \frac{d}{d(\alpha\tau)} x(\alpha\tau) = \tau \frac{d}{d\tau}x(\alpha\tau)\tag{45}
\end{equation}\]
这意味着,当我们将 \(u(t)\) 换成 \(u(\alpha t)\) 时,LegS的ODE形式并没有变化,而对应的解则是 \(x(t)\) 换成了 \(x(\alpha t)\) 。这个性质的直接后果就是:当我们选择更大的步长时,递归格式不需要发生变化,因为结果 \(x_k\) 的步长也会自动放大,这就是LegS型ODE离散化与步长无关的本质原因。
长尾衰减
LegS型ODE的另一个优良性质是,它关于历史信号的记忆是多项式衰减(Polynomial decay)的,这比常规RNN的指数衰减更缓慢,从而理论上能记忆更长的历史,更不容易梯度消失。为了理解这一点,我们可以从精确解(44) 出发,从(44)可以看到,每递归一步,历史信息的衰减效应可以用矩阵指数 \(e^{(\ln(k+1) - \ln k) A}\) 来描述,那么从第 \(m\) 步递归到第 \(n\) 步,总的衰减效应是
\[\begin{equation}\prod_{k=m}^{n-1} e^{(\ln(k+1) - \ln k) A} = e^{(\ln n - \ln m) A}\tag{46}
\end{equation}\]
回顾HiPPO-LegS中 \(A\) 的形式:
\[\begin{equation}A_{n,k} = -\left\{\begin{array}{l}\sqrt{(2n+1)(2k+1)}, &k < n \\ n+1, &k = n \\
0, &k > n\end{array}\right.\tag{47}
\end{equation}\]
从定义可以看出,\(A\)是一个下三角阵,其对角线元素为 \(-1,-2,-3,\cdots\)。我们知道,三角阵的对角线元素正好是它的特征值(参考Triangular matrix),由此可以看到一个 \(d\times d\) 大小的\(A\)矩阵,有 \(d\) 个不同的特征值 \(-1,-2,\cdots,-d\),这说明 \(A\) 矩阵是可对角化的,即存在可逆矩阵 \(P\),使得 \(A = P^{-1}\Lambda P\),其中 \(\Lambda = \text{diag}(-1,-2,\cdots,-d)\),于是我们有
\[\begin{equation}\begin{aligned}
e^{(\ln n - \ln m) A} =&\, e^{(\ln n - \ln m) P^{-1}\Lambda P} \\
=&\, P^{-1} e^{(\ln n - \ln m) \Lambda}P \\
=&\, P^{-1}\,\text{diag}(e^{-(\ln n - \ln m)},e^{-2(\ln n - \ln m)},\cdots,e^{-d(\ln n - \ln m)})\,P \\
=&\, P^{-1}\,\text{diag}\Big(\frac{m}{n},\frac{m^2}{n^2},\cdots,\frac{m^d}{n^d}\Big)\,P \\
\end{aligned}\tag{48}
\end{equation}\]
可见,最终的衰减函数是 \(1/n\) 的 \(1,2,\cdots,d\) 次函数的线性组合,所以LegS型ODE关于历史记忆至多是多项式衰减的,比指数衰减更加长尾,因此理论上有更好的记忆力。
计算高效
最后,我们指出HiPPO-LegS的A矩阵是计算高效(Computational efficiency)的。具体来说,直接按照矩阵乘法的朴素实现的话,一个 \(d\times d\) 的矩阵乘以 \(d\times 1\) 的列向量,需要做 \(d^2\) 次乘法,但LegS的 \(A\) 矩阵与向量相乘则可以降低到 \(\mathcal{O}(d)\) 次,更进一步地,我们还可以证明离散化后的 (39) 也可以在 \(\mathcal{O}(d)\) 完成。
为了理解这一点,我们首先将HiPPO-LegS的 \(A\) 矩阵等价地改写成
\[\begin{equation}A_{n,k} = \left\{\begin{array}{l}n\delta_{n,k} - \sqrt{2n+1}\sqrt{2k+1}, &k \leq n \\ 0, &k > n\end{array}\right.\tag{49}
\end{equation}\]
对于向量 \(v = [v_0,v_1,\cdots,v_{d-1}]\),我们有
\[\begin{equation}\begin{aligned}
(Av)n = \sum{k=0}^n A_{n,k}v_k =&\, \sum_{k=0}^n \left(n\delta_{n,k} - \sqrt{2n+1}\sqrt{2k+1}\right)v_k \\
=&\, n v_n -\sqrt{2n+1}\sum_{k=0}^n \sqrt{2k+1}v_k
\end{aligned}\tag{50}
\end{equation}\]
这包含三种运算,第一项的 \(n v_n\)是向量 \([0,1,2,\cdots,d-1]\)与 \(v\) 做逐位相乘运算,第二项的\(\sqrt{2k+1}v_k\) 则是向量 \([1,\sqrt{3},\sqrt{5},\cdots,\sqrt{2d-1}]\) 与 \(v\) 做逐位相乘,然后 \(\sum\limits_{k=0}^n\) 就是 \(\text{cumsum}\) 运算,最后乘以 \(\sqrt{2n+1}\) 就是再逐位相乘向量 \([1,\sqrt{3},\sqrt{5},\cdots,\sqrt{2d-1}]\),每一步都可以在\(\mathcal{O}(d)\)内完成,因此总的复杂度是 \(\mathcal{O}(d)\) 的。
我们再来看 (44),它包含两步“矩阵-向量”乘法运算,一是 \((I+\lambda A)v\),\(\lambda\) 是任意实数,刚才我们已经证明了 \(Av\) 是计算高效的,自然 \((I+\lambda A)v\) 也是;二是 \((I-\lambda A)^{-1}v\),接下来我们将证明它也是计算高效的。这只需要留意到求 \(z=(I-\lambda A)^{-1}v\) 等价于解方程 \(v = (I-\lambda A)z\),利用上面给出的 \(Av\) 表达式,我们可以得到
\[\begin{equation}v_n = z_n - \lambda \left(n z_n - \sqrt{2n+1}\sum_{k=0}^n \sqrt{2k+1}z_k\right)\tag{51}
\end{equation}\]
记 \(S_n = \sum\limits_{k=0}^n \sqrt{2k+1}z_k\),那么 \(z_n = \frac{S_n - S_{n-1}}{\sqrt{2n+1}}\),代入上式得
\[\begin{equation}v_n = \frac{S_n - S_{n-1}}{\sqrt{2n+1}} - \lambda \left(n \frac{S_n - S_{n-1}}{\sqrt{2n+1}} - \sqrt{2n+1}S_n\right)\tag{52}
\end{equation}\]
整理得
\[\begin{equation}S_n = \frac{1 - \lambda n}{1+\lambda n + \lambda}S_{n-1} + \frac{\sqrt{2n+1}}{1+\lambda n + \lambda}v_n\tag{53}
\end{equation}\]
这是一个标量的递归式,可以完全串行地计算,也可以利用Prefix Sum的相关算法并行计算,计算复杂度为 \(\mathcal{O}(d)\) 或者 \(\mathcal{O}(d\log d)\),总之相比 \(\mathcal{O}(d^2)\) 都会更加高效。
傅立叶基
最后,我们以傅立叶基的一个推导收尾。在前面的推导中,我们以傅立叶级数来引出了线性系统,但只推导了邻近窗口形式的结果,而后面的勒让德多项式基我们则推导了邻近窗口和完整区间两个版本(即LegT和LegS)。那么傅立叶基究竟能不能推导一个跟LegS相当的版本呢?其中会面临什么困难呢?下面我们对此进行探讨。
同样地,相关铺垫我们不再重复,按照上一节的记号,傅立叶基的系数为
\[\begin{equation}c_n(T) = \int_0^1 u(t_{\leq T}(s)) e^{-2i\pi n s}ds\tag{54}
\end{equation}\]
跟LegS一样,为了记忆整个 \([0,T]\) 区间的信号,我们需要一个 \([0,1]\mapsto [0,T]\) 的映射,为此选取最简单的 \(t_{\leq T}(s)=sT\),代入后两边求导得到
\[\begin{equation}\frac{d}{dT}c_n(T) = \int_0^1 u'(sT) s e^{-2i\pi n s}ds\tag{55}
\end{equation}\]
分部积分得到
\[\begin{equation}\begin{aligned}
\frac{d}{dT}c_n(T) =&\, \frac{1}{T}\int_0^1 s e^{-2i\pi n s}d u(sT) \\
=&\, \frac{1}{T} u(sT) s e^{-2i\pi n s}\big|_{s=0}^{s=1} - \frac{1}{T}\int_0^1 u(sT) d(s e^{-2i\pi n s})\\
=&\, \frac{1}{T} u(T) - \frac{1}{T}\int_0^1 u(sT) e^{-2i\pi n s} ds + \frac{2i\pi n}{T}\int_0^1 u(sT) s e^{-2i\pi n s} ds\\
=&\, \frac{1}{T} u(T) - \frac{1}{T}c_n(T) + \frac{2i\pi n}{T}\int_0^1 u(sT) s e^{-2i\pi n s} ds\\
\end{aligned}\tag{56}
\end{equation}\]
前面我们提到,HiPPO选取勒让德多项式为基的重要原因之一是 \((s+1)p_n'(t)\) 可以分解为\(p_0(t),p_1(t),\cdots,p_n(t)\) 的线性组合,而傅里叶基的 \(s e^{-2i\pi n s}\) 则不能做到这一点。但事实上,如果允许误差的话,这个断论是不成立的,因为我们同样可以将 \(s\) 分解为傅里叶级数:
\[\begin{equation}s = \frac{1}{2} + \frac{i}{2\pi}\sum_{k\neq 0} \frac{1}{k} e^{2i\pi k s}\tag{57}
\end{equation}\]
这里的求和有无限项,如果要截断为有限项的话,就会产生误差,但我们可以先不纠结这一点,直接往上代入得到
\[\begin{equation}\begin{aligned}
&\,\frac{2i\pi n}{T}\int_0^1 u(sT) s e^{-2i\pi n s} ds \\
=&\, \frac{2i\pi n}{T}\int_0^1 u(sT) \left(\frac{1}{2} + \frac{i}{2\pi}\sum_{k\neq 0} \frac{1}{k} e^{2i\pi k s}\right) e^{-2i\pi n s} ds \\
=&\, \frac{i\pi n}{T}\int_0^1 u(sT) e^{-2i\pi n s} ds - \frac{1}{T}\sum_{k\neq 0} \frac{n}{k}\int_0^1 u(sT) e^{-2i\pi (n - k) s} ds \\
=&\, \frac{i\pi n}{T}c_n(T) - \frac{1}{T}\sum_{k\neq 0} \frac{n}{k}c_{n-k}(T) \\
=&\, \frac{i\pi n}{T}c_n(T) - \frac{1}{T}\sum_{k\neq n} \frac{n}{n - k}c_k(T) \\
\end{aligned}\tag{58}
\end{equation}\]
这样一来
\[\begin{equation}
\frac{d}{dT}c_n(T) = \frac{1}{T} u(T) + \frac{i\pi n - 1}{T}c_n(T) - \frac{1}{T}\sum_{k\neq n} \frac{n}{n - k}c_k(T)\tag{59}
\end{equation}\]
所以可以写出
\[\begin{equation}\begin{aligned}
x'(t) =&\, \frac{A}{t}x(t) + \frac{B}{t}u(t)\\[8pt]
\quad A_{n,k} =&\, \left\{\begin{array}{l}-\frac{n}{n-k}, &k \neq n \\ i\pi n - 1, &k = n\end{array}\right.\\[8pt]
B_n =&\, 1
\end{aligned}\tag{60}
\end{equation}\]
实际使用的时候,我们只需要截断 \(|n|,|k|\leq N\),就可以得到一个 \((2N+1)\times (2N+1)\)的矩阵。截断带来的误差其实是无所谓的,因为我们在推导HiPPO-LegT的时候同样引入了有限级数近似,那会我们同样也没考虑误差,或者反过来讲,对于特定的任务,我们会选择适当的规模(即 \(N\) 的大小),而这个“适当”的含义之一,就是截断带来误差对于该任务是可以忽略的。
对大多数人来说,傅立叶基的这个推导可能还更容易理解一些,因为勒让德多项式对很多读者来说都比较陌生,尤其是LegT、LegS推导过程中用到的几个恒等式,而对于傅立叶级数大多数读者应该或多或少都有所了解。不过,从结果上来看,傅立叶基的这个结果可能不如LegS实用,一来它引入了复数,这增加了实现的复杂度,二来它推导出的 \(A\) 矩阵不像LegS那样是个相对较淡的下三角阵,因此理论分析起来也更为复杂。所以,大家权当它是一道深化对HiPPO的理解的练习题就好。
文章小结
在这篇文章中,我们介绍了《HiPPO: Recurrent Memory with Optimal Polynomial Projections》(简称HiPPO)的主要推导。HiPPO通过适当的记忆假设,自下而上地导出了线性ODE系统,并且针对勒让德多项式的情形求出了相应的解析解(HiPPO矩阵),其结果被后来诸多SSM(State Space Model)使用,可谓是SSM的重要奠基之作。另外我们还补充介绍了如何对ODE进行离散化、LegS型ODE的一些优良性质,以及利用傅立叶基记忆整个历史区间的结果推导(即LegS的傅立叶版本),以求获得对HiPPO的更全面理解。
Reference
重温SSM(一):线性系统和HiPPO矩阵 - 科学空间|Scientific Spaces
重温SSM(二):HiPPO的一些遗留问题 - 科学空间|Scientific Spaces
HiPPO: Recurrent Memory with Optimal Polynomial Projections