SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shida Wang
Department of Mathematics National University of Singapore e0622338@u.nus.edu
Beichen Xue
Department of Statistics and Data Science National University of Singapore e0773769@u.nus.edu
State-space models have gained popularity in sequence modelling due to their simple and efficient network structures. However, the absence of nonlinear activation along the temporal direction limits the model's capacity. In this paper, we prove that stacking state-space models with layer-wise nonlinear activation is sufficient to approximate any continuous sequence-to-sequence relationship. Our findings demonstrate that the addition of layer-wise nonlinear activation enhances the model's capacity to learn complex sequence patterns. Meanwhile, it can be seen both theoretically and empirically that the state-space models do not fundamentally resolve the issue of exponential decaying memory. Theoretical results are justified by numerical verifications.
State-space model [1-5] is an emerging family of neural networks specialized in learning long sequence relationships. It achieves significantly better performance compared with attention-based transformers in the long range arena (LRA) dataset [6, 7]. Despite its effectiveness, the state-space model is built on a relatively simple foundation of linear-RNN-like layers. One of the key advantages of state-space models is their simple recurrence, which enables efficient acceleration. In fact, this recurrence allows for an asymptotic computational complexity of only  $O(T \log T)$ , which is significantly better than the  $O(T^2)$  complexity of traditional full-attention approaches [8]. A natural question would be whether SSM achieves this speedup with certain sacrifices in model capacity or memory property. It is currently unclear whether the state-space model's linear architecture with layerwise nonlinearity possesses sufficient expressive capacity to approximate any target sequence-to-sequence relationship. This knowledge would be important to answer pertinent questions regarding the model's ability to handle the complexity of real-world datasets characterized by diverse and intricate sequence relationships. In particular, considering the speed advantage of SSM over attention-based transformers, the universal approximation property impacts whether a state-space model could be a suitable replacement for a transformer.
In this paper, we study the universality of state-space model. Furthermore, the memory property is investigated and we show that state-space models also have an asymptotically exponential decaying memory.
The main contributions can be summarized as follow
1. We give a constructive proof for the universal approximation property for multi-layer state-space models. The width dependency on sequence length is analyzed.
Figure 1: Network structure of two-layer state-space model.
2. The state-space models are shown to have an exponentially decaying memory, which coincides with usual recurrent neural networks.
3. Numerical verifications for the memory property are given on synthetic datasets.
In this section, we introduce the general form of the state-space models. The classical results on universal approximation for recurrent neural networks are summarized. Based on the approximation result, the well-documented challenge of learning long-term memory via recurrent networks is summarized. In particular, we give the definition of memory function which we shall use in the derivation of memory decay.
2.1 State-space models
Single-layer state-space model can be viewed as a recurrent neural network without nonlinear recurrent activation. As is shown in Figure 1, the discrete time version of SSM [9] is
$$
h _ {k + 1} = W h _ {k} + U x _ {k + 1}, \quad h _ {0} = 0 \tag {1}
$$
$$
y _ {k} = C h _ {k} + D x _ {k} = C W ^ {k} h _ {0} + \sum_ {i = 1} ^ {k} C W ^ {k - i} U x _ {i} + D x _ {k}. \tag {2}
$$
where  $x \in \mathbb{R}^{d_{in}}$ ,  $y \in \mathbb{R}^{d_{out}}$ ,  $h \in \mathbb{R}^m$ ,  $W \in \mathbb{R}^{m \times m}$ ,  $U \in \mathbb{R}^{m \times d_{in}}$ ,  $C \in \mathbb{R}^{d_{out} \times m}$ ,  $D \in \mathbb{R}^{d_{out} \times d_{in}}$ . Notice that we drop the bias term for simplicity. For multi-layer state-space model, the nonlinear activation is added layer-wise.
The continuous time version of a single layer is
$$
\frac {d h _ {t}}{d t} = W h _ {t} + U x _ {t}, \quad h _ {0} = 0 \tag {3}
$$
$$
y _ {t} = C h _ {t} + D x _ {t} = \int_ {0} ^ {T} C e ^ {W (t - s)} U x _ {s} d s + D x _ {t}. \tag {4}
$$
It can be seen in Equation (4) that the first component of output  $y$  is the convolution between kernel function  $\rho(t) = Ce^{W(t - s)}U$  and  $x$ . In discrete case, the convolution operator  $*$  is represented by
$$
y _ {[ 0, T ]} = \rho (t) * x _ {[ 0, T ]} \Rightarrow y _ {k} = \sum_ {i = 0} ^ {k} \rho_ {k - i} x _ {i}. \tag {5}
$$
Compared with temporal convolution network (TCN) [10], which has a finite kernel size, state-space models can be regarded as implementing global convolution  $(\rho(t) = Ce^{W(t - s)}U, t \geq 0)$  over the
temporal axis. Numerically, since the kernel size is the same as the sequence length, the convolution implemented via fast Fourier transform can lead to significant accelerations [11]. Compared with attention-based transformer, it is shown that SSM can speed up the training 100x in learning sequences with length 64K [12]. The state-space model's temporal efficiency is achieved by eliminating the nonlinear activation function in its recurrent layers  $h_{k + 1} = Wh_k + Ux_k + b$ , resulting in faster processing of sequential data via parallel scan [11].
2.2 Universal approximation in RNN
It is long-known that recurrent neural networks with nonlinear sigmoidal activations (such as tanh and sigmoid) are universal approximators.
Theorem 2.1 (Simplified universality statement). For any sequence to sequence map:  $\mathbf{H}:\mathbf{x}\rightarrow \mathbf{y}$  and tolerance  $\epsilon$ , there exists a hidden dimension  $m$  and weights  $c, W, U, b$  such that the RNN with weights  $(c,W,U,b)$  can approximate the target sequence to sequence map:
$$
\| y - \hat {y} \| \leq \epsilon . \tag {6}
$$
where prediction sequence  $\hat{y}$  is given by
$$
h _ {k + 1} = \sigma \left(W h _ {k} + U x _ {k} + b\right), \tag {7}
$$
$$
\hat {y} _ {k} = C ^ {\top} h _ {k}. \tag {8}
$$
Universal approximation establishes the feasibility of learning sequence to sequence relationships via recurrent neural networks. However, typical approximation rate results depend on the sequence length [13]. Sequence length dependent approximation rate does not generalize to the case of sequence of infinite length. In learning sequences with infinite length, Li et al. [14] shows that linear RNNs have difficulty in learning non-exponential decaying memory. Various numerical experiments [15] confirm that adding nonlinear recurrent activation does not fundamentally change the decay. In state-space models, the nonlinearity is included in a layer-wise approach. It is unknown whether such layer-wise nonlinearity alone is sufficient to approximate any sequence to sequence relationships.
2.3 Memory function and curse of memory
In this paper we study the memory property of state-space model. Before we introduce the main results, we present a simple memory function definition in sequence modelling. Li et al. [14] proves that a bounded causal continuous regular time-homogeneous linear functional has the following Riesz representation:
$$
y _ {t} = H _ {t} (\mathbf {x}) = \int_ {- \infty} ^ {t} \rho (t - s) x _ {s} d s. \tag {9}
$$
Here  $\rho : \mathbb{R}^{+} \to \mathbb{R}$  is an  $L_{1}$ -integrable function. If  $\rho_{t}$  rapidly decreases with  $t$ , then the target sequence relationship has a short-term memory.
Since  $\rho_{t}$  fully captures the memory property of a linear functional [14], we call it the memory function (of the linear functional). In particular, when approximating the linear functional with linear RNNs, the model's memory function is  $\hat{\rho} (t - s) = C^{\top}e^{W(t - s)}U$
The curse of memory refers to the phenomenon that if a target linear functional can be approximated by a sequence of linear RNNs, the memory function  $\rho_{t}$  decays exponentially.
Take the test input to be  $\mathbf{x}_{\mathrm{test}} = \left\{ \begin{array}{ll}1 & t\geq 0,\\ 0 & t < 0. \end{array} \right.$  Notice that the derivative of the linear functional at test input extracts the memory function  $|\frac{d}{dt} H_t(\mathbf{x}_{\mathrm{test}})| = |\rho (t)|_2$  . Therefore a natural extension of the memory function [16] to the bounded causal continuous regular time-homogeneous nonlinear functionals will be:
$$
\hat {\rho} (t) = \left| \frac {d \hat {y} _ {t}}{d t} \right| _ {2}, \quad \hat {y} _ {t} = \widehat {\mathbf {H}} _ {t} (\mathbf {x} _ {\text {t e s t}}). \tag {10}
$$
It can be seen this definition is compatible with the memory function definition of linear functional. Also, this memory function can be evaluated by computing the models' derivative at the test input using finite difference method.
Figure 2: Two-layer state-space model can approximate any continuous element-wise function
In this section, we first give a simple constructive proof to show that any element-wise function on input sequence can be approximated by a two-layer state-space model. Next, we show any temporal convolution can be approximated by a state-space model. The constructive proof for general nonlinear sequence to sequence functional is given based on the above propositions. Moreover, as the Kolmogorov-Arnold-representation-based construction has a weight number dependent on sequence length, it can be highly inefficient to construct a shallow wide network to learn long sequences. To reduce the widths dependency on the sequence length, a Volterra-series-based construction is demonstrated.
3.1 Two-layer SSM approximates element-wise function
By element-wise function, we mean learning sequence relationship with form  $(x_{1},\ldots ,x_{T})\to (f(x_{1}),\ldots ,f(x_{T}))$
Proposition 3.1. For any given continuous function  $f$  over a compact set  $K$ , let the activation  $\sigma$  be sigmoidal function.
$$
\lim  _ {z \rightarrow \infty} \sigma (z) = 1, \quad \lim  _ {z \rightarrow - \infty} \sigma (z) = - 1. \tag {11}
$$
There exists a sequence of two-layer state-space model with weights  $\{W_1, W_2, U_1, U_2, b_1, b_2\}$  that can approximate sequence relationship  $\mathcal{H}:(x_1,\ldots,x_T)\to(y_1,\ldots,y_T)$ .
$$
\sup  _ {t} \sup  _ {\mathbf {x}} | f (\mathbf {x}) - \hat {\mathbf {y}} | \leq \epsilon . \tag {12}
$$
Here the two-layer state-space model is constructed by
$$
h _ {k + 1} ^ {(1)} = W _ {1} h _ {k} ^ {(1)} + U _ {1} x _ {k} + b _ {1}, \tag {13}
$$
$$
h _ {k + 1} ^ {(2)} = W _ {2} h _ {k} ^ {(2)} + U _ {2} \sigma \left(h _ {k + 1} ^ {(1)}\right) + b _ {2} \tag {14}
$$
$$
y _ {k} = h _ {k} ^ {(2)}. \tag {15}
$$
The proof is included in Appendix B.2. The main idea is to approximate element-wise function  $f(x)$  with  $U_{2}\sigma (U_{1}x + b_{1})$ . A graphical demonstration is given in Figure 2. It can be seen the two-layer state-space model can approximate any element-wise function by setting  $W_{1} = W_{2} = b_{2} = 0$ .
3.2 SSM approximates temporal convolution
In this part we show that state-space models can approximate temporal convolution:  $y_{k} = \sum_{i=1}^{k} \rho_{k-i} x_{i}, 1 \leq k \leq T$ . According to Equation (3), the hidden state of SSM is the convolution between input sequence and exponentially decaying function  $\hat{\rho}_{k} = C W^{k} U$ .  $y_{k} = \sum_{i=1}^{k} C W^{k-i} U x_{i} + C W^{k} h_{0} + D x_{k}$ .
Therefore the approximation problem of temporal convolution by state-space models is reduced to the approximation problem of general convolution kernel  $\rho_{k}, 1 \leq k \leq T$  by exponentially decaying convolution kernel  $\hat{\rho}_{k}, 1 \leq k \leq T$ . The optimal weights are defined by  $C, W, U = \arg \min_{C, W, U} \sup_{k} |\rho_{k} - CW^{k}U|$ .
Proposition 3.2. For any given convolution kernel  $\rho_{k}$ ,  $1\leq k\leq T$ , the single-layer state-space model is universal approximator for temporal convolution.
In other words, for any  $\epsilon > 0$ , there exists a hidden dimension  $m$  and corresponding weights  $C, W, U$  such that  $\rho_k = CW^k U$  satisfies
$$
\sup  _ {k} \left| \rho_ {k} - \hat {\rho} _ {k} \right| <   \epsilon . \tag {16}
$$
See Appendix B.3 for the proof. The main idea is to represent the single-layer state-space model output as a convolution between input and kernel function. The approximation of temporal convolution is then reduced to the approximation of general kernel with the SSM-induced kernels.
Remark 3.3. Although the Proposition 3.2 indicates that we can use single-layer state-space model to approximate any convolution, it does not reveal the necessary hidden dimension  $m$  for such approximation.
3.3 Universality of SSM
Now we show that five-layer state-space model is universal. Without loss of generality, assume the output  $y$  is one-dimensional. The main proof is based on the famous Kolmogorov-Arnold representation theorem [17]. By Kolmogorov-Arnold representation theorem, we know any multivariate continuous function  $f:\mathbb{R}^d\to \mathbb{R}$  can be represented by
$$
f \left(x _ {1}, \dots , x _ {d}\right) = \sum_ {q = 0} ^ {2 d} \Phi_ {q} \left(\sum_ {p = 1} ^ {d} \phi_ {q, p} \left(x _ {p}\right)\right). \tag {17}
$$
Remark 3.4. George Lorentz [18] shows that we can use the same  $\Phi: f(x_1, \ldots, x_d) = \sum_{q=0}^{2d} \Phi\left(\sum_{p=1}^{d} \phi_{q,p}(x_p)\right)$ . Sprecher [19] proves that it can be further reduced to the same  $\phi: f(x_1, \ldots, x_d) = \sum_{q=0}^{2d} \Phi\left(\sum_{p=1}^{d} \lambda_p \phi(x_p + \eta p) + c_q\right)$ . Braun and Griebel [20] gives the first constructive proof for the superposition. Moreover, the inner function  $\phi$  is shown to be independent of target function  $f$ . It means the learning of function  $\phi$  can be approximated without retraining for different target functionals.
We summarize the Kolmogorov-Arnold theorem as follow:
Theorem 3.5 ([17, 20]). Fix dimension  $d \geq 2$ . There are real numbers  $a, b_p, c_q$  and a continuous and monotone function  $\phi: \mathbb{R} \to \mathbb{R}$ , such that for any continuous function  $f: [0,1]^d \to \mathbb{R}$ , there exists a continuous function  $\Phi: \mathbb{R} \to \mathbb{R}$  with
$$
f \left(x _ {1}, \dots , x _ {d}\right) = \sum_ {q = 0} ^ {2 d} \Phi \left(\sum_ {p = 1} ^ {d} b _ {p} \phi \left(x _ {p} + q a\right) + c _ {q}\right). \tag {18}
$$
Proposition 3.6. For any bounded causal continuous sequence to sequence relationship  $H: \{x_k\}_{k=1}^T \to \{y_k\}_{k=1}^T$  and tolerance  $\epsilon > 0$ , there exists a hidden dimension  $m$  and corresponding state-space model (as constructed in Figure 3) such that the error of approximation
$$
\left| y _ {k} - \hat {y} _ {k} \right| \leq \epsilon , \quad k \in \{1, \dots , T \}. \tag {19}
$$
See the proof in Appendix B.4. The main idea is demonstrated in Figure 3, the nonlinear functions are separately approximated by two-layer state-space model.
Remark 3.7. The Kolmogorov theorem provides a construction for achieving universality in a five-layer state-space model. However, the quantity of hidden neurons increases linearly with the sequence length. This can become exceedingly burdensome when the sequence length escalates.
Another approach is from the perspective of Volterra Series, which features the sequence-length independent neurons.
Figure 3: Multi-layer state-space models are universal approximators: Drawing from the Kolmogorov-Arnold representation theorem, we have  $f(x_{1},\ldots ,x_{d}) = \sum_{q = 0}^{2d}\Phi \left(\sum_{p = 1}^{d}b_{p}\phi (x_{p} + qa) + c_{q}\right)$ . Here, both element-wise nonlinear functions  $\phi$  and  $\Phi$  can be approximated by a two-layer state-space model, as shown in Proposition 3.1. Additionally, the temporal convolution is represented by a single-layer state-space model, as detailed in Proposition 3.2.
Theorem 3.8 ([21]). For any continuous time-invariant system with  $x(t)$  as input and  $y(t)$  as output can be expanded in the Volterra series as follow
$$
y (t) = h _ {0} + \sum_ {n = 1} ^ {N} \int_ {0} ^ {t} \dots \int_ {0} ^ {t} h _ {n} \left(\tau_ {1}, \dots , \tau_ {n}\right) \prod_ {j = 1} ^ {n} x \left(t - \tau_ {j}\right) d \tau_ {j}. \tag {20}
$$
In particular, we call the expansion order  $N$  to be the series' order.
A simplified interpretation of the  $N$ -th Volterra Series expansion is the "  $N$ -th Taylor expansion in the sequence variable  $\mathbf{x}$ "
Proposition 3.9. For any bounded causal continuous time-homogeneous sequence to sequence relationship  $H: \{x_k\}_{k=1}^T \to \{y_k\}_{k=1}^T$  and tolerance  $\epsilon > 0$ , there exists a hidden dimension  $m$  and corresponding state-space model (as constructed in Figure 4) such that the error of approximation
$$
\left| y _ {k} - \hat {y} _ {k} \right| \leq \epsilon , \quad k \in \{1, \dots , T \}. \tag {21}
$$
Moreover, the neurons of the state-space model do not explicitly depend on the sequence length  $T$ .
See the proof in Appendix B.5. The main idea is to approximate the convolution kernels  $h_n(\tau_1, \ldots, \tau_n)$  by the low-rank tensor product of first-order convolution kernel.
Figure 4: Volterra-series-type construction for state-space models
Remark 3.10. The advantage of Volterra-series-type construction is the approximation of the convolution kernel with SSM-induced kernel as well as the approximation of multivariable convolution kernel with tensor product of one-dimensional convolution kernel do not explicitly depend on the sequence length. Similar approaches have been adopted by idea of implicit convolution in CKConv [22].
3.4 Memory decay of SSM
In the ensuing discourse, we turn our focus towards an examination of the memory property inherent in state-space models. It has been thoroughly studied in literature that recurrent neural networks exhibit a phenomenon of exponential memory decay [14]. An intriguing question that naturally arises in this context is whether state-space models are plagued by similar challenges. Upon careful investigation, it is concluded that state-space models, much like their neural network counterparts, do possess an asymptotically exponential decaying memory.
Proposition 3.11. Assume there exists a constant  $c_0 > 0$  such that
$$
\lim  _ {t \rightarrow \infty} e ^ {c _ {0} t} \| x _ {t} - x ^ {*} \| \rightarrow 0, \quad x ^ {*} := \lim  _ {t \rightarrow \infty} x _ {t}. \tag {22}
$$
Assume the output  $y_{t}$  is the output of single-layer state-space model with parameters  $C, W, U$ . Then, for the same constant  $c_{0}$ , if the output's derivative satisfies  $\frac{dy_t}{dt} \to 0$
$$
\lim  _ {t \rightarrow \infty} e ^ {c _ {0} t} \| y _ {t} - y ^ {*} \| \rightarrow 0. \tag {23}
$$
At the same time, general smooth nonlinear activation does not change the exponential decay property of a sequence:
Proposition 3.12. Assume there exists a constant  $c_0 > 0$  such that
$$
\lim  _ {t \rightarrow \infty} e ^ {c _ {0} t} \| x _ {t} - x ^ {*} \| \rightarrow 0, \quad x ^ {*} := \lim  _ {t \rightarrow \infty} x _ {t}. \tag {24}
$$
For given Lipschitz continuous layer-wise activations  $\sigma$ , there exists a positive constant  $c_{0}$  such that the output memory function  $\frac{dy_t}{dt} \to 0$
$$
\lim  _ {t \rightarrow \infty} e ^ {c _ {0} t} \| y _ {t} - y ^ {*} \| \rightarrow 0. \tag {25}
$$
See the proofs for above two propositions in Appendix B.6 and Appendix B.7.
Based on the above two propositions, by induction we have the following theorem:
Theorem 3.13. Assume  $\mathbf{H}$  is a multi-layer state-space model with Lipschitz continuous function as the layer-wise activations. Assume the state-space model is stable in the sense that matrix  $W$  's eigenvalue are bounded by 1.
There exists a positive constant  $c_0$  such that the memory function (defined in Equation (10)) of state-space model is decaying exponentially
$$
\lim  _ {t \rightarrow \infty} e ^ {c _ {0} t} \hat {\rho} (t) \rightarrow 0. \tag {26}
$$
4 Numerical verifications
Based on the generalization of memory function  $\rho$  in linear functional, we verify the asymptotic memory decay of state-space models with simple randomly generated models. The definition is given in one-dimensional case, but it can be generalized to multi-variable case by taking different unit inputs in various coordinates. The motivation for the above definition is based on the idea of measuring the earlier input at the later output.
In our experiment, we construct various RNN models and SSM with random generated weights. It can be seen in Figure 5 that the memory of naive state-space models also has an exponentially decay memory. It is consistent with the previous theorem that state-space model has an asymptotic exponentially decaying memory. Notice that here the naive SSM is simply adding a tanh activation across layers without specially tuning the weights. Such random initialization can expose the memory issue more significantly as the S4 layer is constructed with several parameterization techniques. However, the manually constructed S4 still has an asymptotic exponential decaying memory as is shown in Figure 6.
In this section, we introduce the previous works on state-space models. As the single-layer state-space model is a linear RNN, we summarize the related approximation work on RNN. In particular, the approximation result and memory result is emphasized as this paper works on the universal approximation property and memory decay property of SSM.
State-space models State-space models originate from the HIPPO matrix which is optimal in the online function approximation sense [2-4]. The Hippo matrix initialization for recurrent matrix  $W$  enables the state-space model to have a slow decaying memory. The universal approximation idea is heuristically demonstrated in Orvieto et al. [23]. However, the proof from the Koopman theory perspective only guarantees the existence of universal approximation. Our proof is a constructive proof which can be further generalized to study the approximation rate with respect to the hidden dimensions and network depths.
Recurrent neural networks Recurrent neural networks (RNNs) [24] are one of the most popular neural networks for sequence modelling. Various results have been established in RNNs approximation theory, see Sontag [25], Hanson et al. [13]. Apart from the universal approximation, the exponential decaying memory property is the notorious phenomenon in recurrent neural networks which prohibits the scale up of the models in terms of the sequence length [14, 26].
In Table 1 we compare the classical sequence models including RNN, TCN and attention-based transformer. The state-space model can be considered as an enhancement of Recurrent Neural Networks (RNNs) due to its superior optimization and inference speed. Despite maintaining a similar network topology, inference cost, and memory pattern, it provides a more efficient and streamlined approach. The effort to extend the long-memory learning is also carried out in convolutional networks and attention-based transformers. Romero et al. [22] proposes to parameterize the convolution
(a) RNN
(b) GRU
(c) LSTM
(d) Naive SSM
Figure 5: Memory functions of a randomly initialized recurrent networks. The shadow indicates the error bar for 100 repeats.
kernel implicitly, which utilizes the power of spline function approximation  $y = \sin (w_0(Wx + b))$ . Similar idea has been adopted in Poli et al. [12]. To summarize, we regard SSM, CKConv and linear transformer as the sequence models' improvement in learning long-memory for long sequences.