Recurrent, LSTM, and structured state-space models

Fix positive finite dimensions $d_x,d_h,d_y$, a finite horizon $T\in \mathbb{N}$, a parameter set $\mathrm{\Theta }$, an initial state $h_0\in {\mathbb{R}}^{d_h}$, and inputs $x_1,\dots ,x_T\in {\mathbb{R}}^{d_x}$. A shared-parameter recurrent neural network is specified by a state map $F:{\mathbb{R}}^{d_h}\times {\mathbb{R}}^{d_x}\times \mathrm{\Theta }\to {\mathbb{R}}^{d_h}$ and, optionally, an output map $G:{\mathbb{R}}^{d_h}\times {\mathbb{R}}^{d_x}\times \mathrm{\Theta }\to {\mathbb{R}}^{d_y}$. For a fixed $\theta \in \mathrm{\Theta }$, it computes $h_t=F(h_{t-1},x_t;\theta )$ and, when $G$ is supplied, $y_t=G(h_t,x_t;\theta )$ for $1\le t\le T$, using the same $\theta$ at every time. Its unrolling is the finite directed acyclic computation graph whose vertices are the input, state, output, and map-evaluation occurrences and whose edges record which values each evaluation uses. The folded schematic may show a feedback edge, while the time-indexed unrolled graph is acyclic.

For horizon $T=2$, the following TikZ graph shows the dependency edges in the unrolling. The boxes $F_1,F_2$ are two time-indexed occurrences of the same map $F$; the dashed edges indicate that both occurrences receive the identical parameter $\theta$. The graph omits the optional output map $G$.

\begin{center} \begin{tikzpicture}[>=stealth, every node/.style={draw, rounded corners, inner sep=4pt}] \node (h0) at (0,0) {$h_0$}; \node (x1) at (0,-1.1) {$x_1$}; \node (f1) at (2,-0.55) {$F_1$}; \node (h1) at (4,-0.55) {$h_1$}; \node (x2) at (4,-1.65) {$x_2$}; \node (f2) at (6,-1.1) {$F_2$}; \node (h2) at (8,-1.1) {$h_2$}; \node (theta) at (4,1) {$\theta$}; \draw[->] (h0) -- (f1); \draw[->] (x1) -- (f1); \draw[->] (f1) -- (h1); \draw[->] (h1) -- (f2); \draw[->] (x2) -- (f2); \draw[->] (f2) -- (h2); \draw[->, dashed] (theta) -- (f1); \draw[->, dashed] (theta) -- (f2); \end{tikzpicture} \end{center}

Goodfellow, Bengio, and Courville, Deep Learning, Chapter 10, §§10.1–10.2

Assume the Axiom of Countable Choice, so the usual Lebesgue-measurable sets on $\mathbb{R}$ form a complete measure space with Lebesgue measure $\lambda$. Every bounded interval is Lebesgue measurable, and any interval is the countable union of its bounded intersections with $[-n,n]$; since the Lebesgue-measurable sets form a sigma-algebra, $I$ is measurable and $\lambda$ restricts to $I$. Let $\mathbb{F}$ be $\mathbb{R}$ or $\mathbb{C}$, let $I\subseteq \mathbb{R}$ be an interval with $t_0\in I$, and let $N,m,p\ge 1$ be finite. A continuous-time linear state-space model consists of matrices $A\in {\mathbb{F}}^{N\times N}$, $B\in {\mathbb{F}}^{N\times m}$, $C\in {\mathbb{F}}^{p\times N}$, $D\in {\mathbb{F}}^{p\times m}$, an initial value $x_0\in {\mathbb{F}}^N$, and an input $u:I\to {\mathbb{F}}^m$ whose real and imaginary coordinates are Lebesgue measurable and integrable on every compact subinterval of $I$. A state trajectory is a function $x:I\to {\mathbb{F}}^N$ whose real and imaginary coordinates are absolutely continuous on every compact subinterval of $I$ (with absolute continuity defined by the $\epsilon$–$\delta$ condition on finite families of disjoint subintervals), with $x(t_0)=x_0$ and $x^{\mathrm{\prime }}(t)=Ax(t)+Bu(t)$ outside a Lebesgue-null subset of $I$. Its output is $y(t)=Cx(t)+Du(t)$. The term $Du(t)$ is the direct input-to-output term and is separate from the state contribution $Cx(t)$.

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

Integrable real and complex functions, and their integrals

Absolute continuity on a compact interval

Lebesgue measurable sets, the family $\mathcal{L}({\mathbb{R}}^n)$, and the restricted set function ${\lambda }_n$

Assuming countable choice, $\mathcal{L}({\mathbb{R}}^n)$ is a sigma-algebra containing every elementary set and ${\lambda }_n$ is a complete measure extending elementary volume

A box in ${\mathbb{R}}^n$ with parameters $a_i\le b_i$ is Lebesgue measurable of measure ${\prod }_{i<n}(b_i-a_i)$, whichever of its faces are included

The Axiom of Countable Choice (${\mathrm{A}\mathrm{C}}_{\omega }$)

For a nonnegative integer $q$, let $M=2^q$, set $\omega =1$ when $M=1$ and $\omega =e^{-2\pi i\mathrm{/}M}$ when $M>1$, and let $a=(a_0,\dots ,a_{M-1})\in {\mathbb{C}}^M$. The forward discrete Fourier transform (DFT) is the vector $\widehat{a}\in {\mathbb{C}}^M$ with ${\widehat{a}}_k={\sum }_{j=0}^{M-1}{a}_j{\omega }^{jk}$. The radix-two procedure returns this vector by the following exact recursion. If $M=1$, return $(a_0)$. If $M>1$, recursively compute the length-$M\mathrm{/}2$ DFTs $E$ and $O$ of $(a_0,a_2,\dots ,a_{M-2})$ and $(a_1,a_3,\dots ,a_{M-1})$, using the primitive root ${\omega }^2$. For $0\le k<M\mathrm{/}2$, return ${\widehat{a}}_k=E_k+{\omega }^k{O}_k$ and ${\widehat{a}}_{k+M\mathrm{/}2}=E_k-{\omega }^k{O}_k$. The inverse transform is obtained by running the same procedure with $\bar{\omega }$ and dividing every output by $M$.

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

Let $D\in {\mathbb{C}}^{N\times N}$ be invertible, let $U,V\in {\mathbb{C}}^{N\times r}$, and set $W=I_r+V^{\ast }{D}^{-1}U$, where ${}^{\ast }$ denotes conjugate transpose. Then $D+U{V}^{\ast }$ is invertible if and only if $W$ is invertible. When these matrices are invertible,

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

Fix positive finite dimensions $d_x,d_h$, inputs $x_t\in {\mathbb{R}}^{d_x}$, hidden states $h_t\in {\mathbb{R}}^{d_h}$, cell states $c_t\in {\mathbb{R}}^{d_h}$, and initial values $h_0,c_0$. A standard long short-term memory (LSTM) cell without peephole connections or projection has shared matrices $W_i,W_f,W_o,W_c\in {\mathbb{R}}^{d_h\times d_x}$, $U_i,U_f,U_o,U_c\in {\mathbb{R}}^{d_h\times d_h}$, and biases $b_i,b_f,b_o,b_c\in {\mathbb{R}}^{d_h}$. Apply $\sigma (s)=(1+e^{-s}{)}^{-1}$ and $\tanh (s)=(e^s-e^{-s})\mathrm{/}(e^s+e^{-s})$ coordinatewise. At each time $t\ge 1$, set

$i_t=\sigma (W_i{x}_t+U_i{h}_{t-1}+b_i),f_t=\sigma (W_f{x}_t+U_f{h}_{t-1}+b_f),o_t=\sigma (W_o{x}_t+U_o{h}_{t-1}+b_o),$

$g_t=\tanh (W_c{x}_t+U_c{h}_{t-1}+b_c),c_t=f_t\odot c_{t-1}+i_t\odot g_t,h_t=o_t\odot \tanh (c_t),$

where $\odot$ is coordinatewise multiplication. The same matrices and biases are reused at every time step. The vectors $i_t,f_t,o_t$ are the input, forget, and output gates, and $g_t$ is the candidate state update.

Goodfellow, Bengio, and Courville, Deep Learning, Chapter 10, §10.10.1, Eqs. 10.40–10.44

An $N$-state single-input, single-output complex linear state-space model has $B,C\in {\mathbb{C}}^N$, state equation $x^{\mathrm{\prime }}(t)=Ax(t)+Bu(t)$, and output $y(t)=C^{\ast }x(t)$; its direct term is set to zero in the S4 convolution convention. Here $X^{\ast }$ denotes conjugate transpose; a square complex matrix $H$ is normal when $H^{\ast }H=H{H}^{\ast }$; and a square matrix $V$ is unitary when $V^{\ast }V=V{V}^{\ast }=I$. A matrix is in diagonal-plus-low-rank (DPLR) form of rank $r$ if

where $\mathrm{\Lambda }\in {\mathbb{C}}^{N\times N}$ is diagonal and $P,Q\in {\mathbb{C}}^{N\times r}$. It is in normal-plus-low-rank (NPLR) form of rank $r$ if

where $V\in {\mathbb{C}}^{N\times N}$ is unitary, $\mathrm{\Lambda }$ is diagonal, and $P,Q\in {\mathbb{C}}^{N\times r}$. The unitary state change $\tilde{x}=V^{\ast }x$ converts this NPLR representation to DPLR form and changes $B,C$ to $V^{\ast }B,V^{\ast }C$, respectively.

For a step size $\mathrm{\Delta }>0$ such that $I-\mathrm{\Delta }A\mathrm{/}2$ is invertible, the bilinear recurrent matrices are $\bar{A}=(I-\mathrm{\Delta }A\mathrm{/}2{)}^{-1}(I+\mathrm{\Delta }A\mathrm{/}2),\bar{B}=(I-\mathrm{\Delta }A\mathrm{/}2{)}^{-1}\mathrm{\Delta }B,\bar{C}=C\mathrm{.}$ With initial state $x_{-1}=0$, the discrete recurrence is $x_t=\bar{A}{x}_{t-1}+\bar{B}{u}_t$, $y_t={\bar{C}}^{\ast }{x}_t$. Its length-$L$ impulse kernel is the vector $K_L=({\bar{C}}^{\ast }{\bar{A}}^j\bar{B}{)}_{j=0}^{L-1}$. For a nonzero direct term $D$, the output has the additional separate skip term $D{u}_t$.

For a filter evaluation of transform length $M\ge L$, the output vector for the length-$M$ truncated generating function is ${\tilde{C}}_M=(I-{\bar{A}}^M{)}^{\ast }C$. A filter procedure may receive ${\tilde{C}}_M$ as given data; an operation count for evaluating the filter from this transformed vector does not include the separate work of forming it from $C$.

Continuous-time linear state-space model

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

Let $M=2^q$ for an integer $q\ge 0$, and let $\omega =e^{-2\pi i\mathrm{/}M}$ (with $\omega =1$ when $M=1$). Define the forward DFT matrix $F\in {\mathbb{C}}^{M\times M}$ by $F_{kj}={\omega }^{jk}$, where ${}^{\ast }$ denotes conjugate transpose. Then $F^{\ast }F=MI$, so the inverse DFT is $M^{-1}{F}^{\ast }$. The radix-two procedure defined on this page computes either transform exactly in $O(M\log (2+M))$ complex arithmetic operations and $O(M)$ space. Consequently, for every integer $n\ge 1$, two polynomials of degree less than $n$ can be multiplied exactly in $O(n\log (2+n))$ complex arithmetic operations and $O(n)$ space by zero-padding to a radix-two transform length at least $2n-1$.

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

Pan, “Fast Approximate Computations with Cauchy Matrices and Polynomials,” §3.3

Radix-two discrete Fourier transform

The proofs use only finite sums and arithmetic in $\mathbb{C}$; no choice principle is needed.

Input an integer $M=2^q$ with $q\ge 0$, a primitive $M$th root of unity $\omega$, a finite list of poles ${\mu }_1,\dots ,{\mu }_n\in \mathbb{C}$, and weights $a_1,\dots ,a_n\in \mathbb{C}$. The required output is $y_k={\sum }_{j=1}^n{a}_j\mathrm{/}({\omega }^k-{\mu }_j)$ for $0\le k<M$, provided no denominator is zero. If $n=0$, return the zero vector.

Here a primitive $M$th root means ${\omega }^M=1$ and ${\omega }^k\mathrm{\ne }1$ for $1\le k<M$. For $n\ge 1$, form a balanced binary tree on the pole indices, splitting each consecutive index interval into two intervals whose sizes differ by at most one. At a leaf containing index $j$, store polynomials $D_j(z)=z-{\mu }_j$ and $P_j(z)=a_j$. At an internal node with disjoint child index sets $I_1,I_2$, store $D_I(z)=D_{I_1}(z)D_{I_2}(z),P_I(z)=P_{I_1}(z)D_{I_2}(z)+D_{I_1}(z)P_{I_2}(z)\mathrm{.}$ At the root, fold the coefficients of each polynomial modulo $z^M-1$: the coefficient at degree $r$ is replaced by the sum of coefficients at degrees congruent to $r$ modulo $M$. Let ${\omega }_0=e^{-2\pi i\mathrm{/}M}$, which equals $1$ when $M=1$. Apply the radix-two DFT procedure to each folded coefficient vector to evaluate $P_I$ and $D_I$ at ${\omega }_0^j$, $0\le j<M$. If $M>1$, find the unique odd integer $s\in \{1,3,\dots ,M-1\}$ with $\omega ={\omega }_0^s$ by generating and comparing the $M$ powers of ${\omega }_0$, then place the value at index $sk<mrow><mi\; mathvariant="normal">m</mi><mi\; mathvariant="normal">o</mi><mi\; mathvariant="normal">d</mi></mrow>M$ at output index $k$; this is the supplied root order because ${\omega }^k={\omega }_0^{sk}$. For $M=1$, the only root is $1$ and no reindexing is needed. This root search and reindexing use $O(M)$ operations. If some reindexed $D_I({\omega }^k)=0$, return status collision; otherwise return $y_k=P_I({\omega }^k)\mathrm{/}{D}_I({\omega }^k)$.

Pan, “Fast Approximate Computations with Cauchy Matrices and Polynomials,” §§3.2–3.3

Radix-two discrete Fourier transform

Index rows and columns by $0,\dots ,N-1$. Define the real LegS coefficient matrix $G^S$ by $G_{nk}^S=-\sqrt{(2n+1)(2k+1)}$ when $n>k$, $G_{nn}^S=-(n+1)$, and $G_{nk}^S=0$ when $n<k$. In the LegS coefficient ODE for $t>0$ this matrix is multiplied by the time-dependent scalar $1\mathrm{/}t$; the matrix assertion below does not make that ODE time-invariant. For $\theta >0$, define the LegT state generator by $G_{nk}^T=-(2n+1)(-1{)}^{n-k}\mathrm{/}\theta$ when $n\ge k$ and $G_{nk}^T=-(2n+1)\mathrm{/}\theta$ when $n<k$. Define the original HiPPO LagT state generator by $G_{nk}^L=-1$ when $n\ge k$ and $G_{nk}^L=0$ when $n<k$. Here $X^{\ast }$ is conjugate transpose, a complex matrix $H$ is normal when $H^{\ast }H=H{H}^{\ast }$, and a square matrix $V$ is unitary when $V^{\ast }V=V{V}^{\ast }=I$. For LegS and LagT use the identity similarity; for LegT use the invertible diagonal similarity $S_0=\mathrm{diag}(\sqrt{2n+1}{)}_{n=0}^{N-1}\mathrm{diag}((-1{)}^n{)}_{n=0}^{N-1}$. Each transformed matrix is a difference $H-U{V}^{\ast }$ with $H$ normal and $\mathrm{rank}(U{V}^{\ast })\le 1$ for LegS and LagT and $\mathrm{rank}(U{V}^{\ast })\le 2$ for LegT. Consequently the equivalent finite-dimensional complex state-space matrix has a normal-plus-low-rank representation, and a unitary change of state coordinates gives diagonal-plus-low-rank form. Under any state change $x=S\tilde{x}$, the input and output vectors change to $S^{-1}B$ and $S^{\ast }C$, respectively.

Gu, Goel, and Ré, “Efficiently Modeling Long Sequences with Structured State Spaces,” Theorem 1 and Appendix C.1

Gu et al., “HiPPO: Recurrent Memory with Optimal Polynomial Projections,” §2.3 and Appendices D.1–D.3

Let $N,r\ge 1$ and $A=\mathrm{\Lambda }-P{Q}^{\ast }$ with $\mathrm{\Lambda }\in {\mathbb{C}}^{N\times N}$ diagonal and $P,Q\in {\mathbb{C}}^{N\times r}$, where ${}^{\ast }$ denotes conjugate transpose; let $B,C\in {\mathbb{C}}^N$, let $\mathrm{\Delta }>0$, and put $\alpha =\mathrm{\Delta }\mathrm{/}2$. The bilinear step is defined only when $G=I-\alpha A$ is invertible. Given an input sequence $(u_t)$ and state $x_{-1}$, its recurrent mode computes $b_t=(I+\alpha A)x_{t-1}+\mathrm{\Delta }B{u}_t,G{x}_t=b_t,y_t=C^{\ast }{x}_t\mathrm{.}$ For a zero initial state, the equivalent matrices and length-$L$ kernel are $\bar{A}=G^{-1}(I+\alpha A)$, $\bar{B}=G^{-1}\mathrm{\Delta }B$, and $K_j=C^{\ast }{\bar{A}}^j\bar{B}$ for $0\le j<L$; a direct term $D{u}_t$, if present, is added separately.

For the filter mode, fix an integer $L\ge 1$, let $M$ be the least power of two with $M\ge L$, let $\omega =e^{-2\pi i\mathrm{/}M}$, and let $\tilde{C}\in {\mathbb{C}}^N$ be supplied as the length-$M$ truncated-output parameter. When a conventional output vector $C$ is used, this parameter is $\tilde{C}=(I-{\bar{A}}^M{)}^{\ast }C$. Define the $M$ root values $z_k={\omega }^k$. The algorithm returns the first $L$ coefficients of the inverse radix-two transform of the values $H(z_k)$ below.

For $z_k=-1$, set $H(z_k)=(\mathrm{\Delta }\mathrm{/}2){\tilde{C}}^{\ast }B$. For each $z=z_k\mathrm{\ne }-1$, put $s(z)=(2\mathrm{/}\mathrm{\Delta })(1-z)\mathrm{/}(1+z)$. If $s(z)-{\lambda }_j=0$ for some diagonal entry ${\lambda }_j$ of $\mathrm{\Lambda }$, return status resonant-diagonal. If no diagonal resonance is present, write $R(z)=\mathrm{diag}((s(z)-{\lambda }_j{)}^{-1})$ and $W(z)=I_r+Q^{\ast }R(z)P$. If $W(z)$ is singular, return status resonant-correction. When this correction is invertible, the target root value is $H(z)=\frac{2}{1+z}({\tilde{C}}^{\ast }R(z)B-{\tilde{C}}^{\ast }R(z)PW(z{)}^{-1}{Q}^{\ast }R(z)B)\mathrm{.}$ The preceding resolvent expression specifies the target root values. The filter evaluation computes them using $1+2r+r^2=(r+1{)}^2$ scalar Cauchy products as follows. Put $a=2\mathrm{/}\mathrm{\Delta }$, and for each $j$ with ${\lambda }_j\mathrm{\ne }-a$ define ${\mu }_j=(a-{\lambda }_j)\mathrm{/}(a+{\lambda }_j)$. For any pair of vectors $u,v\in {\mathbb{C}}^N$, compute on the root grid $S_{uv}(z)=\frac{2}{1+z}{u}^{\ast }R(z)v=\frac{\mathrm{\Delta }}{2}{\sum }_{j:{\lambda }_j=-a}\overline{u_j}{v}_j-2{\sum }_{j:{\lambda }_j\mathrm{\ne }-a}\frac{\overline{u_j}{v}_j}{a+{\lambda }_j}\frac{1}{z-{\mu }_j}$ by calling the roots-of-unity Cauchy procedure with poles ${\mu }_j$ and weights $-2\overline{u_j}{v}_j\mathrm{/}(a+{\lambda }_j)$, then adding the displayed constant. A collision ${\mu }_j=z_k$ returns status resonant-diagonal. For $z=z_k\mathrm{\ne }-1$, set $W(z)=I_r+\frac{1+z}{2}{S}_{Q,P}(z),H(z)=S_{\tilde{C},B}(z)-\frac{1+z}{2}{S}_{\tilde{C},P}(z)W(z{)}^{-1}{S}_{Q,B}(z)\mathrm{.}$ If $W(z)$ is singular, return status resonant-correction. The four families $S_{\tilde{C},B}$, $S_{\tilde{C},P}$, $S_{Q,B}$, and $S_{Q,P}$ contain $1$, $r$, $r$, and $r^2$ scalar products, respectively. The root $-1$ is handled directly and is never included in this Cauchy evaluation. Matrix products and inversions of the $r\times r$ corrections are performed separately at each root.

The recurrent solver may exploit the diagonal-plus-low-rank structure as follows. Write $d_j=1-\alpha {\lambda }_j$, let $J=\{j:d_j\mathrm{\ne }0\}$ and $Z=\{j:d_j=0\}$. If $Z$ is empty, use the Woodbury inverse with $D=\mathrm{diag}(d_j)$ and correction $I_r+\alpha Q^{\ast }{D}^{-1}P$; if this correction is singular, return status singular-bilinear-matrix. If $s=\mathrm{\mid }Z\mathrm{\mid }>0$, then when $s\le r$ solve the bordered system $\left[\begin{array}{cc}{I}_r+\alpha Q_J^{\ast }{D}_J^{-1}{P}_J& -Q_Z^{\ast }\\ \alpha P_Z& 0\end{array}\right]\left[\begin{array}{c}z\\ x_Z\end{array}\right]=\left[\begin{array}{c}{Q}_J^{\ast }{D}_J^{-1}{b}_J\\ b_Z\end{array}\right],x_J=D_J^{-1}(b_J-\alpha P_{J}z),$ where $D_J=\mathrm{diag}(d_j{)}_{j\in J}$; if $s>r$ or the bordered system is singular, return status singular-bilinear-matrix. Here $z=Q^{\ast }x$. This branch solves $Gx=b$ without assuming the diagonal base is invertible.

Gu, Goel, and Ré, “Efficiently Modeling Long Sequences with Structured State Spaces,” Algorithm 1 and Appendix C

DPLR structured state-space parameterization

Radix-two discrete Fourier transform

Exact Cauchy matrix-vector product on a radix-two root grid

Woodbury inverse for a diagonal-plus-low-rank update

Let $n\ge 1$, $M=2^q\ge 1$, $\omega$ be a primitive $M$th root of unity, and let ${\mu }_1,\dots ,{\mu }_n\in \mathbb{C}$ and $a_1,\dots ,a_n\in \mathbb{C}$ satisfy ${\mu }_j\mathrm{\ne }{\omega }^k$ for every $j$ and $0\le k<M$. Then the Cauchy matrix-vector product $y_k={\sum }_{j=1}^n{a}_j\mathrm{/}({\omega }^k-{\mu }_j)$, $0\le k<M$, is computed exactly by the balanced product-tree procedure using $O((n+M){\log }^2(2+n+M))$ complex arithmetic operations and $O(n+M)$ space. Repeated poles are allowed.

Pan, “Fast Approximate Computations with Cauchy Matrices and Polynomials,” §3.3

Exact Cauchy matrix-vector product on a radix-two root grid

Roots-of-unity orthogonality and radix-two FFT cost