はじめに
この前のブログでは、mambaの論文を翻訳した。
https://izmyon.hatenablog.com/entry/2023/12/11/155551
本シリーズでは、mambaの理論的背景を理解するために、それらの先行研究を順々にまとめて解説していく。重要な先行研究としては、LMU, HiPPO, LSSL, S4, H3などがあり、主に以下の流れで解説する。
HiPPOフレームワークとLSSL
S4のアルゴリズム
H3とHyena Hierarchy
mamba
本記事は、第一段として、HiPPOフレームワークとLSSLについてまとめる。
元論文を以下に記す。
HiPPO
https://arxiv.org/abs/2008.07669
LSSL
https://arxiv.org/abs/2110.13985
この記事の続きはこちら
https://zenn.dev/izmyon/articles/c56a2fd6670546
読者の方へ
補足や訂正などがあれば、コメントにて優しく丁寧にご教示いただけると喜びます。
書くのけっこう大変だったのでバッチを送っていただけると嬉しいです。たくさんもらえると僕のやる気が上がって投稿頻度が上がるかもしれません。（逆にもらえないと下がるかもしれない。）
状態空間モデル
時刻t u(t) N x(t) y(t) L L \mathbb{R}^L \rightarrow \mathbb{R}^L
\begin{align}
\dot{x}(t) = A x(t) + B u(t) \quad
\\
y(t) = C x(t) + D u(t) \quad
\end{align}
状態空間モデルは、Hidden Markov Models (HMM）のように、制御工学や機械学習などの幅広い分野でみられる潜在空間を利用するモデルであり、行列A \in \mathbb{R}^{N \times N} B \in \mathbb{R}^{N \times 1} C \in \mathbb{R}^{1 \times N} D \in \mathbb{R}^{1 \times 1}
これは連続時間のシステムであるが、コンピュータ上では離散時間信号を扱うため、式(1)、式(2)の離散化を行う。ここで、式(2)の離散化は簡単だが問題なのは式(1)の常微分方程式（ODE）を離散化することである。
ゼロ次ホールド（ZOH）と同様に、時刻t u(t) \Delta t u(t+\Delta t) = u(t) t x(t) \dot{x}(t) = f(t, x(t)) t_i x(t_0), x(t_1), \ldots x(t_{i+1}) = x(t_i) + \int^{t_{i+1}}_{t_i} f(s, x(s))ds \alpha
\begin{align*}
\begin{split}
&\int^{t+\Delta}_{t} f(s, x(s))ds \\
&\simeq \alpha \Delta t f(t+\Delta, x(t+\Delta t)) + (1-\alpha) \Delta t f(t, x(t)) \\
&= \alpha \Delta t (A x(t+\Delta) + B u(t+\Delta)) + (1-\alpha) \Delta t (A x(t) + B u(t)) \\
&= \alpha \Delta t (A x(t+\Delta) + B u(t)) + (1-\alpha) \Delta t (A x(t) + B u(t)) \quad (\because u(t+\Delta) = u(t)) \\
&= \alpha \Delta t A x(t+ \Delta) + (1-\alpha) \Delta t A x(t) + B \Delta t u(t)
\end{split}
\end{align*}
この結果を用いると、x(t+\Delta t) = x(t) + \int^{t+\Delta t}_{t} f(s, x(s))ds
\begin{align*}
x(t+ \Delta t)
&= x(t) + \alpha \Delta t A x(t+ \Delta t) + (1-\alpha) \Delta t A x(t) + B \Delta t u(t) \notag \\
(I- \alpha \Delta t A) x(t+\Delta t)
&= (I + (1-\alpha) \Delta t A) x(t) + B \Delta t u(t) \notag \\
\therefore x(t+\Delta t)
&= \left(I - \alpha \Delta t A \right)^{-1} \left(I + (1- \alpha) \Delta t A\right) x(t) + \left( I - \alpha \Delta t A \right)^{-1} B \Delta t u(t)
\end{align*}
ここでは、\alpha = \frac{1}{2}
\begin{align}
\begin{split}
x(t+\Delta t)
= \left(I - \frac{1}{2} \Delta t A \right)^{-1} \left(I + \frac{1}{2} \Delta t A\right) x(t) + \left( I - \frac{1}{2} \Delta t A \right)^{-1} B \Delta t u(t)
\end{split}
\end{align}
さらに、以下のように行列A, B \bar{A}, \bar{B}
\begin{align}
\bar{A} = \left(I - \frac{1}{2} \Delta t A \right)^{-1} \left(I + \frac{1}{2} \Delta t A\right)
\\
\bar{B} = \left( I - \frac{1}{2} \Delta t A \right)^{-1} B \Delta t
\end{align}
すると、以下のように式(1), (2)の連続時間システムを離散化することが出来る。
\begin{align}
x_t &= \bar{A} x_{t-1} + \bar{B} u_t \quad
\\
y_t &= C x_t + D u_t \quad
\end{align}
!
GBTは、\alpha = 0 x(t+\Delta t)= x(t) + \Delta t f(t, x(t)) \alpha = 1 x(t) = x(t + \Delta t) - \Delta t f(t+\Delta t, x(t+\Delta t)) \alpha \in [0, 1]
積分\int^{t+\Delta}_{t} f(s, x(s))ds \alpha=0 f(t, x(t)) \alpha=1 f(t+\Delta t, x(t+\Delta t)) \alpha = \frac{1}{2} f(t, x(t)) f(t+\Delta t, x(t+\Delta t))
HiPPO (High-Order Polynomial Projection Operator; 高次多項式投影演算子)
ここまで線形システム(1)、(2)を離散化し、新たなシステムである式(6),(7)を得た。次に、本節で説明するHiPPOと呼ばれるフレームワークを適用し、以下のように係数行列A
A_{nk} = \left\{
\begin{array}{ll}
(2n+1)^{1/2} (2k+1)^{1/2} & \text{if} \ n>k \\
n+1 & \text{if} \ n=k \\
0 & \text{if} \ n<k \\
\end{array}
\right.
なぜこのような初期化を行うのかをこれから説明化するが、長いのでLSSLの全体像を先につかみたい場合は、次の節「畳み込みによる高速化」に進んでもらって構わない。
HiPPOの定義
時間変化する(-\infty, t] \mu^{(t)} N \mathcal{G} u: \mathbb{R}_{\geq 0} \rightarrow \mathbb{R} u x: \mathbb{R}_{\geq 0} \rightarrow \mathbb{R}^N \text{hippo} \text{proj}_t \text{coef}_t t \text{coef}_t \circ \text{proj}_t (\text{hippo}(u))(t)=\text{coef}_t (\text{proj}_t (u))
\text{proj}_t t u u_{<t}:=\{u(y)\}_{y \leq t} \| u_{\leq t} - g^{(t)} \|_{L_2 (\mu^{(t)})} g^{(t)} \in \mathcal{G}
\text{coef}_t: \mathcal{G} \rightarrow \mathbb{R}^N g^{(t)} \mu^{(t)} x(t) \in \mathbb{R}^N
HiPPOを導くためのHiPPOフレームワーク
\{P_n^{(t)}\}_{n \in \mathbb{N}} \mu^{(t)} p_n^{(t)} P_n^{(t)} p_n^{(t)} = P_n^{(t)}/ \| P_n^{(t)} \|_{\mu}^2
\begin{align}
\int^{\infty}_{-\infty} p_n^{(t)}(y) p_m^{(t)}(y) \omega^{(t)}(y) dy = \delta_{m,n}
\end{align}
ここで、\omega^{(t)}(y) d \mu^{(t)} = \omega^{(t)}(y) dy, \int^{\infty}_{-\infty} d \mu^{(t)} = 1
入力信号の履歴u_{\leq t}:=\{u(y)\}_{y \leq t} \Delta t u(y) \{p_n^{(t)}\}_{n=0,1,\ldots, N-1} u_{\leq t} u(y) g^{(t)} u_{\leq t} \{p_n^{(t)}\}_{n=0,1,\ldots, N-1} g^{(t)}
\begin{align}
u_{\leq t} \simeq g^{(t)} = \sum_{n=0}^{N-1} x_n (t) p_n^{(t)}
\end{align}
このように近似すると、時刻t N \{x_n^{(t)}\}_{n=0,1,\ldots, N-1} x_n (t) t u_{\leq t}
\begin{align}
x_n (t) = \int^{\infty}_{-\infty} g^{(t)} p_n^{(t)} \omega^{(t)}(y) dy
&\simeq \int^{\infty}_{-\infty} u_{\leq t} p_n^{(t)} \omega^{(t)}(y) dy
\end{align}
実は、式(10)に従って、以下の入力信号u_{\leq t} N \text{hippo}
\begin{align*}
\text{hippo}(u_{\leq t}) := (x_0, x_1, \ldots, x_{N-1}) \\
\text{where} \ u_{\leq t} \simeq g^{(t)} = \sum_{n=0}^{N-1} x_n (t) p_n^{(t)}
\end{align*}
しかしながら、これらの係数を各時刻で計算するのは容易ではない。そこでHiPPOフレームワークでは、式(1)を離散化して式(6)の更新式を導いたように、x(t) x(t)
!
ということでHiPPOでは式(10)を用いてx_n(t)
まず、u_{\leq t} k u[k], k = 0, \ldots, L-1 u_{\leq t}
\begin{align*}
u_{\leq t}(y) = \left\{
\begin{array}{ll}
u[k] & \text{if} \quad y = k\Delta t, \quad k = 0, \ldots, L-1\\
0 & otherwise \\
\end{array}
\right.
\end{align*}
これを用いて、
\begin{align*}
\int^{\infty}_{-\infty} u_{\leq t} p_n (t, y) \omega (t,y) dy
&= \int^{\infty}_{-\infty} ( \sum^{L-1}_{k=0} u[k]\delta(y-k \Delta t) ) p_n (t, y) \omega (t,y) dy
\\
&= \sum^{L-1}_{k=0} u[k] \int^{\infty}_{-\infty} \delta(y-k \Delta t) p_n (t, y) \omega (t,y) dy
\\
&= \sum^{L-1}_{k=0} u[k] p_n (t, k \Delta t) \omega (t, k \Delta t)
\end{align*}
従って、以下の式でx_n(t)
\begin{align*}
x_n (t) \simeq \sum^{L-1}_{k=0} u[k] p_n (t, k \Delta t) \omega (t, k \Delta t)
\end{align*}
ちなみにここで、n \omega p_{\omega}^{(t)} (y) = e^{-j \omega y} \omega (t, y) = 1 x^{(t)} (\omega)
\begin{align*}
x^{(t)} (\omega) = \sum^{L-1}_{k=0} u[k] e^{-j \omega k \Delta t}
\end{align*}
となり、これは離散フーリエ変換（DFT）と考えることが出来る。
係数が満たすODE
以上より、下式が成り立つことが分かった。
\begin{align}
x_n (t) = \int^{\infty}_{-\infty} u(y) p_n (t, y) \omega(t, y) dy
\end{align}
両辺を微分して、係数x^{(t)}_n
\begin{align*}
\dot{x}_n (t) &= \frac{\partial}{\partial t} \int^{\infty}_{-\infty} u(y) p_n (t, y) \omega(t, y) dy
\\
&= \int^{\infty}_{-\infty} \frac{\partial}{\partial t} (u(y) p_n (t, y) \omega(t, y) ) dy
\\
&= \int^{\infty}_{-\infty} u(y) (\frac{\partial}{\partial t} p_n (t, y)) \omega (t, y) dy + \int^{\infty}_{-\infty} u(y) p_n (t, y) (\frac{\partial}{\partial t} \omega (t, y) ) dy
\end{align*}
以上より、以下の式(11)が成り立つ。
\begin{align}
\dot{x}_n (t)
= \int^{\infty}_{-\infty} u(y) (\frac{\partial}{\partial t} p_n (t, y)) \omega (t, y) dy + \int^{\infty}_{-\infty} u(y) p_n (t, y) (\frac{\partial}{\partial t} \omega (t, y) ) dy
\end{align}
HiPPOフレームワークでは、具体的に用いる基底\{P_n^{(t)}\}_{n \in \mathbb{N}} \mu^{(t)} \dot{x}(t) x(t) u(t) A B
HiPPO-LegS
HiPPOの論文では、様々な基底や測度を用いた場合について検証されており、LSSLではその中で最も良い結果を示したScaled Legendre Measure (LegS)という設定を用いて得られる行列A \omega_t = \frac{1}{t} \mathbb{1}_{[0, t]}
ここではまず、ルシャンドル多項式の性質や、基底として用いた場合の\omega_t
正規直交規定および測度
ルシャンドル多項式は、以下のn
\begin{align}
P_n (x) := \frac{1}{2^n n!} \frac{d^n}{{dx}^n} [ (x^2 - 1 )^n ]
\end{align}
測度\omega^{\text{leg}} = \bold{1}_{[-1, 1]}
\begin{align}
\frac{2n+1}{2} \int^1_{-1} P_n(x) P_m(x) dx = \delta_{nm}
\end{align}
さらに、以下の性質を持つ。
\begin{align}
P_n (1) &= 1
\\
P_n (-1) &= (-1)^n
\\
(2n+1) P_n &= P'_{n+1} - P'_{n-1}
\\
P'_{n+1} &= (n+1) P_n + x P'_n
\end{align}
以上のルシャンドル多項式は区間x \in [-1, 1] y = \frac{t}{2}(x+1) x = \frac{2}{t} y - 1 dx = \frac{2}{t} dy
\begin{align*}
& \frac{2n+1}{2} \int^1_{-1} P_n(x) P_m(x) dx
\\
=& \frac{2n+1}{2} \int^t_{0} P_n( \frac{2}{t} y - 1 ) P_m( \frac{2}{t} y - 1 ) \frac{2}{t}dy
\\
=& (2n+1) \int^t_{0} P_n( \frac{2}{t} y - 1 ) P_m( \frac{2}{t} y - 1 ) \omega^{\text{leg}} (\frac{2}{t} y - 1 ) \frac{1}{t} dy
\end{align*}
従って、以下が成り立つ。
\begin{align}
(2n+1) \int^t_{0} P_n( \frac{2}{t} y - 1 ) P_m( \frac{2}{t} y - 1 ) \omega^{\text{leg}} (\frac{2}{t} y - 1 ) \frac{1}{t} dy = \delta_{nm}
\end{align}
この式より、新たな測度\omega_t = \frac{1}{t} \mathbb{1}_{[0, t]}
\begin{align}
p_n (t, y) = (2n+1)^{1/2} P_n ( \frac{2}{t} y - 1 )
\end{align}
すると、以下が成り立つ。
\begin{align}
\int^t_{0} p_n(t, y) p_m(t, y) \omega_t dy = \delta_{nm}
\end{align}
ODEの導出
HiPPO-LegSが用いるODEを導出する前に、導出で用いるルシャンドル基底の微分が満たす性質について述べる。まず、式(17)から、以下が導かれる。
\begin{align}
P'_{n+1} = (2n+1) P_n + (2n-3)P_{n-2} + \ldots,
\end{align}
次数を一つずらすと、
\begin{align}
P'_{n} = (2n-1) P_{n-1} + (2n-5)P_{n-3} + \ldots,
\end{align}
これらの性質を用いて(x+1)P'_n (x)
\begin{align}
(x+1) P'_{n}
=& x P'_{n} + P'_{n} = P'_{n+1} + P'_n - (n+1) P_n \notag
\\
&(\because \text{式(18)より、} xP'_n = P'_n - (n+1) P_n) \notag
\\
=& nP_n + (2n-1) P_{n-1} + (2n-3)P_{n-2} + \ldots
\\
&(\because \text{式(22), (23)を代入}) \notag
\\
\end{align}
まず、\omega_t p_n(t, y)
\begin{align}
\frac{\partial}{\partial t} \omega(t) =& -t^{-2} \mathcal{1}_{[0,1]} + t^{-1} \delta_t \notag
\\
=& t^{-1} ( - \omega (t) + \delta_t)
\\
\frac{\partial}{\partial t} p_n (t, y) &= - (2n+1)^{\frac{1}{2}} 2y t^{-2} P_n' (\frac{2}{t} y - 1 ) \notag
\\
=& - (2n+1)^{\frac{1}{2}} t^{-1} ((\frac{2}{t} y - 1) + 1 ) P_n' (\frac{2}{t} y - 1 ) \notag
\\
=& - (2n+1)^{\frac{1}{2}} t^{-1} (x + 1 ) P_n' (x) \ (\because x = \frac{2}{t} y - 1) \notag
\\
=& - (2n+1)^{\frac{1}{2}} t^{-1} [nP_n (x) + (2n-1) P_{n-1}(x) + (2n-3) P_{n-2}(x) + \ldots ] \ (\because 式(24)を用いた。) \notag
\\
=& - (2n+1)^{\frac{1}{2}} t^{-1} [n (2n+1)^{-\frac{1}{2}} p_n (t, y) + (2n-1)^{\frac{1}{2}} p_{n-1} (t, y) + (2n-3)^{\frac{1}{2}} p_{n-2}(t, y) + \ldots ] \notag
\\
& (\because 式(20)より、P_n (x) = (2n+1)^{-\frac{1}{2}} p_n (t, y) )
\end{align}
これらを用いて、式(12)のODEの右辺を計算する。
まず、第一項を求める。
\begin{align}
&\int^t_0 u(y) (\frac{\partial}{\partial t} p_n (t, y)) \omega (t, y) dy \notag
\\
\simeq & \int^t_0 g (t, y) (\frac{\partial}{\partial t} p_n (t, y)) \omega (t, y) dy \notag
\\
=&\int^t_0 ( \sum_{n=0}^{N-1} x_n (t) p_n (t, y) ) \{ - (2n+1)^{\frac{1}{2}} t^{-1} [n (2n+1)^{-\frac{1}{2}} p_n (t, y) + (2n-1)^{\frac{1}{2}} p_{n-1} (t, y) + (2n-3)^{\frac{1}{2}} p_{n-2}(t, y) + \ldots ] \} \omega (t, y) dy \notag
\\
&(\because 式(26)を用いた。) \notag
\\
=& - (2n+1)^{\frac{1}{2}} t^{-1} [n (2n+1)^{-\frac{1}{2}} x_n (t) + (2n-1)^{\frac{1}{2}} x_{n-1} (t) + (2n-3)^{\frac{1}{2}} x_{n-2}(t) + \ldots ]
\\
&(\because 式(21)の正規直交性を用いた。) \notag
\\
\end{align}
つぎに、第二項を求める。
\begin{align}
&\int^t_0 u(y) p_n (t, y) (\frac{\partial}{\partial t} \omega (t, y) ) dy \notag
\\
=&\int^t_0 u(y) p_n (t, y) ( t^{-1} ( - \omega (t) + \delta_t) ) dy \ (\because 式(25)を用いた。) \notag
\\
=& -t^{-1} \int^t_0 u(y) p_n (t, y) \omega (t) dy + t^{-1} \int^t_0 u(y) p_n (t, y) \delta_t dy \notag
\\
=& -t^{-1} x_n (t) + t^{-1} u(t) p_n (t, t) \notag
\\
=& -t^{-1} x_n (t) + t^{-1} (2n+1)^{1/2} u(t)
\end{align}
式(12)の右辺に式(27), (28)を代入して、以下が成り立つ。
\begin{align}
\dot{x}_n (t)
=& - (2n+1)^{\frac{1}{2}} t^{-1} [n (2n+1)^{-\frac{1}{2}} x_n (t) + (2n-1)^{\frac{1}{2}} x_{n-1} (t) + (2n-3)^{\frac{1}{2}} x_{n-2}(t) + \ldots ] \notag
\\
&-t^{-1} x_n (t) + t^{-1} (2n+1)^{1/2} u(t) \notag
\\
=& - (2n+1)^{\frac{1}{2}} t^{-1} [(n+1) (2n+1)^{-\frac{1}{2}} x_n (t) + (2n-1)^{\frac{1}{2}} x_{n-1} (t) + (2n-3)^{\frac{1}{2}} x_{n-2}(t) + \ldots ] \notag
\\
& + t^{-1} (2n+1)^{1/2} u(t)
\end{align}
式(29)をベクトル化すると、以下の式(30)が成立する。
\begin{align}
\dot{x}(t) &= - \frac{1}{t} A x(t) + \frac{1}{t} B u(t)
\\
A_{nk} &= \left\{
\begin{array}{ll}
(2n+1)^{1/2} (2k+1)^{1/2} & \text{if} \ n>k \\
n+1 & \text{if} \ n=k \\
0 & \text{if} \ n<k \\
\end{array}
\right.
\\
B_n &= (2n+1)^{\frac{1}{2}}
\\
\end{align}
式(31)の行列A
!
式(21)の直交性を用いると、式(11)より、
\begin{align}
x_n (t) &= \int^{\infty}_{-\infty} u(y) p_n (t, y) \omega(t, y) dy \notag
\\
&= \int^t_0 u(y) p_n (t, y) \omega(t, y) dy
\end{align}
としてみる。
ライプニッツの積分法則
\begin{align*}
\frac{\partial}{\partial t} \int^{\beta (t)}_{\alpha (t)} f(t, y) dy
= \int^{\beta (t)}_{\alpha (t)} \frac{\partial}{\partial t} f(t, y) dy - \alpha' (t) f(\alpha(t), t) + \beta' (t) f( \beta(t), t)
\end{align*}
について考え、ここで、\beta (t) = t \alpha (t) = 0
\begin{align*}
\frac{d}{dt} \int^t_0 f(t, y) dy
&= f(t, t) \frac{d}{dt} t - F(t,0)\frac{d}{dt} 0 + \int^t_0 \frac{\partial}{\partial t} f(t, y) dy
\\
&= \int^t_0 \frac{\partial}{\partial t} f(t, y) dy + f(t, t)
\end{align*}
f (t y) = u(y) p_n (t, y) \omega(t, y)
\begin{align*}
\dot{x}_n (t)
&= \frac{d}{dt} \int^t_0 u(y) p_n (t, y) \omega(t, y) dy
\\
&= \frac{d}{dt} \int^t_0 f(t, y) dy
\\
&= \int^t_0 \frac{\partial}{\partial t} f(t, y) dy + f(t, t) \quad (\because ライプニッツの積分法則）
\\
&= \int^t_0 u(y) (\frac{\partial}{\partial t} p_n (t, y)) \omega (t, y) dy + \int^t_0 u(y) p_n (t, y) (\frac{\partial}{\partial t} \omega (t, y) ) dy
\\ & \quad + u(t) p_n(t,t) \omega (t, t)
\end{align*}
となり、式(12)と比較すると、第三項が増える。
実は、HiPPO-LegSにおいて、この第三項を計算すると、u(t) p_n(t,t) \omega (t, t) = (2n+1)^{\frac{1}{2}} t^{-1} u(t) B_n = 2(2n+1)^{\frac{1}{2}}
https://github.com/HazyResearch/hippo-code/issues/11
HiPPO-LegSは時間スケールに依存しない。
式(30)のODEを計算するにあたり、式(6)を求めたように両辺を積分し、GBT（\alpha=\frac{1}{2} x_n (t)
\begin{align}
x(t+\Delta t) - x(t) &= \int^{t+\Delta t}_t - \frac{1}{t} A x(t) + \frac{1}{t} B u(t) dt \notag
\\
&\simeq \frac{\Delta t}{2}{(-\frac{1}{t}Ac(t) + \frac{1}{t}Bu(t))+(-\frac{1}{t+\Delta t}Ac(t+\Delta t) + \frac{1}{t+\Delta t}Bu(t+\Delta t))} \notag
\end{align}
ゼロ次ホールドを仮定し、u(t) = u(t+\Delta t)
\begin{align}
x(t+\Delta t) - x(t) &= \frac{\Delta t}{2}(-\frac{1}{t}Ax(t) + \frac{1}{t}Bu(t)) + \frac{\Delta t}{2}(-\frac{1}{t+\Delta t}Ax(t+\Delta t) + \frac{1}{t+\Delta t}Bu(t)) \notag
\\
x(t+\Delta t) + \frac{\Delta t}{2(t+\Delta t)}Ax(t+\Delta t) &= x(t) - \frac{\Delta t}{2t} Ax(t) + (\frac{\Delta t}{2t} + \frac{\Delta t}{2(t+\Delta t)})Bu(t) \notag
\\
(I + \frac{\Delta t}{2(t+\Delta t)}A) x(t+\Delta t) &= (I - \frac{\Delta t}{2t} A)x(t) + (\frac{\Delta t}{2t} + \frac{\Delta t}{2(t+\Delta t)})Bu(t)
\end{align}
ここで、時間を離散化し、t=k\Delta t x_k := x(k\Delta t) u_k := u(k \Delta t)
\begin{align*}
(I + \frac{\Delta t}{2(k\Delta t+\Delta t)}A) x_{k+1} &= (I - \frac{\Delta t}{2k\Delta t} A)x_k + (\frac{\Delta t}{2k\Delta t} + \frac{\Delta t}{2(k\Delta t+\Delta t)})Bu_k
\end{align*}
従って、以下が成り立つ。
\begin{align}
(I + \frac{1}{2(k+1)}A) x_{k+1} &= (I - \frac{1}{2k} A)x_k + (\frac{1}{2k} + \frac{1}{2(k+1)})Bu_k
\end{align}
式(35)では、時間スケール\Delta t
この式に従ってx_n (t) u_{\leq t}
\begin{align*}
u_{\leq t} \simeq g^{(t)} &= \sum_{n=0}^{N-1} x_n (t) p_n (t, y)
\\
&= \sum_{n=0}^{N-1} x_n (t) (2n+1)^{\frac{1}{2}} P_n (\frac{2y}{t} - 1)
\end{align*}
LSSL(Linear State-Space Layers)
ここまで、HiPPOおよびHiPPO-LegSについて説明してきた。HiPPO-LegSでは、式(30)のODEが導かれ、式(35)のように離散化された。LSSLやその発展であるS4、H3、mambaなどでは、式(1),(2)の線形システムを考え、式(6)、(7)の形に離散化するが、行列A x(t) u_{\leq t}
畳み込みによる高速化
式(6),(7)を見た時に思ったかもしれないが、式(6)を式(7)に代入することで、以下のようにx u_{\leq t} y_k x_{-1} = 0
\begin{align*}
y_k =& C x_k + D u_k
\\
=& C (\bar{A} x_{k-1} + \bar{B} u_k) + D u_k
\\
=& C \bar{A} (\bar{A} x_{k-2} + \bar{B} u_{k-1}) + C \bar{B} u_k + D u_k
\\
=& C (\bar{A})^2 (\bar{A} x_{k-3} + \bar{B} u_{k-2}) + C \bar{A} \bar{B} u_{k-1} + C \bar{B} u_k + D u_k
\\
=& C (\bar{A})^3 (\bar{A} x_{k-4} + \bar{B} u_{k-3}) + C (\bar{A})^2 \bar{B} u_{k-2} + C \bar{A} \bar{B} u_{k-1} + C \bar{B} u_k + D u_k
\\
\vdots
\\
=& C (\bar{A})^k x_0 + \cdots + C (\bar{A})^2 \bar{B} u_{k-2} + C \bar{A} \bar{B} u_{k-1} + C \bar{B} u_k + D u_k
\\
=& C (\bar{A})^k (\bar{A} x_{-1} + \bar{B} u_0) + \cdots + C (\bar{A})^2 \bar{B} u_{k-2} + C \bar{A} \bar{B} u_{k-1} + C \bar{B} u_k + D u_k
\\
=& C (\bar{A})^k \bar{B} u_0 + \cdots + C (\bar{A})^2 \bar{B} u_{k-2} + C \bar{A} \bar{B} u_{k-1} + C \bar{B} u_k + D u_k
\\
&(\because x_{-1} = 0)
\end{align*}
従って、\mathcal{K}_k (\bar{A}, \bar{B}, C) = (C \bar{B}, C \bar{A} \bar{B}, \ldots, C (\bar{A})^k \bar{B})
\begin{align}
y_k &= (C \bar{B}, C \bar{A} \bar{B}, \ldots, C (\bar{A})^k \bar{B}) * u_{\leq k} + D u_k \notag
\\
&= \mathcal{K}_k (\bar{A}, \bar{B}, C) * u_{\leq k} + D u_k
\end{align}
ここで、以下のように置いた。
\begin{align}
\mathcal{K}_k (\bar{A}, \bar{B}, C)
&= (C \bar{A}^i B)_{i = 0, \ldots, k} \notag
\\
&= (C \bar{B}, C \bar{A} \bar{B}, \ldots, C (\bar{A})^k \bar{B})
\\
u_{\leq k} &= (u_0, u_1, \ldots, u_k)
\end{align}
また、系列長L u_{\leq L-1} = (u_0, u_1, \ldots, u_{L-1}) y_{L-1}
\begin{align}
y_{L-1} = \mathcal{K}_{L-1} (\bar{A}, \bar{B}, C) * u_{\leq L-1} + D u_{L-1}
\end{align}
このように畳み込み形式で計算することで、全時刻にわたるy \in \mathbb{R}^{H \times L}
ただし、LSSLの計算にはボトルネックがある。式(6)の再帰形式では、離散化された状態行列A \bar{A} \bar{A} \mathcal{K}_L
LSSLの表現力
ここまで、LSSLが式(6)、(7)の再帰形式、式(39)の畳み込み形式の二つの形式を持つことを見てきた。LSSLでは線形計算のみを行い、RNNやTransformerなどで見られるような非線形関数は用いていない。それでもLSSLは高い表現力を持つことを示す。
畳み込みはLSSLである。
Under Construction
RNNはLSSLである。
Under Construction
Lemma [3, p.6, 3.1]
一次元ゲート付き再帰\bold{x_t = (1 - \sigma(z)) x_{t-1} + \sigma(z) u_t} \bold{\sigma} \bold{z} \bold{\dot{x}(t) = -x(t) + u(t)} \bold{\alpha=1}
証明：\dot{x}(t) = -x(t) + u(t)
\begin{align*}
x_{t-1} = x_t - \Delta t (- x_t + u_t)
\end{align*}
ここで、\Delta t = e^z
\begin{align*}
x_{t-1} = x_t - e^z (- x_t + u_t) = (1 + e^z)x_t - e^z u_t
\end{align*}
従って、以下が成り立つ。
\begin{align}
x_t &= \frac{1}{1 + e^z} x_{t-1} + \frac{e^z}{1 + e^z} u_t \notag
\\
&= \frac{e^{-z}}{1 + e^{-z}} x_{t-1} + \frac{1}{1 + e^{-z}} u_t \notag
\\
&= (1 - \sigma(z)) x_{t-1} + \sigma(z) u_t \quad (\because \sigma(z) = \frac{1}{1 + e^{-z}})
\end{align}
(証明終)
Lemma [3, p.6, 3.2]
Under Construction
Deep LSSLs
Under Construction
参考文献
[1] Voelker, A., Kajić, I., & Eliasmith, C. (2019). Legendre memory units: Continuous-time representation in recurrent neural networks. Advances in neural information processing systems, 32.https://techblog.morphoinc.com/entry/2022/05/24/102648 , 2024年2月10日閲覧
Discussion