HiPPO推导

Published on: 2026-03-04

Written by: Siriuns
description: Mathematical foundation of S4: a derivation of HiPPO.

关于S4架构HiPPO的推导

在隐空间中考虑LLM, 此时我们需要学习/预测隐空间中一个高维向量的走势, 模型的核心部分输入记为 $𝐱(t)\in {ℝ}^d$; 为了减少存储开销/简化模型, 我们将$𝐱(t)$的通道之间解耦, 从而仅需考虑一个输入为 $x(t)\in ℝ$ 的模型. SSM 模型假设系统为

$$\{\begin{array}{cc}{𝐡}^{\prime }(t)& =𝐀𝐡(t)+𝐁x(t)\\ 𝐡(0)& =0\end{array}$$

我们需要对该 ODE 离散化. 该方程为 一阶线性常系数微分方程, 有通用解法如下:

$$\begin{array}{cc}{𝐡}^{\prime }(t)-𝐀𝐡(t)& =𝐁x(t)\\ e^{-𝐀t}({𝐡}^{\prime }(t)-\mathrm{𝐀𝐡}(t))& =e^{-𝐀t}𝐁x(t)\\ e^{-𝐀t}𝐡(t)& ={\int }_0^t{e}^{-𝐀\tau }𝐁x(\tau )d\tau \\ 𝐡(t)& =e^{𝐀t}{\int }_0^t{e}^{-𝐀\tau }𝐁x(\tau )d\tau \\ \end{array}$$

此时我们采用 ZOH (即零阶保持假设) 来离散化, 记时间步为 $\mathrm{\Delta }\in ℝ$, 将 $𝐡(k\mathrm{\Delta }),x(k\mathrm{\Delta })$ 简记为 ${𝐡}_k,x_k$.

$$\begin{array}{cc}{𝐡}_k& =e^{k\mathrm{\Delta }𝐀}{\int }_0^{k\mathrm{\Delta }}{e}^{-𝐀\tau }𝐁x(\tau )d\tau \\ & =e^{k\mathrm{\Delta }𝐀}\sum _{r=0}^{k-1}{\int }_r^{r+1}{e}^{-\mathrm{\Delta }𝐀\tau }𝐁x_{r}d\tau \\ & =e^{k\mathrm{\Delta }𝐀}\sum _{r=0}^{k-1}(e^{-r\mathrm{\Delta }𝐀}-e^{-(r+1)\mathrm{\Delta }𝐀})\left({𝐀}^{-1}𝐁\frac{x_r}{\mathrm{\Delta }}\right)\\ \end{array}$$

于是不难发现

$$\begin{array}{cc}{𝐡}_{k+1}& =e^{(k+1)\mathrm{\Delta }𝐀}\sum _{r=0}^k(e^{-r\mathrm{\Delta }𝐀}-e^{-(r+1)\mathrm{\Delta }𝐀})\left({𝐀}^{-1}𝐁\frac{x_r}{\mathrm{\Delta }}\right)\\ & =e^{\mathrm{\Delta }𝐀}{𝐡}_k+(e^{\mathrm{\Delta }𝐀}-1){\left(\mathrm{\Delta }𝐀\right)}^{-1}𝐁x_r\\ \end{array}$$

此时我们记 $e^{\mathrm{\Delta }𝐀}$ 为 ${𝐀}_d$, 记 $(e^{\mathrm{\Delta }𝐀}-1){\left(\mathrm{\Delta }𝐀\right)}^{-1}𝐁$ 为 ${𝐁}_d$, 从而得到形式简洁的 一阶线性常系数差分方程.

$${𝐡}_{k+1}={𝐀}_d{𝐡}_k+{𝐁}_d{x}_r$$

对这个式子直接展开就可以得到

$${𝐡}_{k+1}=\sum _{r=0}^k{𝐀}_d^r{𝐁}_d\cdot x_{k-r}$$

此时方程可以进行并行运算, 相比原来的 RNN 可以更快训练. 我们记${𝐠}_r={𝐀}_d^r{𝐁}_d$, 可以把这个式子写成卷积的形式.

$$\begin{array}{cc}{𝐡}_{k+1}& =\sum _{r=0}^k{𝐠}_r{x}_{k-r}\\ & =(𝐠\ast x)[k]\\ \end{array}$$

这个卷积可以用 FFT 等算法来加速, 并且支持并行计算.