S4 state-space recurrence and convolution procedure

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