Recurrent, LSTM, and structured state-space models — Examples

Take $d_x=d_h=1$, set every LSTM weight and bias to zero, allow any scalar input sequence and any initial values $h_0,c_0\in \mathbb{R}$. Then for every $t\ge 1$, the input, forget, and output gates equal $1\mathrm{/}2$, the candidate update is zero, and

Standard LSTM cell

Goodfellow, Bengio, and Courville, Deep Learning, Chapter 10, §10.10.1

Take the scalar state-space model $A=0$, $B=C=1$, direct term $D=0$, step size $\mathrm{\Delta }=1$, and initial state $x_{-1}=0$. For any integer $L\ge 1$ and any length-$L$ scalar input sequence $(u_0,\dots ,u_{L-1})$, the bilinear matrices are $\bar{A}=1$ and $\bar{B}=1$, so the recurrence gives $x_t=x_{t-1}+u_t$ and $y_t={\sum }_{s=0}^t{u}_s$ for $0\le t<L$. The length-$L$ impulse kernel is $(1,\dots ,1)$, and its causal convolution ${\sum }_{j=0}^t{K}_j{u}_{t-j}$ gives the same outputs for $0\le t<L$.

DPLR structured state-space parameterization

S4 state-space recurrence and convolution procedure

Gu, Goel, and Ré, “Efficiently Modeling Long Sequences with Structured State Spaces,” §§2.3–2.4

Assume $N,r,L\ge 1$ with fixed $r$, $A=\mathrm{\Lambda }-P{Q}^{\ast }$ for diagonal $\mathrm{\Lambda }\in {\mathbb{C}}^{N\times N}$ and $P,Q\in {\mathbb{C}}^{N\times r}$, $B,C\in {\mathbb{C}}^N$, and $\mathrm{\Delta }>0$. Let $G=I-\mathrm{\Delta }A\mathrm{/}2$ be invertible, let $M$ be the least power of two at least $L$, set $\bar{A}=G^{-1}(I+\mathrm{\Delta }A\mathrm{/}2)$ and $\bar{B}=G^{-1}\mathrm{\Delta }B$, and supply ${\tilde{C}}_M=(I-{\bar{A}}^M{)}^{\ast }C$. Assume the root-grid diagonal resolvents $R(z)=(a(1-z)\mathrm{/}(1+z)I-\mathrm{\Lambda }{)}^{-1}$ and corrections $I_r+Q^{\ast }R(z)P$ are invertible for every $M$th root $z\mathrm{\ne }-1$, where $a=2\mathrm{/}\mathrm{\Delta }$; the value at $-1$, when that root occurs, is handled by its finite formula. In the exact complex arithmetic model, each scalar addition, multiplication, division, and zero test has unit cost, and the filter cost counts evaluation from the supplied ${\tilde{C}}_M$, not its formation. A nonzero direct input-to-output term is separate from the kernel.

For numerical evaluation, at each root $z\mathrm{\ne }-1$ put $\beta (z)=(1+z)\mathrm{/}2$ and $\mathcal{D}(z)=2R(z)\mathrm{/}(1+z)$, and define $x={\tilde{C}}_M^{\ast }\mathcal{D}B$, $u={\tilde{C}}_M^{\ast }\mathcal{D}P$, $v=Q^{\ast }\mathcal{D}B$, $T=Q^{\ast }\mathcal{D}P$, and $W=I_r+\beta T$. Matrix infinity norm is maximum absolute row sum, and for a column it is maximum absolute entry. Suppose certified constants $b,w,s\ge 1$ bound $\mathrm{\mid }x\mathrm{\mid },\mathrm{\parallel }u{\mathrm{\parallel }}_{\mathrm{\infty }},\mathrm{\parallel }v{\mathrm{\parallel }}_{\mathrm{\infty }},\mathrm{\parallel }T{\mathrm{\parallel }}_{\mathrm{\infty }}$ by $b$ and $\mathrm{\parallel }{W}^{-1}{\mathrm{\parallel }}_{\mathrm{\infty }}$ by $w$ over these roots. For every scalar Cauchy sum used for a scaled block, supply a certified pole/root separation $d>0$ and a weight bound $A_{\mathrm{w}\mathrm{t}}\ge {\sum }_j\mathrm{\mid }{a}_j\mathrm{\mid }$ with $A_{\mathrm{w}\mathrm{t}}\mathrm{/}d\le s$; also supply the exact index partition ${\lambda }_j=-a$ or ${\lambda }_j\mathrm{\ne }-a$ so the exceptional constant terms are selected, and form the finite constants and the $-1$ value separately. Empty Cauchy sums need no separation certificate. Put $\kappa =64r{w}^2(b+1{)}^{2}s$.

For $0<\epsilon \le 1\mathrm{/}2$, assume variable-precision scalar arithmetic with unrestricted exponent range, certified input, root and arc-center approximations, certified enclosures sufficient for the geometric routing decisions, and the finite rational-circuit precision certificate. Explicitly, for the selected finite circuit let $T_0$ bound its arithmetic gates, $V_0\ge 1$ bound exact intermediate moduli, and $0<D_0\le 1$ bound below every exact divisor modulus; with $K_0=64(V_0+1{)}^2\mathrm{/}{D}_0^2$, choose $u\le \min (1,D_0\mathrm{/}2,\epsilon \mathrm{/}2)\mathrm{/}(2K_0^{T_0+1})$. Each supplied input/constant approximation and each scalar arithmetic operation has absolute error at most $u$; each scalar operation and word access has unit cost. Precision and word lengths may depend on the input separation, correction conditioning, intermediate magnitudes and target accuracy. The operation and word counts exclude bit cost, computing the supplied certificates, and forming ${\tilde{C}}_M$.

For the stated DPLR data, the length-$L$ kernel is $K_L=(C^{\ast }{\bar{A}}^j\bar{B}{)}_{j=0}^{L-1}$. From the supplied ${\tilde{C}}_M$, the exact filter procedure evaluates it using at most $(r+1{)}^2$ roots-of-unity Cauchy products, $O((N+L){\log }^2(2+N+L))$ exact complex arithmetic operations, and $O(N+L)$ working space. Under the numerical hypotheses above, for $0<\epsilon \le 1\mathrm{/}2$ it returns ${\widehat{K}}_L$ with $\mathrm{\parallel }{\widehat{K}}_L-K_L{\mathrm{\parallel }}_{\mathrm{\infty }}\le \epsilon$ using $O((N+L)\log (2+N+L)\log (\kappa \mathrm{/}\epsilon ))$ scalar operations and $O(N+L)$ scalar words; bounded-$\kappa$ families have the factor $\log (1\mathrm{/}\epsilon )$. These evaluation counts exclude forming ${\tilde{C}}_M$, computing the supplied certificates, and bit cost. The filter procedure does not compute parameter gradients or optimizer updates, so its count is not a bound for a complete training step. Neither the exact arithmetic count nor the conditional numerical guarantee supplies a fixed-word bound or an empirical long-context accuracy guarantee.

S4 convolution-filter arithmetic and numerical cost

S4 state-space recurrence and convolution procedure

Gu, Goel, and Ré, “Efficiently Modeling Long Sequences with Structured State Spaces,” §3.3 and Appendix C.3