HiPPO推导

Published on: 2026-03-04

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

关于S4架构HiPPO的推导

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

{h′(t)=Ah(t)+Bx(t)h(0)=0\begin{cases} \mathbf{h}'(t) &= \mathbf{A}\mathbf{h}(t) + \mathbf{B}x(t) \\ \mathbf{h}(0) &= 0 \end{cases}

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

h′(t)−Ah(t)=Bx(t)e−At(h′(t)−Ah(t))=e−AtBx(t)e−Ath(t)=∫0te−AτBx(τ)dτh(t)=eAt∫0te−AτBx(τ)dτ\begin{aligned} \mathbf{h}'(t) - \mathbf{A}\mathbf{h}(t) &= \mathbf{B}x(t) \\ e^{-\mathbf{A}t} \left( \mathbf{h}'(t) - \mathbf{Ah}(t) \right) &= e^{-\mathbf{A}t}\mathbf{B}x(t) \\ e^{-\mathbf{A}t}\mathbf{h}(t) &= \int_0^t e^{-\mathbf{A}\tau}\mathbf{B}x(\tau) d\tau \\ \mathbf{h}(t) &= e^{\mathbf{A}t} \int_0^t e^{-\mathbf{A}\tau}\mathbf{B}x(\tau) d\tau \\ \end{aligned}

此时我们采用 ZOH (即零阶保持假设) 来离散化, 记时间步为 Δ∈R\Delta \in \mathbb{R}, 将 h(kΔ),x(kΔ)\mathbf{h}(k \Delta), x(k \Delta) 简记为 hk,xk\mathbf{h}_k, x_k.

hk=ekΔA∫0kΔe−AτBx(τ)dτ=ekΔA∑r=0k−1∫rr+1e−ΔAτBxrdτ=ekΔA∑r=0k−1(e−rΔA−e−(r+1)ΔA)(A−1BxrΔ)\begin{aligned} \mathbf{h}_k &= e^{k \Delta \mathbf{A}} \int_0^{k \Delta} e^{-\mathbf{A}\tau}\mathbf{B}x(\tau) d\tau \\ &= e^{k \Delta \mathbf{A}} \sum_{r=0}^{k-1} \int_{r}^{r+1} e^{-\Delta \mathbf{A}\tau}\mathbf{B}x_r d\tau \\ &= e^{k \Delta \mathbf{A}} \sum_{r=0}^{k-1} \left(e^{- r \Delta \mathbf{A}} - e^{- (r+1) \Delta \mathbf{A}} \right) \left(\mathbf{A}^{-1}\mathbf{B}\frac{x_r}{\Delta}\right) \\ \end{aligned}

于是不难发现

hk+1=e(k+1)ΔA∑r=0k(e−rΔA−e−(r+1)ΔA)(A−1BxrΔ)=eΔAhk+(eΔA−1)(ΔA)−1Bxr\begin{aligned} \mathbf{h}_{k+1} &= e^{(k+1) \Delta \mathbf{A}} \sum_{r=0}^{k} \left(e^{- r \Delta \mathbf{A}} - e^{- (r+1) \Delta \mathbf{A}} \right) \left(\mathbf{A}^{-1}\mathbf{B}\frac{x_r}{\Delta}\right) \\ &= e^{\Delta \mathbf{A}} \mathbf{h}_k + \left(e^{\Delta \mathbf{A}} - 1\right) \left(\Delta\mathbf{A}\right)^{-1}\mathbf{B}x_r \\ \end{aligned}

此时我们记 eΔAe^{\Delta \mathbf{A}} 为 Ad\mathbf{A}_d, 记 (eΔA−1)(ΔA)−1B\left(e^{\Delta \mathbf{A}} - 1\right) \left(\Delta\mathbf{A}\right)^{-1}\mathbf{B} 为 Bd\mathbf{B}_d, 从而得到形式简洁的 一阶线性常系数差分方程.

hk+1=Adhk+Bdxr\mathbf{h}_{k+1} = \mathbf{A}_d \mathbf{h}_k + \mathbf{B}_d x_r

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

hk+1=∑r=0kAdrBd⋅xk−r\mathbf{h}_{k+1} = \sum_{r = 0}^{k} \mathbf{A}_d^r \mathbf{B}_d \cdot x_{k - r}

此时方程可以进行并行运算, 相比原来的 RNN 可以更快训练. 我们记gr=AdrBd\mathbf{g}_r = \mathbf{A}_d^r\mathbf{B}_d, 可以把这个式子写成卷积的形式.

hk+1=∑r=0kgrxk−r=(g∗x)[k]\begin{aligned} \mathbf{h}_{k+1} &= \sum_{r=0}^{k} \mathbf{g}_r x_{k-r} \\ &= (\mathbf{g} * x)[k] \\ \end{aligned}

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