Model Compression Method for S4 with Diagonal State Space Layers using Balanced Truncation

A crucial component of a deep learning model discussed in our study, as outlined in Section III, is the hidden layer known as the DSS layer. This layer is defined by using the SSM

where $u(t)\in\mathbb{R}$, $y(t)\in\mathbb{C}$, and $x(t)\in\mathbb{C}^{N}$ denote the input, output, and state, respectively, and $(A,B,C)\in\mathbb{C}^{N\times N}\times\mathbb{C}^{N\times 1}\times\mathbb{C}^{1\times N}$. The matrix $A$ denotes a state transition matrix, which describes the internal influence on the time evolution of the internal state $x$. The matrix $B$ is an input matrix, which describes how the external input $u$ affects the internal state $x$. The matrix $C$ serves as an output matrix, which describes how the internal state $x$ is transformed into the observable output $y$.

Notably, SSM (1) is a single-input and single-output system, and the matrices $A$, $B$, $C$, state $x(t)$, and output $y(t)$ are all complex-valued. Moreover, the state dimension $N$ is relatively large, unlike the standard setting in control theory [38].

To address computational issues arising from the large state dimension of SSM (1), we consider using the balanced truncation method [33], a reduction technique. This method focuses on the controllability and observability of SSM (1) to derive another SSM with dimension $r~{}(\leq N)$, which gives almost the same output as (1):

where $A_{r}\in\mathbb{C}^{r\times r},B_{r}\in\mathbb{C}^{r\times 1},C_{r}\in\mathbb{C}^{1\times r}$. Below is a brief explanation of the balanced truncation method. For more detailed information, refer to Appendix A.

For SSM (1) with asymptotic stability, where all the real parts of the eigenvalues of matrix $A$ are negative, the controllability Gramian $P$ and observability Gramian $Q$ are defined as

These are the unique solutions to the Lyapunov equations

as shown in [38, Theorem 4.1].

In the balanced truncation method, the subspace spanned by eigenvectors corresponding to small eigenvalues of the controllability Gramian $P$ and the observability Gramian $Q$ is ignored. The minimum input energy required to achieve $\lim_{t\to\infty}x(t)=x_{f}$ from the initial condition $x(0)=0$ is expressed using the controllability Gramian $P$ as

Thus, eigenvectors corresponding to smaller eigenvalues of $P$ correspond to directions in the state space that are less influenced by the input $u$. On the other hand, when $x(0)=x_{0}$ and $u(0)=0$, the output energy can be expressed using the observability Gramian $Q$ as

Thus, eigenvectors corresponding to smaller eigenvalues of $Q$ correspond to directions in the state space that have less impact on the output $y$.

A coordinate transformation $\bar{x}(t)=Tx(t)$ is applied to SSM (1), obtaining another SSM where the controllability Gramian and observability Gramian coincide as a diagonal matrix $\Sigma$:

SSM (9) is referred to as the balanced realization of SSM (1). Here, the diagonal elements of $\Sigma$ are denoted as $\sigma_{1},\cdots,\sigma_{N}$, which are called the Hankel singular values, satisfying $\sigma_{1}\geq\cdots\geq\sigma_{N}>0$ under the assumption that SSM (1) is controllable and observable. By partitioning the matrices $TAT^{-1}$, $TB$, $CT^{-1}$ in SSM (9) as

we define the parameters $(A_{r},B_{r},C_{r})$ of reduced model (2) as

The resulting system (2) with (13) can be interpreted as a reduced SSM of SSM (1), obtained by truncating the state space associated with the smaller Hankel singular values $\sigma_{r+1},\cdots,\sigma_{N}$, which correspond to the subspace spanned by eigenvectors that are less influenced by the input or have less influence on the output. Moreover, if SSM (1) is asymptotically stable, the reduced SSM (2) with (13) is also asymptotically stable, as shown in [38, Proposition 4.15].

For the model explained in Section III, the parameters of the matrices $(A,B,C)$ in the SSM (1) are trained using a suitable optimization algorithm. The initialization of matrix $A$ significantly influences the performance of trained models, as it sets the initial state for the optimization process. The High-order Polynomial Projection Operators (HiPPO) matrix [7] is derived from a method for online compression of continuous signals using projections onto subspaces spanned by polynomial bases. It is well-established that the HiPPO matrix is an effective choice for the initial $A$ [6, 12, 9].

The derivation of the HiPPO matrix is explained below. With respect to a measure $\mu$ on $[0,\infty)$, let

The inner product and norm on $L_{2}(\mu)$ are defined as

respectively.

For an input signal $u(t)$ defined on $t\geq 0$, the history $u_{\leq t}\coloneqq u(\tau)\mid_{\tau\leq t}$ at each time $t>0$ is approximated by projecting it onto a subspace spanned by polynomial bases, and the corresponding coefficient vector $x(t)$ represents the history of input signal. This compression is useful because storing $u_{\leq t}$ requires a significant amount of memory. Thus, at each time $t$, $x(t)$ contains sufficient information to reconstruct $u_{\leq t}$, even though it requires less memory compared to directly storing $u_{\leq t}$.

The vector $x(t)$ is expressed as the optimal solution to a convex optimization problem, which is defined by a measure $\mu^{(t)}$ on $(-\infty,t]$ and orthogonal polynomial basis of the subspace of $L_{2}(\mu^{(t)})$ denoted as $\{g_{n}^{(t)}\}_{1\leq n\leq N}$ (i.e. $\mathrm{deg}(g_{i}^{(t)})=i~{}(i=1,\cdots,N),\langle g_{i}^{(t)},g_{j}^{(t)}\rangle_{\mu^{(t)}}=0~{}(i\neq j)$). The optimization problem is

If $\{g_{n}^{(t)}\}_{1\leq n\leq N}$ is a normalized orthogonal basis, i.e. $\lVert g_{i}^{(t)}\rVert_{L_{2}(\mu^{(t)})}=1~{}(i=1,\cdots,N)$, the optimal solution is given by

The vector $x(t)$ defines the approximation $g=\sum_{k=1}^{N}x_{k}(t)g^{(t)}_{k}$ of $u_{\leq t}$, thus it retains information necessary for reconstructing the history $u_{\leq t}$ of the input $u(t)$ at time $t$. This property of memorizing the input history in the state vector will be useful for modeling long sequential data, as capturing dependencies in sequential data requires referencing information from previous inputs to compute each output at every time step. Moreover, by Equation (18), the measure $\mu^{(t)}$ represents the importance of each time step when compressing the history $u_{\leq t}$. This $x(t)$ satisfies the differential equation

where $A(t)\in\mathbb{R}^{N\times N},B(t)\in\mathbb{R}^{N\times 1}$ depend on the polynomial basis $\{g^{(t)}_{n}\}_{1\leq n\leq N}$ and the measure $\mu^{(t)}$. Unlike SSM (1) introduced in Subsection II-A, $A(t)$ and $B(t)$ are time-dependent.

For HiPPO-LegS that is a variant of HiPPO [7], the measure $\mu^{(t)}$ is defined as the scaled Legendre (LegS) measure

This assigns uniform importance to the entire history $[0,t]$ at each time $t$. Furthermore, the polynomial basis $g^{(t)}_{n}(\tau)$ is the normalized orthogonal basis

where $P_{n}(\alpha)$ is the Legendre polynomial

In this case, as shown in [7, Appendix D.3], $x(t)$ satisfies

This matrix $A$ of Equation (24) is called the HiPPO matrix. For the SSM incorporating the HiPPO matrix, the state vector $x(t)$ retains information about the history $u_{\leq t}$ of the input $u$ at each time $t$ [7].

The most important component of the deep learning model illustrated in Fig. 2 is the DSS layer. As shown in Fig. 3, the DSS layer consists of:

• •

  Independent $H$ DSS models.

• •

  Nonlinear connection blocks.

• •

  Linear combination block.

The details of each of these components are explained below.

By restricting the matrix $A\in\mathbb{C}^{N\times N}$ in (1) to be diagonal, assuming that the diagonal elements do not lie on the imaginary axis, the DSS model is defined by the following discretization for a sample time $\Delta\in\mathbb{R}_{>0}$:

where $\bar{A}=e^{A\Delta}$, $\bar{B}=(\bar{A}-I)A^{-1}B$, and $\bar{C}=C$. The diagonal elements of the matrix $\bar{A}$ do not lie on the unit circle in the complex plane, due to the assumption on the matrix $A$.

The nonlinear connection block receives the input $u_{k}$ and output $y_{k}$ from the DSS model and outputs 1-dimensional sequential data

where $D\in\mathbb{R}$, and GELU[39] is a nonlinear activation function expressed as

Here, $\Phi(\alpha)$ is the cumulative distribution function of the standard normal distribution. This approach is expected to enhance the performance of the model. Note that $\mathrm{Re}(y_{k})$ is used in Equation (27) since $y_{k}$ can be a complex number.

Finally, the $H$ 1-dimensional sequential data $(y_{k}^{\prime(h)})_{1\leq h\leq H}$ outputted from each nonlinear connection block is mixed to obtain the final output of the DSS layer, resulting in $H$ 1-dimensional sequential data $(u_{k}^{\prime(h)})_{1\leq h\leq H}$. With parameters of weight $W_{out}\in\mathbb{R}^{H\times H}$ and bias $b\in\mathbb{R}^{H}$, the output is expressed as

where $\bm{1}$ is a 1-dimensional sequential data of the same length as $y_{k}^{\prime(h)}$ with all elements $1$.

The output of DSS (26) can be calculated as

with $h_{i}\coloneqq\bar{C}\bar{A}^{i}\bar{B}$, which is referred to as the impulse response of DSS (26). Given a sample time $\Delta$, $h_{i}$ is determined by $(A,B,C)$, and different sets of parameters $(A,B,C)$ may result in the same sequence of impulse responses. In fact,

are the same for different $(A,B,C)$, as described below [9].

Proposition 1 implies that, under weak assumptions, the impulse responses $h_{0},h_{1},\cdots,h_{L-1}$ of DSS (26) can be achieved with special structure of $(A,B,C)$. DSS (26) with $(B,C)=((1)_{1\leq i\leq N},w)$, as stated in Proposition 1(a), is referred to as DSS${}_{\text{EXP}}$, and DSS (26) with $(B,C)=(((e^{L\lambda_{i}\Delta}-1)^{-1})_{1\leq i\leq N},w)$, as stated in Proposition 1(b), is referred to as DSS${}_{\text{SOFTMAX}}$ [9]. DSS${}_{\text{EXP}}$ and DSS${}_{\text{SOFTMAX}}$ offer different approaches to modeling the impulse responses, with potential implications for the performance and interpretability of the DSS model. In the following sections, we utilize DSS${}_{\text{EXP}}$ or DSS${}_{\text{SOFTMAX}}$ as DSS (26).

Among the parameters $W$ trained in our deep learning model, those in the DSS layer include:

• •

  $(A,B,C,\Delta)$ for each of the $H$ DSS models, as defined in (26).

• •

  $D$ for each of the $H$ nonlinear connection blocks, as defined in (27).

• •

  Weight $W_{out}\in\mathbb{R}^{H\times H}$ and bias $b\in\mathbb{R}^{H}$ for the linear combination block, as defined in (29).

For DSS${}_{\text{EXP}}$ defined in Subsection III-B, the parameters $(A,B,C)$ are defined as

where

For DSS${}_{\text{EXP}}$, the parameters $\Lambda_{re},\Lambda_{im}\in\mathbb{R}^{N}$ and $w\in\mathbb{C}^{N}$ are trained to determine $(A,B,C)$.

For DSS${}_{\text{SOFTMAX}}$ defined in Subsection III-B, the parameters $(A,B,C)$ are defined as:

where

Similarly to DSS${}_{\text{EXP}}$, the parameters $\Lambda_{re},\Lambda_{im}\in\mathbb{R}^{N}$ and $w\in\mathbb{C}^{N}$ are trained to determine $(A,B,C)$ for DSS${}_{\text{SOFTMAX}}$, with a different expression for $A$ and $B$.

The performance of S4 with DSS layers is sensitive to initialization of the state matrix $A$. To obtain an effective initial value for $A$, the HiPPO matrix is decomposed into a normal matrix and a low-rank matrix. This decomposition allows for a more structured and interpretable initialization of the state matrix $A$, which can improve the performance of the model. The eigenvalues of the normal matrix are employed to initialize the diagonal elements of $A$.

In more detail, the HiPPO matrix $\mathcal{H}\in\mathbb{R}^{N^{\prime}\times N^{\prime}}$ is defined as explained in Subsection II-C:

For SSM (1) incorporating this HiPPO matrix, the state $x(t)$ retains information about the history of the input $u(t)$ [7]. The HiPPO matrix $\mathcal{H}$ can be decomposed into a normal matrix and a low-rank matrix as

where $\mathcal{H}^{\prime}\in\mathbb{R}^{N^{\prime}\times N^{\prime}},P\in\mathbb{R}^{N^{\prime}},Q\in\mathbb{R}^{N^{\prime}}$ are defined as

This $\mathcal{H}^{\prime}$ is a normal matrix. Under the assumption that $N^{\prime}=2N$ and $\mu_{1},\ldots,\mu_{N}$ are the eigenvalues of $\mathcal{H}^{\prime}\in\mathbb{R}^{2N\times 2N}$ with positive imaginary parts, $\Lambda_{re},\Lambda_{im}\in\mathbb{R}^{N}$ are defined as follows:

• •

  For DSS${}_{\text{EXP}}$,

  $\displaystyle\Lambda_{re,i}=\log(-\mathrm{Re}(\mu_{i})),\quad\Lambda_{im,i}=\mathrm{Im}(\mu_{i}).$

  (43)

• •

  For DSS${}_{\text{SOFTMAX}}$,

  $\displaystyle\Lambda_{re,i}=\mathrm{Re}(\mu_{i}),\quad\Lambda_{im,i}=\mathrm{Im}(\mu_{i}).$

  (44)

Using $\Lambda_{re}$ and $\Lambda_{im}$, the matrix $A$ is initialized as $A\coloneqq\mathrm{diag}(\lambda_{1},\ldots,\lambda_{N})$, where each $\lambda_{i}$ is derived from Equation (36) for DSS${}_{\text{EXP}}$ or Equation (38) for DSS${}_{\text{SOFTMAX}}$. This process is known as the Skew-HiPPO initialization [12, 9]. Other parameters within the DSS are randomly sampled, as detailed in Section VI. According to [9], models utilizing Skew-HiPPO initialization demonstrate superior prediction accuracy compared to those initialized with values from $\mathcal{N}(0,1)$.

When the input is entered one by one or all at once, reducing $H$ (the hidden size) and $N$ (the state dimension) facilitates a reduction in the computational cost of the DSS layer output. In fact, the time and space complexities of the DSS layer output are as illustrated in Table I and Table II, respectively, as discussed below. Here, the input length of the 1-dimensional sequence is denoted as $L$.

In the case where input $u_{k}$ at each time step is entered one by one into DSS (26), the output $y_{k}$ can be computed using the previous state vector $x_{k-1}$ according to (26). The time and space complexities per step are both $O(N)$. The time complexity for processing the entire input of length $L$ is $O(NL)$, and the space complexity involves overwriting at each step, thus remaining $O(N)$. For the nonlinear connection block, both the time and space complexities per step are $O(1)$. The time complexity for processing the entire input of length $L$ is $O(L)$, and the space complexity involves overwriting at each step, thus remaining $O(1)$. Regarding the following linear combination block, the time complexity per step is $O(H^{2})$, and the space complexity is $O(H)$. The time complexity for processing the entire input of length $L$ is $O(H^{2}L)$, and the space complexity involves overwriting at each step, thus remaining $O(H)$. Therefore, adding up $H$ DSS, $H$ nonlinear connection blocks, and one linear combination block, the time complexity of the DSS layer output when input is entered one by one is $O(HNL)+O(H^{2}L)$, and the space complexity is $O(HN)$ (Table I).

In the case where whole input $(u_{k})_{0\leq k<L}$ is entered all at once, the output

can be efficiently computed. In fact, leveraging the fast Fourier transform [41] implies that the time complexity is $O(L\log(L))$ and the space complexity is $O(L)$. Furthermore, the computation is parallelizable. Besides, the impulse response $h_{i}$ in Equation (45) can be easily computed. In fact, for a diagonal matrix $A\in\mathbb{C}^{N\times N}$ with diagonal elements $\lambda_{i}$, $h_{i}$ can be calculated as

The computation time for the sequence of impulse responses $(h_{0},h_{1},\cdots,h_{L-1})$ is $O(NL)$. As for the nonlinear connection block, the time and space complexities are both $O(L)$ Regarding the following linear combination block, the time complexity is $O(H^{2}L)$, and the space complexity is $O(HL)$. Therefore, adding up $H$ DSS, $H$ nonlinear connection blocks, and one linear combination block, the time complexity of the DSS layer output when input is entered all at once is $O(HL\log(L))+O(H^{2}L)$, and the space complexity is $O(HL)$ (Table II).

Additionally, Equation (45) is an approximation that holds when DSS (26) is asymptotically stable and $L$ is sufficiently large. The exact output of DSS (26) is given by (30). However, when (26) is asymptotically stable and $L$ is sufficiently large, $h_{i}\approx 0$ for $i\geq L$, making the approximation in (45) valid.

Let us consider the issues that hinder the application of S4 with DSS layers to EI [17, 18], as explained in Section I. These issues include memory constraints, computational complexity, and the trade-offs between model performance and resource efficiency.

The following can be concluded from the arguments of Subsection IV-C:

• •

  When processing the entire input of large length $L$ at once, the application of S4 with DSS layers to EI is challenging, even if $H$ and $N$ of the trained model are sufficiently small. This is because the space complexity is $O(HL)$ (Table II), making it difficult to conduct inference in devices with small capacity memory (e.g. sensors set in factories).

• •

  When processing the input one-by-one, which corresponds to $L=1$, S4 with DSS layers can be applied to EI if $H$ and $N$ of the trained model are sufficiently small. Specifically, the time and space complexities are those shown in Table  I. That is, when $L=1$, the time complexity is $O(HN)+O(H^{2})$, and the space complexity is $O(HN)$. This implies that even for very large input sequences, inference can be conducted in devices with small capacity memory.

In summary, to apply S4 with DSS layers to EI, the input needs to be processed one-by-one, and it is desirable to keep the values of $H$ and $N$ as small as possible. However, excessively small values of $H$ and $N$ may limit the model’s capacity to capture complex patterns in the data, leading to a deterioration in performance.

Table V and Table VI show the accuracy of models obtained through various training methods, using DSS${}_{\text{EXP}}$ and DSS${}_{\text{SOFTMAX}}$ respectively, where $N$ denotes the dimension of the state vector of DSS. The number of DSS layers is $4$ and the hidden size $H$ is $16$. The columns labeled “before” and “after” denote the accuracy of the model with the initial parameter values and the accuracy of the model after training from that initial state, respectively.

In the “HiPPO” column on the left, the Skew-HiPPO initialization explained in Subsection IV-B was used to initialize the state matrix $A$ of each DSS. In the middle column “Random”, the initial values of $A$ were randomly sampled. For DSS${}_{\text{EXP}}$, the real and imaginary parts of the diagonal elements of matrix $A$ were sampled from $\mathcal{U}(-1,0)$ and $\mathcal{N}(0,1)$, respectively. For DSS${}_{\text{SOFTMAX}}$, the real and imaginary parts of the diagonal elements were sampled from $\mathcal{N}(0,1)$. Other parameters within DSS were randomly sampled for both “HiPPO” and “Random”. The real and imaginary parts of each element in $w\in\mathbb{C}^{N}$ were sampled from $\mathcal{N}(0,1)$.