如何解决大模型长距离依赖问题？HiPPO 技术深度解析

前言

HiPPO 是包括 S4、H3 等一系列 state space model (SSM) 相关模型的理论基石，提出了全新架构，旨在解决序列建模中的长距离依赖问题（long-term dependencies）。长距离依赖建模的核心问题在于如何用有限空间记录累计历史数据的信息，并随输入在线更新。当前主流模型大多有各种各样的问题，包括：

• 记忆范围有限，有 vanishing gradient 等问题
• 需要先验，并一定程度上受限于此
• 缺乏处理长距离依赖关系的理论保证

而这篇论文经数学推导，证明了HiPPO理论上可较好地处理长距离依赖关系，并提供了初始化参数，对后续研究产生了深远影响。

HiPPO: Recurrent Memory with Optimal Polynomial Projections

论文地址：https://arxiv.org/abs/2008.07669

正文

1.基本思路 —— 在线函数近似

我们希望通过某种方式将历史信息存储在一个向量中，借此处理长距离依赖关系。为便于分析，将所有历史输入看作一个关于时间的连续函数 $$f(t)$$，最后再对其进行离散化处理。一个自然的思路是用一个性质更好的函数 $$g(t)$$ 对其进行近似，随输入不断更新，达到预测未来的效果。而为将函数转换更容易处理的形式，希望将 $$g(t)$$ 用一个向量 $$\mathbf{c}=(c_i{)}_i$$ 表示，这就引出了取一组空间上的基，用系数表示函数的想法。我们还希望 $$(c_i{)}_i$$ 满足某种递推关系，从而可以只考虑当前输入和向量的历史累计值。

以下将对这一过程进行更详细的说明。

1.1 目的

给定输入函数 $$f(t)\in \mathbb{R},(t\ge 0)$$，对给定的 $$t_0$$，记此时的历史为 $$f_{\le t_0}:=f(x){|}_{x\le t_0}$$，它是否能用固定长度的向量 $$\mathbf{c}(t)=(c_i(t_0){)}_i$$ 近似地表示，并将其写成一个常微分方程的形式（离散条件下是递推式）？

1.2 如何判断近似的质量

评判近似质量需要定义一个在函数空间上的距离。我们选择一个概率测度 $$\mu (t)$$ 代表每个时间点在近似时所占的权重（也即每个时间点的影响力有多大），而任何测度都诱导一个内积为 $${⟨f,g⟩}_{\mu }={\int }_0^{+\infty }f(x)g(x)d\mu (x)$$，平方可积的函数空间（即希尔伯特空间，记为 $${\mathcal{H}}_{\mu }$$），对应范数为 $${\Vert f\Vert }_{L_2(\mu )}={⟨f,f⟩}_{\mu }^{1/2}$$

我们最终想要一个次数有限的结果，因此我们取一个 $${\mathcal{H}}_{\mu }$$ 的 $$N$$ 维子空间 $$\mathcal{G}$$，$$N$$ 代表了近似程度。在以下的论证中，均选择多项式作为自然基，$$\mathcal{G}$$ 为次数低于 $$N$$ 的多项式空间。

实际上并不要求如此选取，此处是出于多项式受到广泛研究，有更多现成结论的缘故。

由于对于时刻 $$t$$，我们仅关心 $$f_{\le t}$$ 的影响，令选择的测度随时间变化，记作 $${\mu }^{(t)}$$，在 $$(-\infty ,t]$$ 上度量。此时，近似程度便可以用距离 $${\Vert f_{\le t}-g^{(t)}\Vert }_{L_2({\mu }^{(t)})}$$ 来表示。我们希望找到 $$g^{(t)}\in \mathcal{G}$$ 使其取得最小值。

1.3 如何找到近似函数

近似函数可用其在 $$\mathcal{G}$$ 的一组基下的展开中的 $$N$$ 个系数来表示，关键在于如何选取适当的基。一个自然的选择是正交基 KaTeX parse error: Expected '}', got 'EOF' at end of input: …_n\}_{n，它可以扩展为 \mathcal{H} _\mu$$ 的一组正交基。

一个希尔伯特空间的正交基 { q n } \{q_n\} {qn​} 也即满足

$${⟨q_i,q_j⟩}_{{\mu }^{(t)}}=\{\begin{array}{ll}1,& i=j\\ 0,& i\ne j\end{array}$$

的一组基（可用 Gram-Schmidt 正交化方法将非正交基转换为正交基得到）。特别地，对于 $$\mathcal{G}$$，一般 $$q_n$$ 为一个 $$(n-1)$$ 次多项式。

$$f_{\le t}$$ 可以表示为在 $${\mathcal{H}}_{\mu }$$ 的正交基下的分解，因此 $$g^{(t)}$$ 为 $$f_{\le t}$$ 在 $$\mathcal{G}$$ 上的限制，即 $$g^{(t)}$$ 在空间 $$\mathcal{G}$$ 的正交基上分解的系数表达式为 $$c_n^{(t)}:={⟨f_{\le t},q_n⟩}_{{\mu }^{(t)}}$$

2 HiPPO 架构

HiPPO 定义了两个映射：

1. $${\mathsf{proj}}_t$$：给定连续输入函数 $$f$$ 与概率测度 $${\mu }^{(t)}$$，将 $$f_{\le t}$$ 映射到上文中的近似函数 $$g^{(t)}$$
2. $${\mathsf{coef}}_t$$：将 $$g^{(t)}$$ 映射到系数向量 $$\mathbf{c}(t)\in {\mathbb{R}}^N$$

这两个映射的合成称为 $$\mathsf{hippo}:=\mathsf{coef}\circ \mathsf{proj}$$，也即 $$(\mathsf{hippo}(f))(t)=\mathbf{c}(t)$$

2.1 ODE

我们的最终目标是将系数向量表示为递推形式，便于训练与推理。为了达到这一离散的目标，需要在连续情形下先写成 ODE 的形式，再对其离散化。

HiPPO 的另一个关键思想是将系数的计算式对 $$t$$求导。由于正交多项式由 Gram-Schmidt 计算得来，这往往会导致自相似关系，从而写成一个ODE的形式。具体地，HiPPO 证明了对于文中提及的几个概率测度 $${\mu }^{(t)}$$，有特殊的 $$A(t)\in {\mathbb{R}}^{N\times N},B(t)\in {\mathbb{R}}^{N\times 1}$$ 满足：

$$\dot{\mathbf{c}}(t)=A(t)\mathbf{c}(t)+B(t)f(t)$$

2.2 如何离散化

在实际操作中，输入输出都相对时间离散，因此需要对上式进行离散化操作，整体分为四步：

1. 接收输入序列 $$(f_k{)}_{k\in \mathbb{N}}$$
2. 对于时间离散步长 $$\Delta t$$，隐性定义连续函数 $$f(t)$$，满足 $$f(k\cdot \Delta t)=f_k$$
3. 写出函数 $$\mathbf{c}(t)$$ 的 ODE 表达式
4. 重新离散为输出序列 $${\mathbf{c}}_k:=\mathbf{c}(k\cdot \Delta t)$$

通常，这一过程对于时间离散步长 $$\Delta t$$ 的选择是敏感的。

注意此处的 $${\mathbf{c}}_k$$ 为向量，与上文中有所不同。

2.3 广义双线性变换法（GBT）

离散化算法的选取有许多种，考虑到后续工作中使用了双线性变换法，此处介绍更为普适的广义双线性变换法（Generalized Bilinear Transformation, GBT）。

记 ODE 为 $$\dot{\mathbf{c}}(t):=h(t,\mathbf{c}(t))$$，可将其改写为

$$\mathbf{c}(t+\Delta t)-\mathbf{c}(t)={\int }_t^{t+\Delta t}h(s,\mathbf{c}(s))ds$$

GBT 通过取两端点的加权平均将右式近似为 $$\Delta t[(1-\alpha )h(t,\mathbf{c}(t))+\alpha h(t+\Delta t,\mathbf{c}(t+\Delta t))]$$ $$(\ast )$$，其中 $$\alpha \in [0,1]$$。

当 $$\dot{\mathbf{c}}(t)=A\mathbf{c}(t)+Bf(t)$$ 时将其代入 $$(\ast )$$，整理可得：

$$(I-\Delta t\alpha A)\mathbf{c}(t+\Delta t)=(I+\Delta t(1-\alpha )A)\mathbf{c}(t)+\Delta tBf(t)$$

因此

$$\mathbf{c}(t+\Delta t)=(I-\Delta t\alpha A{)}^{-1}(I+\Delta t(1-\alpha )A)\mathbf{c}(t)+\Delta t(I-\Delta t\alpha A{)}^{-1}Bf(t)$$

注意到 $${\mathbf{c}}_k:=\mathbf{c}(k\cdot \Delta t)$$，上式也即：

$${\mathbf{c}}_{k+1}=(I-\Delta t\alpha A{)}^{-1}(I+\Delta t(1-\alpha )A){\mathbf{c}}_k+\Delta t(I-\Delta t\alpha A{)}^{-1}B{f}_k$$

为一个关于 { c k } \{\boldsymbol{c}_k\} {ck​} 的递推式。

此处假定 $$A$$ 与 $$B$$ 均为时间不变量，若随时间改变，也可用相同思路计算。

在 GBT 中，$$\alpha$$ 取特殊值时有不同别名。当 $$\alpha =0$$ 时称前向欧拉法，$$\alpha =1$$ 时称反向欧拉法，$$\alpha =1/2$$ 时称双线性变换法。

3 实例：LegT，LagT，LegS

在选择多项式空间的基础上，根据不同的概率测度，该论文着重分析了三个实例：translated Legendre (LegT)，translated Laguerre (LagT)，scaled Legendre (LegS)。

下图为三种概率测度的图示，蓝色与紫色部分分别代表 $$t=t_0$$ 与 $$t=t_1$$ 时刻的 $${\mu }^{(t)}(x)$$

由于 LegS 的表现最好，下文将简要带过前两者，并主要关注于 LegS。

3.1 LegS

LegS 选取的概率测度为 $${\mu }^{(t)}=\frac{1}{t}{\mathbb{I}}_{[0,t]}$$，为所有历史分配相同的权重，随时间的推移缩放窗口。这是最直观的记忆方式。以下，将从这一多项式空间的正交基出发，计算 $$\mathbf{c}$$ 的导数，并得出 ODE 和递推式的表达。

3.1.1 正交基

首先，我们需要知道 $${\mathcal{H}}_{{\mu }^{(t)}}$$ 的正交基。考虑到测度的形式， $${\mathcal{H}}_{{\mu }^{(t)}}$$ 的正交基与 Legendre 多项式的数学表达式非常类似。

Legendre 多项式 { P n ( x ) } \{P_n(x)\} {Pn​(x)} 是在内积取 $$⟨f,g⟩={\int }_{-1}^{1}f(x)g(x)dx$$ 时，从自然基 { 1 , x , x 2 , ⋯   } \{1,x,x^2,\cdots\} {1,x,x2,⋯} 出发进行 Gram-Schmidt 计算得到的结果，使得：

$$\frac{2n+1}{2}{\int }_{-1}^1{P}_n(x)P_m(x)dx={\delta }_{nm}，{\delta }_{nm}=\{\begin{array}{ll}1,& i=j\\ 0,& i\ne j\end{array}$$

且有 $$P_n(1)，P_n(-1)=(-1{)}^n$$

$$P_n(x)$$ 也可用求导表示：

$$P_0(x)=1，P_n(x)=\frac{1}{2^{n}n!}\frac{d^n}{d{x}^n}(x^2-1{)}^n，x\in [-1,1]，n=1,2,\cdots$$

Legendre 多项式有以下性质：

$$\begin{array}{rl}(2n+1)P_n& =P_{n+1}^{\prime }-P_{n-1}^{\prime }\\ P_{n+1}^{\prime }& =(n+1)P_n+x{P}_n^{\prime }\end{array}$$

在当前的测度之下，考虑内积表达式 $${⟨q_n,q_m⟩}_{\mu }={\int }_0^t{q}_n(u)q_m(u)d\mu (u)$$，由于 $${\mu }^{(t)}=\frac{1}{t}{\mathbb{I}}_{[0,t]}$$，我们希望

$${\int }_0^t{q}_n(u)q_m(u)\frac{1}{t}du={\delta }_{nm}，{\delta }_{nm}=\{\begin{array}{ll}1,& i=j\\ 0,& i\ne j\end{array}$$

与 Legendre 多项式计算式中的积分表达与上下限进行比较可知，对 Legendre 多项式进行 $$x:=\frac{2u}{t}-1$$ 的变量替换，得到

$$\frac{2n+1}{t}{\int }_0^t{P}_n(\frac{2u}{t}-1)P_m(\frac{2u}{t}-1)du={\delta }_{nm}，{\delta }_{nm}=\{\begin{array}{ll}1,& i=j\\ 0,& i\ne j\end{array}$$

此时相差常数倍，因此 $${\mathcal{H}}_{{\mu }^{(t)}}$$ 的正交基即为将 Legendre 多项式进行缩放与平移的结果，可写作 $$q_n^{(t)}(x)=(2n+1{)}^{1/2}{P}_n(\frac{2x}{t}-1)$$

3.1.2 系数求导

记 $$\omega (t,x)=\frac{1}{t}{\mathbb{I}}_{[0,t]},q_n(t,x)=q_n^{(t)}$$ ，对 $$t$$ 求偏导，结合 Legendre 多项式的性质可知：

$$\frac{\partial }{\partial t}{q}_n(t,x)=-t^{-1}(2n+1{)}^{1/2}[n(2n+1{)}^{-\frac{1}{2}}{q}_n(t,x)+(2n-1{)}^{\frac{1}{2}}{q}_{n-1}(t,x)+\cdots ]$$

回忆到 $$c_n(t)={⟨f_{\le t},q_n⟩}_{{\mu }^{(t)}}$$，再让 $$c_n(t)$$ 对 $$t$$ 求偏导，整理可得：

$$\frac{d}{dt}{c}_n(t)=-t^{-1}(2n+1{)}^{1/2}[(n+1)(2n+1{)}^{-\frac{1}{2}}{c}_n(t)+(2n-1{)}^{\frac{1}{2}}{c}_{n-1}(t)+\cdots ]+t^{-1}(2n+1{)}^{1/2}f(t)$$

观察等式，发现 $$\mathbf{c}(t)$$ 每个维度的导数都可以表示成 $$c_k(t)$$ 与 $$f(t)$$ 的线性组合，这正是我们想要的形式。因此，可以写作矩阵形式 $$\dot{\mathbf{c}}(t)=-\frac{1}{t}A\mathbf{c}(t)+\frac{1}{t}Bf(t)$$ ，其中 $$A_{nk}$$ 为 $$\frac{d}{dt}{c}_n(t)$$ 中 $$c_k(t)$$ 的系数。

而对这一结果运用双线性变换法即可得到离散表达，完整结论见下文。需要注意的是， LegS 的离散形式独立于时间离散步长 $$\Delta t$$

3.1.3 结论

HiPPO-LegS 中在连续与离散情况下的变化为：$$\dot{\mathbf{c}}(t)=-\frac{1}{t}A\mathbf{c}(t)+\frac{1}{t}Bf(t)，{\mathbf{c}}_{k+1}=(1-\frac{A}{k}){\mathbf{c}}_k+\frac{1}{k}B{f}_k$$

其中 $$B_n=(2n+1{)}^{\frac{1}{2}}$$，$$A_{nk}=\begin{cases} (2n+1)^{1/2}(2k+1){1/2}, &n>k\ n+1, &n=k\ 0, &n

3.2 LegT & LagT

LegT 选取的概率测度为 $${\mu }^{(t)}(x)=\frac{1}{\theta }{\mathbb{I}}_{[t-\theta ,t]}(x)$$，为一长度固定的窗口。LagT 选取的测度为 $${\mu }^{(t)}(x)=e^{-(t-x)}{\mathbb{I}}_{(-\infty ,t]}(x)=\{\begin{array}{ll}{e}^{-(t-x)},& x\le t\\ 0,& x>t\end{array}$$，从当前时刻向前指数递减。

对于这两种测度，均有 ODE 为 $$\dot{\mathbf{c}}(t)=-A\mathbf{c}(t)+Bf(t)$$ 的形式，其中 $$A$$ 与 $$B$$ 如下（证明类似 LegS 的情形）：

LegT： $$A_{nk}=\frac{1}{\theta }\{\begin{array}{ll}(-1{)}^{n-k}(2n+1),& n\ge k\\ 2n+1,& n\le k\end{array}$$， $$B_n=\frac{1}{\theta }(2n+1)(-1{)}^n$$

LagT： KaTeX parse error: Expected & or \\ or \cr or \end at end of input: … k\\ 0, &n， B_n=1$$

4 LegS的优势与特性

在论文中，指出了三个 HiPPO-LegS 的优势，并分析了近似造成的误差幅度，以下对此进行介绍。

4.1 时间尺度的鲁棒性

LegS 具有时间等变性，也就是说，对任意的 $$\alpha >0$$，如果 $$h(t):=f(\alpha t)$$，则必有 $$\mathsf{hippo}(h)(t)=\mathsf{hippo}(f)(\alpha t)$$。这意味着它可以从某个时间尺度自然地迁移到其他尺度。

4.2 计算效率高

回忆 GBT 的推导过程，在更新 $${\mathbf{c}}_k$$ 时实质上在计算形如 $$(I-\Delta t\alpha A{)}^{-1}$$ 和 $$(I+\Delta t(1-\alpha )A)$$ 的矩阵与向量的乘积。由于 $$A$$ 的特殊结构，这一过程可以被大幅压缩，加快计算速度。这一结论在 Efficiently Modeling Long Sequences with Structured State Spaces 中有详细说明。

4.3 对历史信息遗忘速度慢

RNN 的遗忘速度为指数级，而 LegS 为多项式级，具体地有 $$\Vert \frac{\partial c_{l+1}}{\partial f_k}\Vert =\Theta (1/l)$$。

这是由于 $${\mathbf{c}}_l$$ 的递推形式，可以将其展开至由 $$f_k$$ 表示，求偏导知 $$\frac{\partial c_{l+1}}{\partial f_k}=(I-\frac{A}{l})\cdots (I-\frac{A}{k+1})\frac{B}{k}$$ 由于 $$I-\frac{A}{n}$$ 为可交换的可对角化矩阵（考虑到 $$A$$ 下三角），可写作 $$\frac{\partial c_{l+1}}{\partial f_k}=P^{-1}DPB$$，其中 $$D$$ 为对角矩阵，它的最大对角元为 $$\rho =(1-\frac{1}{l})\cdots (1-\frac{1}{k+1})\frac{1}{k}$$ 取对数后分析可知 $$\rho =\Theta (1/l)$$，因而原式成立。

这说明了 LegS 可以处理更长距离的内容。

4.4 误差与输入平滑度的关联

在 HiPPO 中有一步近似操作，其导致的误差与 $$f$$ 的光滑度密切相关。

• 当 $$f$$ 满足 L-Lipschitz 条件时（即对于 $$\forall x,y\in [0,t]$$，有 $$\left|f(x)-f(y)\right|\le L\left|x-y\right|$$ ），则 $$\Vert f_{\le t}-g^{(t)}\Vert =O(tL/\sqrt{N})$$
• 当 $$f$$ 有 $$k$$ 阶有界导数时，则 $$\Vert f_{\le t}-g^{(t)}\Vert =O(t^k{N}^{-k+1/2})$$

结语

文章根据在线函数近似的思路提出了 HiPPO 架构，解释了可以通过递推式反映历史的根本原因，主要理论核心是对若干概率测度的情形给出了 HiPPO 初始化实例，其中论证了取测度为 $${\mu }^{(t)}=\frac{1}{t}{\mathbb{I}}_{[0,t]}$$ 的 HiPPO-LegS 的优势。同时，ODE 的形式也启发后续对 SSM 的使用。

参考：https://zhuanlan.zhihu.com/p/613713261