知差（chai）

作者： admin

探索线性Attention的局限性：从“集中注意力”角度出发
近年来，Transformer架构在自然语言处理领域取得了显著的成果，而Attention机制则是其核心所在。然而，随着研究的深入，传统的标准Attention机制暴露出了一些计算复杂度和资源需求上的问题，这促使研究者们开始探索更高效的线性Attention。然而，线性Attention在实际效果上却一直不如标准Attention。本文将从“集中注意力”的角度，探讨线性Attention的局限性，并尝试给出一个合理的解释。

Attention机制的稀疏性

什么是稀疏性？

在《从熵不变性看Attention的Scale操作》一文中，我们用信息熵来度量Attention的“集中注意力”程度。熵越低，Attention越有可能集中在某个token上。然而，信息熵只能用于归一化且非负的Attention矩阵，这限制了其适用范围。因此，我们引入了另一种稀疏性度量指标：S(x) = E[|x|] / sqrt(E[x^2])。这个指标与信息熵类似，S(x)越小，表示对应的随机矢量越稀疏，即越有可能“一家独大”。

标准Attention的稀疏性

对于标准Attention机制，f = exp，我们可以推导出其稀疏性：
```
S(a) = exp(-1/2 * σ^2 * ||q||^2)
```
当σ或||q||趋向无穷大时，S(a)趋向于0，这意味着标准Attention可以任意稀疏地“集中注意力”。

GAU的稀疏性

对于Gated Attention Unit (GAU)机制，f = relu2，其稀疏性较为复杂，但通过计算可以发现，只有当偏置项β小于0时，稀疏性才有机会趋于0。这表明，GAU在某些条件下也能实现较高的稀疏性。

线性Attention的局限性

极简线性Attention

对于最简单的线性Attention，即不加任何激活函数f = identical，其稀疏性为：
```
S(a) = sqrt(2/π * γ * exp(-β^2/(2γ^2)) + β * erf(β/(2√γ)) / (β^2 + γ^2))
```
从图像可以看出，极简线性Attention的稀疏性存在一个较高的下限，这意味着它难以“集中注意力”到关键位置上。

一般线性Attention

线性Attention的一般形式为a_j ∝ g(q) ⋅ h(k_j)，其稀疏性为：
```
S(a) = 1 / sqrt(1 + (σ~ * μ~ * ||q~||_2 / ||q~||_1)^2)
```
这表明，要想线性Attention变得稀疏，可以通过降低k~串行的信噪比或增大q的模长。然而，非负型线性Attention通常只能表示绝对位置的重要性，难以表达相对位置的重要性。

线性衰减Attention

对于带有显式递归的线性RNN模型，其稀疏性为：
```
S(a) = 1 - λ^n / n(1 - λ) * sqrt(1 + λ / (1 + λ^n))
```
当λ < 1时，随着n趋向无穷大，S(a)趋向0。这意味着这种模型可以实现较高的稀疏性，但其注意力仅能表达固定不变的注意力衰减，难以自适应地关注到长距离的上下文。

结论

本文通过Attention矩阵的稀疏程度，考察了不同Attention机制的潜力，得出以下结论：
1. 标准Attention可以实现任意稀疏的注意力矩阵。
2. 线性Attention难以实现高稀疏性，尤其是在表示相对位置的重要性时。
3. 带有显式衰减的线性Attention可以实现稀疏性，但其注意力固定，难以自适应。
这些发现或许能够解释线性Attention在实际效果上略逊一筹的原因。线性Attention在“集中注意力”方面存在固有的局限性，这使得它在处理复杂上下文时表现不如标准Attention。未来的研究或许需要在如何提高线性Attention的稀疏性和灵活性上继续努力，以期实现更高效且性能优越的Transformer模型。

参考文献
1. 《Transformer升级之路：3、从Performer到线性Attention》
2. 《为什么现在的LLM都是Decoder-only的架构？》
3. 《从熵不变性看Attention的Scale操作》
4. 《FLASH：可能是近来最有意思的高效Transformer设计》
5. 《相对位置编码Transformer的一个理论缺陷与对策》
6. 《如何度量数据的稀疏程度？》
7. 《线性Attention的探索：Attention必须有个Softmax吗？》
8. 《Google新作试图“复活”RNN：RNN能否再次辉煌？》
通过上述分析，我们不仅理解了不同Attention机制的稀疏性差异，还揭示了线性Attention在实际应用中的局限性。希望本文的讨论能够为未来的研究提供一些新的思路和方向。
2024 年 5 月 29 日
深度学习中的状态空间模型（SSM）初探
引言

前几天，笔者看了几篇介绍SSM（State Space Model）的文章，才发现原来自己从未认真了解过SSM，于是打算认真去学习一下SSM的相关内容，顺便开了这个新坑，记录一下学习所得。

SSM的概念由来已久，但这里我们特指深度学习中的SSM，一般认为其开篇之作是2021年的S4，不算太老，而SSM最新最火的变体大概是去年的Mamba。当然，当我们谈到SSM时，也可能泛指一切线性RNN模型，这样RWKV、RetNet还有此前我们在《Google新作试图“复活”RNN：RNN能否再次辉煌？》介绍过的LRU都可以归入此类。不少SSM变体致力于成为Transformer的竞争者，尽管笔者并不认为有完全替代的可能性，但SSM本身优雅的数学性质也值得学习一番。

尽管我们说SSM起源于S4，但在S4之前，SSM有一篇非常强大的奠基之作《HiPPO: Recurrent Memory with Optimal Polynomial Projections》（简称HiPPO），所以本文从HiPPO开始说起。

基本形式

先插句题外话，上面提到的SSM代表作HiPPO、S4、Mamba的一作都是Albert Gu，他还有很多篇SSM相关的作品，毫不夸张地说，这些工作筑起了SSM大厦的基础。不论SSM前景如何，这种坚持不懈地钻研同一个课题的精神都值得我们由衷地敬佩。

言归正传。对于事先已经对SSM有所了解的读者，想必知道SSM建模所用的是线性ODE系统：

$[ x′(t) = Ax(t) + Bu(t) ]$
$[ y(t) = Cx(t) + Du(t) ]$

其中 $( 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）级数，它的系数计算公式为：

$[ c_n = \int_0^1 u(t) e^{-2i\pi nt} dt ]$

这时候取一个足够大的整数 $( N )$ ，只保留 $( |n| \leq N )$ 的系数，那么就将 $( u(t) )$ 压缩为一个 ( 2N+1 ) 维的矢量了。

接下来，问题难度就要升级了。刚才我们说 $( t \in [0, 1] )$ ，这是一个静态的区间，而实际中 $( u(t) )$ 代表的是持续采集的信号，所以它是不断有新数据进入的，比如现在我们近似了 $( [0, 1] )$ 区间的数据，马上就有 $( [1, 2] )$ 的数据进来，你需要更新逼近结果来试图记忆整个 ( [0, 2] ) 区间，接下来是 ( [0, 3] )、( [0, 4] ) 等等，这我们称为“在线函数逼近”。而上面的傅里叶系数公式只适用于区间 ( [0, 1] )，因此需要将它进行推广。

为此，我们设 $( t \in [0, T] )，( s \mapsto t \leq T(s) )$ 是 ( [0, 1] ) 到 ( [0, T] ) 的一个映射，那么 $( u(t \leq T(s)) )$ 作为 ( s ) 的函数时，它的定义区间就是 $( [0, 1] )$ ，于是就可以复用傅里叶系数公式：

$[ c_n(T) = \int_0^1 u(t \leq T(s)) e^{-2i\pi ns} ds ]$

这里我们已经给系数加了标记 $( (T) )$ ，以表明此时的系数会随着 $( T )$ 的变化而变化。

线性初现

能将 $( [0, 1] )$ 映射到 $( [0, T] )$ 的函数有无穷多，而最终结果也因 $( t \leq T(s) )$ 而异，一些比较直观且相对简单的选择如下：
1. $( t \leq T(s) = sT )$ ，即将 ( [0, 1] ) 均匀地映射到 ( [0, T] )；
2. 注意 $( t \leq T(s) )$ 并不必须是满射，所以像 $( t \leq T(s) = s + T - 1 )$ 也是允许的，这意味着只保留了最邻近窗口 ( [T – 1, T] ) 的信息，丢掉了更早的部分，更一般地有 $( t \leq T(s) = sw + T - w )$ ，其中 ( w ) 是一个常数，这意味着 ( T – w ) 前的信息被丢掉了；
3. 也可以选择非均匀映射，比如 $( t \leq T(s) = T\sqrt{s} )$ ，它同样是 ( [0, 1] ) 到 ( [0, T] ) 的满射，但 ( s = 1/4 ) 时就映射到 ( T/2 ) 了，这意味着我们虽然关注全局的历史，但同时更侧重于 ( T ) 时刻附近的信息。
现在我们以 $( t \leq T(s) = (s + 1)w/2 + T - w )$ 为例，代入傅里叶系数公式得到：

$[ c_n(T) = \int_0^1 u(sw + T - w) e^{-2i\pi ns} ds ]$

现在我们两边求关于 ( T ) 的导数：

$[ \frac{d}{dT} c_n(T) = \frac{1}{w} \left[ u(T) - u(T - w) \right] + \frac{2i\pi n}{w} c_n(T) ]$

其中我们用了分部积分公式。由于我们只保留了 $( |n| \leq N )$ 的系数，所以根据傅里叶级数的公式，可以认为如下是 $( u(sw + T - w) )$ 的一个良好近似：

$[ u(sw + T - w) \approx \sum_{k=-N}^N c_k(T) e^{2i\pi ks} ]$

那么 $( u(T - w) = u(sw + T - w) \big|{s=0} \approx \sum{k=-N}^N c_k(T) )$ ，代入上式得：

$[ \frac{d}{dT} c_n(T) \approx \frac{1}{w} \left[ u(T) - \sum_{k=-N}^N c_k(T) \right] + \frac{2i\pi n}{w} c_n(T) ]$

将 ( T ) 换成 ( t )，然后所有的 $( c_n(t) )$ 堆在一起记为 $( x(t) = (c_{-N}, c_{-(N-1)}, \ldots, c_0, \ldots, c_{N-1}, c_N) )$ ，并且不区分 $( \approx )$ 和 $( = )$ ，那么就可以写出：

$[ x′(t) = A x(t) + B u(t) ]$

其中：

$[ A_{n,k} = \begin{cases} \frac{2i\pi n - 1}{w}, & k = n \ -\frac{1}{w}, & k \ne n \end{cases}, \quad B_n = \frac{1}{w} ]$

这就出现了如上所示的线性ODE系统。即当我们试图用傅里叶级数去记忆一个实时函数的最邻近窗口内的状态时，结果自然而然地导致了一个线性ODE系统。

一般框架

当然，目前只是选择了一个特殊的 $( t \leq T(s) )$ ，换一个 $( t \leq T(s) )$ 就不一定有这么简单的结果了。此外，傅里叶级数的结论是在复数范围内的，进一步实数化也可以，但形式会变得复杂起来。所以，我们要将这一过程推广成一个一般化的框架，从而得到更一般、更简单的纯实数结论。

设 $( t \in [a, b] )$ ，并且有目标函数 $( u(t) )$ 和函数基 $( { g_n(t) }_{n=0}^N )$ ，我们希望有后者的线性组合来逼近前者，目标是最小化 $( L^2 )$ 距离：

$[ \arg\min_{c_1, \ldots, c_N} \int_a^b \left[ u(t) - \sum_{n=0}^N c_n g_n(t) \right]^2 dt ]$

这里我们主要在实数范围内考虑，所以方括号直接平方就行，不用取模。更一般化的目标函数还可以再加个权重函数 $( \rho(t) )$ ，但我们这里就不考虑了，毕竟HiPPO的主要结论其实也没考虑这个权重函数。

对目标函数展开，得到：

$[ \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 ]$

这里我们只考虑标准正交函数基，其定义为 $( \int_a^b g_m(t) g_n(t) dt = \delta_{m,n} )，( \delta_{m,n} )$ 是克罗内克δ函数，此时上式可以简化成：

$[ \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 ]$

这只是一个关于 $( c_n )$ 的二次函数，它的最小值是有解析解的：

$[ c^*_n = \int_a^b u(t) g_n(t) dt ]$

这也被称为 $( 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) )$ ，然后计算系数：

$[ c_n(T) = \int_a^b u(t \leq T(s)) g_n(s) ds ]$

同样是两边求 $( T )$ 的导数，然后用分部积分法：

$[ \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} / \frac{\partial t \leq T(s)}{\partial s} \right) g_n(s) du(t \leq T(s)) ]$

$[ = \left( \frac{\partial t \leq T(s)}{\partial T} / \frac{\partial t \leq T(s)}{\partial s} \right) g_n(s) \bigg|_{s=b}^{s=a} - \int_a^b u(t \leq T(s)) d \left[ \left( \frac{\partial t \leq T(s)}{\partial T} / \frac{\partial t \leq T(s)}{\partial s} \right) g_n(s) \right] ]$

请勒让德

接下来的计算，就依赖于 $( g_n(t) )$ 和 $( t \leq T(s) )$ 的具体形式了。HiPPO的全称是High-order Polynomial Projection Operators，第一个P正是多项式（Polynomial）的首字母，所以HiPPO的关键是选取多项式为基。现在我们请出继傅里叶之后又一位大牛——勒让德（Legendre），接下来我们要选取的函数基正是以他命名的“勒让德多项式”。

勒让德多项式 $( p_n(t) )$ 是关于 $( t )$ 的 $( n )$ 次函数，定义域为 $( [-1, 1] )$ ，满足：

$[ \int_{-1}^1 p_m(t) p_n(t) dt = \frac{2}{2n+1} \delta_{m,n} ]$

所以 $( p_n(t) )$ 之间只是正交，还不是标准（平分积分为1）， $( g_n(t) = \sqrt{\frac{2n+1}{2}} p_n(t) )$ 才是标准正交基。

当我们对函数基 $( {1, t, t^2, \ldots, t^n} )$ 执行施密特正交化时，其结果正是勒让德多项式。相比傅里叶基，勒让德多项式的好处是它是纯粹定义在实数空间中的，并且多项式的形式能够有助于简化部分 $( t \leq T(s) )$ 的推导过程，这一点我们后面就可以看到。勒让德多项式有很多不同的定义和性质，这里我们不一一展开，有兴趣的读者自行看维基百科介绍即可。

接下来我们用到两个递归公式来推导一个恒等公式，这两个递归公式是：

$[ p′_{n+1}(t) - p′_{n-1}(t) = (2n+1) p_n(t) ]$
$[ p′_{n+1}(t) = (n+1) p_n(t) + t p′_n(t) ]$

由第一个公式迭代得到：

$[ p′_{n+1}(t) = (2n+1) p_n(t) + (2n−3) p_{n−2}(t) + (2n-7) p_{n-4}(t) + ⋯ = \sum_{k=0}^n (2k+1) χ_{n−k} p_k(t) ]$

其中当 $( k )$ 是偶数时 $( χ_k = 1 )$ 否则 $( χ_k = 0 )$ 。代入第二个公式得到：

$[ t p′_n(t) = n p_n(t) + (2n−3) p{n-2}(t) + (2n-7) p_{n-4}(t) + ⋯ ]$

继而有：

$[ (t+1) p′_n(t) = n p_n(t) + (2n-1) p_{n−1}(t) + (2n-3) p_{n-2}(t) + ⋯ = - (n+1) p_n(t) + \sum_{k=0}^n (2k+1) p_k(t) ]$

这些就是等会要用到的恒等式。此外，勒让德多项式满足 ( p_n(1) = 1, p_n(-1) = (-1)^n )，这个边界值后面也会用到。

正如 ( n ) 维空间中不止有一组正交基也一样，正交多项式也不止有勒让德多项式一种，比如还有切比雪夫（Chebyshev）多项式，如果算上加权的目标函数（即 $( ρ(t) ≢ 1 )$ ），还有拉盖尔多项式等，这些在原论文中都有提及，但HiPPO的主要结论还是基于勒让德多项式展开的，所以剩余部分这里也不展开讨论了。

邻近窗口

完成准备工作后，我们就可以代入具体的 $( t \le T(s) )$ 进行计算了，计算过程跟傅里叶级数的例子大同小异，只不过基函数换成了勒让德多项式构造的标准正交基 $( g_n(t) = \sqrt{\frac{2n+1}{2}} p_n(t) )$ 。作为第一个例子，我们同样先考虑只保留最邻近窗口的信息，此时 $( t \le T(s) = \frac{(s+1)w}{2} + T - w )$ 将 $( [−1, 1] )$ 映射到 $( [T−w, T] )$ ，原论文将这种情形称为“LegT（Translated Legendre）”。

直接代入之前得到的公式，马上得到：

$[ \frac{d}{dT} c_n(T) = \sqrt{\frac{2(2n+1)}{w}} [u(T) - (−1)^n u(T−w)] - \frac{2}{w} \int_{-1}^1 u\left(\frac{(s+1)w}{2} + T - w\right) g′_n(s) ds ]$

我们首先处理 $( u(T−w) )$ 项，跟傅里叶级数那里同样的思路，我们截断 ( n ≤ N ) 作为 $( u\left(\frac{(s+1)w}{2} + T - w\right) )$ 的一个近似：

$[ u\left(\frac{(s+1)w}{2} + T - w\right) ≈ \sum_{k=0}^N c_k(T) g_k(s) ]$

从而有 $( u(T−w) ≈ \sum_{k=0}^N c_k(T) g_k(−1) = \sum_{k=0}^N (−1)^k c_k(T) \sqrt{\frac{2k+1}{2}} )$ 。接着，利用之前的递归公式得到：

$[ \int_{-1}^1 u\left(\frac{(s+1)w}{2} + T - w\right) g′n(s) ds = \int{-1}^1 u\left(\frac{(s+1)w}{2} + T - w\right) \sqrt{\frac{2n+1}{2}} p′_n(s) ds ]$

$[ = \int_{-1}^1 u\left(\frac{(s+1)w}{2} + T - w\right) \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\left(\frac{(s+1)w}{2} + T - w\right) \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{\frac{2n+1}{2}} \sum_{k=0}^{n-1} \sqrt{2(2k+1)} \chi_{n-1-k} c_k(T) ]$

将这些结果集成起来，就有：

$[ \frac{d}{dT} c_n(T) ≈ \sqrt{\frac{2(2n+1)}{w}} u(T) - \sqrt{\frac{2(2n+1)}{w}} (−1)^n \sum_{k=0}^N (−1)^k c_k(T) \sqrt{\frac{2k+1}{2}} - \frac{2}{w} \sqrt{\frac{2n+1}{2}} \sum_{k=0}^{n-1} \sqrt{2(2k+1)} \chi_{n-1-k} c_k(T) ]$

再次地，将 $( T )$ 换回 $( t )$ ，并将所有的 $( c_n(t) )$ 堆在一起记为 $( x(t) = (c_0, c_1, ⋯, c_N) )$ ，那么根据上式可以写出：

$[ x′(t) = A x(t) + B u(t) ]$

其中：

$[ A_{n,k} = \begin{cases} -\frac{2}{w} \sqrt{\frac{(2n+1)(2k+1)}{2}}, & k < n \ -\frac{2}{w} \sqrt{\frac{(2n+1)(2k+1)}{2}} (−1)^{n−k}, & k \ge n \end{cases}, \quad B_n = \sqrt{\frac{2(2n+1)}{w}} ]$

我们还可以给每个 $( c_n(T) )$ 都引入一个缩放因子，来使得上述结果更一般化。比如我们设 $( c_n(T) = λ_n \tilde{c}_n(T) )$ ，代入上式整理得：

$[ \frac{d}{dT} \tilde{c}n(T) ≈ \sqrt{\frac{2(2n+1)}{w}} \frac{u(T)}{λ_n} - \sqrt{\frac{2(2n+1)}{w}} \frac{(−1)^n}{λ_n} \sum{k=0}^N \frac{(−1)^k c_k(T)}{λk} \sqrt{\frac{2k+1}{2}} - \frac{2}{w} \sqrt{\frac{2n+1}{2}} \sum{k=0}^{n-1} \frac{λk}{λ_n} \sqrt{2(2k+1)} \chi{n-1-k} \tilde{c}_k(T) ]$

如果取 $( λ_n = \sqrt{2} )$ ，那么 ( A ) 不变， $( B_n = \sqrt{2(2n+1)} )$ ，这就对齐了原论文的结果。如果取 $( λ_n = \sqrt{\frac{2}{2n+1}} )$ ，那么就得到了Legendre Memory Units中的结果：

$[ x′(t) = A x(t) + B u(t) ]$

其中：

$[ A_{n,k} = \begin{cases} 2n+1, & k < n \ (−1)^{n−k} (2n+1), & k \ge n \end{cases}, \quad B_n = 2n+1 ]$

这些形式在理论上都是等价的，但可能存在不同的数值稳定性。比如一般来说当 $( u(t) )$ 的性态不是特别糟糕时，我们可以预期 $( n )$ 越大， $( |c_n| )$ 的值就相对越小，这样直接用 $( c_n )$ 的话 $( x(t) )$ 矢量的每个分量的尺度就不大对等，这样的系统在实际计算时容易出现数值稳定问题，而取 $( λ_n = \sqrt{\frac{2}{2n+1}} )$ 改用 $( \tilde{c}_n )$ 的话意味着数值小的分量会被适当放大，可能有助于缓解多尺度问题从而使得数值计算更稳定。

整个区间

现在我们继续计算另一个例子： $( t \le T(s) = \frac{(s+1)T}{2} )$ ，它将 $( [−1, 1] )$ 均匀映射到 $( [0, T] )$ ，这意味着我们没有舍弃任何历史信息，并且平等地对待所有历史，原论文将这种情形称为“LegS（Scaled Legendre）”。

同样地，通过代入之前得到的公式：

$[ \frac{d}{dT} c_n(T) = \sqrt{\frac{2(2n+1)}{T}} u(T) - \frac{1}{T} \int_{-1}^1 u\left(\frac{(s+1)T}{2}\right) (s+1) g′_n(s) ds ]$

利用之前的递归公式得到：

$[ \int_{-1}^1 u\left(\frac{(s+1)T}{2}\right) (s+1) g′n(s) ds = \int{-1}^1 u\left(\frac{(s+1)T}{2}\right) \left[g_n(s) + (s+1) g′_n(s)\right] ds ]$

$[ = c_n(T) + \int_{-1}^1 u\left(\frac{(s+1)T}{2}\right) \sqrt{\frac{2n+1}{2}} p′_n(s) ds ]$

$[ = c_n(T) + \int_{-1}^1 u\left(\frac{(s+1)T}{2}\right) \left[-(n+1) g_n(s) + \sum_{k=0}^n \sqrt{(2n+1)(2k+1)} g_k(s)\right] ds ]$

$[ = c_n(T) - n c_n(T) + \sum_{k=0}^n \sqrt{(2n+1)(2k+1)} c_k(T) ]$

于是有：

$[ \frac{d}{dT} c_n(T) = \sqrt{\frac{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) ]$

将 ( T ) 换回 ( t )，将所有的 ( c_n(t) ) 堆在一起记为 $( x(t) = (c_0, c_1, ⋯, c_N) )$ ，那么根据上式可以写出：

$[ x′(t) = A x(t) + B u(t) ]$

其中：

$[ A_{n,k} = \begin{cases} \sqrt{(2n+1)(2k+1)}, & k < n \ n+1, & k = n \ 0, & k > n \end{cases}, \quad B_n = \sqrt{2(2n+1)} ]$

引入缩放因子来一般化结果也是可行的：设 $( c_n(T) = λ_n \tilde{c}_n(T) )$ ，代入上式整理得：

$[ \frac{d}{dT} \tilde{c}n(T) = \sqrt{\frac{2(2n+1)}{T}} \frac{u(T)}{λ_n} - \frac{1}{T} \left(-n \tilde{c}_n(T) + \sum{k=0}^n \sqrt{(2n+1)(2k+1)} \frac{λ_k}{λ_n} \tilde{c}_k(T)\right) ]$

取 $( λ_n = \sqrt{\frac{2}{2n+1}} )$ ，就可以让 ( A ) 不变， $( B_n = \sqrt{2(2n+1)} )$ ，就对齐了原论文的结果。如果取 $( λ_n = \sqrt{\frac{2}{2n+1}} )$ ，就可以像上一节LegT的结果一样去掉根号：

$[ x′(t) = A x(t) + B u(t) ]$

其中：

$[ A_{n,k} = \begin{cases} 2(2n+1), & k < n \ n+1, & k = n \ 0, & k > n \end{cases}, \quad B_n = 2(2n+1) ]$

但原论文没有考虑这种情况，原因不详。

延伸思考

回顾Leg-S的整个推导，我们可以发现其中关键一步是将 $( (s+1) g′_n(s) )$ 拆成 $( g_0(s), g_1(s), ⋯, 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），这大概是所有线性系统的特性。

文章小结

本文以尽可能简单的方式重复了《HiPPO: Recurrent Memory with Optimal Polynomial Projections》（简称HiPPO）的主要推导。HiPPO通过适当的记忆假设，自下而上地导出了线性ODE系统，并且针对勒让德多项式的情形求出了相应的解析解（HiPPO矩阵），其结果被后来诸多SSM（State Space Model）使用，可谓是SSM的重要奠基之作。

HiPPO框架展现了优雅的数学结构和强大的应用潜力，在处理时间序列数据时提供了一种高效的记忆机制。未来的研究可以进一步探索其在不同领域中的应用和改进。

参考文献： https://spaces.ac.cn/archives/10114
2024 年 5 月 28 日

人生梦想 - 关注前沿的计算机技术 acejoy.com 🐾 步子哥の博客 🐾 背多分论坛 🐾 借一步网沪ICP备2024052574号-1