Robustifying State-space Models for Long Sequences via Approximate Diagonalization
Authors: Annan Yu
Arnur Nigmetov
Dmitriy Morozov
Michael W. Mahoney
N. Benjamin Erichson
Venue: ICLR 2024 spotlight (2024)
Keywords: state-space models
sequence models
Long-Range Arena
recurrent neural networks
Primary area: general machine learning (i.e., none of the above)
TL;DR: In this work, we introduce a "perturb-then-diagonalize" (PTD) methodology to address diagonalization issues of the HiPPO initialization framework to robustify state-space models (SSMs) for long-range sequence tasks.

Abstract: State-space models (SSMs) have recently emerged as a framework for learning long-range sequence tasks. An example is the structured state-space sequence (S4) layer, which uses the diagonal-plus-low-rank structure of the HiPPO initialization framework. However, the complicated structure of the S4 layer poses challenges; and, in an effort to address these challenges, models such as S4D and S5 have considered a purely diagonal structure. This choice simplifies the implementation, improves computational efficiency, and allows channel communication. However, diagonalizing the HiPPO framework is itself an ill-posed problem. In this paper, we propose a general solution for this and related ill-posed diagonalization problems in machine learning. We introduce a generic, backward-stable ``perturb-then-diagonalize'' (PTD) methodology, which is based on the pseudospectral theory of non-normal operators, and which may be interpreted as the approximate diagonalization of the non-normal matrices defining SSMs. Based on this, we introduce the S4-PTD and S5-PTD models. Through theoretical analysis of the transfer functions of different initialization schemes, we demonstrate that the S4-PTD/S5-PTD initialization strongly converges to the HiPPO framework, while the S4D/S5 initialization only achieves weak convergences. As a result, our new models show resilience to Fourier-mode noise-perturbed inputs, a crucial property not achieved by the S4D/S5 models. In addition to improved robustness, our S5-PTD model averages 87.6% accuracy on the Long-Range Arena benchmark, demonstrating that the PTD methodology helps to improve the accuracy of deep learning models.

Published as a conference paper at ICLR 2024
ROBUSTIFYING STATE-SPACE MODELS FOR LONG SE- QUENCES VIA APPROXIMATE DIAGONALIZATION
Annan Yu, Arnur Nigmetov, Dmitriy Morozov, 2
Michael W. Mahoney, 2,3,4 N. Benjamin Erichson 2:3
Center for Applied Mathematics, Cornell University, Ithaca, NY 14853, USA
2 Lawrence Berkeley National Laboratory, Berkeley, CA 94720, USA
3 International Computer Science Institute, Berkeley, CA 94704, USA
Department of Statisties, University of California at Berkeley, Berkeley, CA 94720, USA ay2628corne edu, {ani gmetov, dmorozor/@lll.. gov, mmahoney@stat wew.eley.......k.. edu, erichson@icsi htt.ssiibbeeeeee..... edu
ABSTRACT
State-space models (SSMs) have recently emerged as framework for learning long-range sequence tasks. An example is the structured state-space sequence (S4) layer, which uses the diagonal aesoonal-plus-looe plus-low-rank structure of the HiPPO initialization framework. However, the complicated structure of the S4 layer poses challenges; and, in an effort to address these challenges, models such as S4D and S5 have considered purely diagonal structure. This choice simplifies the implementation, improves computational efficiency, and allows channel communication. However, diagonalizing the HiPPO framework is itself ill-posed problem. In this paper, we propose general solution for this and related ill-posed diagonalization prob- lems in machine learning. We introduce generic, backward-stan stable "perturb-then- diagonalize" (PTD) methodology, which is based on the pseudospectral theory of non-normal operators, and which may be interpreted as the approximate diagonal- ization of the non-normal matrices defining SSMs. Based on this, we introduce the S4-PTD and S5-PTD models. Through theoretical analysis of the transfer functions of different initialization schemes, we demonstrate that the S4-PTD/S5- PTD initialization strongly converges to the HiPPO framework, while the S4D/S5 initialization only achieves weak convergences. As result, our new models show resilience Fourier-mode noise-perturbed inputs, a crucial property not achieved by the S4D/S5 models. In addition to improved robustness, our S5-PTD model av- erages 37.6% accuracy on the Long-Range Arena benchmark, demonstrating that the PTD methodology helps to improve the accuracy of deep learning models.
INTRODUCTION
Sequential data are pervasive across a wide range of fields, including natural language processing, speech recognition, robotics and autonomous systems, well as scientific machine learning and financial time-series analysis, among others. Given that many of these applications produce exceed- ingly long sequences, sequential models need to capture long-range temporal dependencies order to yield accurate predictions.
To this end, many specialized deep learning methods have been de- veloped to deal with long sequences, including recurrent neural networks (RNNs) (Arjovsky al., 2016; Chang et al., 2019; Erichson et al., 2021; Rusch & Mishra, 2021; Orvicto al., 2023), con- volutional neural networks (CNNs) (Bai et al., 2018; Romero et al., 2022), continuous-time models (CTMs) (Gu et al., 2021; Yildiz al., 2021), and transformers (Katharopoulos et al., 2020; Choro- manski et al., 2020; Kitaev et al., 2020; Zhou et al., 2022; Nie et al., 2023).
Over the past few years, the new class of state-space models (SSMs) gained vast popularity for sequential modeling due to their outstanding performance on the dataset (Tay al., 2021). An SSM is built upon continuous-time linear time-invariant (LTTI) dy-
Arena (LRA)

---

Published as a conference paper at ICLR 2024
namical system Σ = (A C.D), which is system of linear ODEs given by
x (1)) - Ax(() + Bu((), y(t) = Cx(() + Du(t),
(1)
where A E Cuxn, BEC CE Coxa, D E C are the state, input, output and feedthrough matrices; and u(t) E C", x(()) E C", y(t) E C" are the inputs, states, and outputs of the system, respectively. The system can be discretized at time steps jål, where Al > 0 and j = 1,..., L, to be fed with sequential inputs of length L. To store and process the information of the long sequential inputs online, the SSMs are often initialized by a pre-designed LTI system. One of the most popular schemes is called "HiPPO initialization" (Voelker et al., 2019; Gu et al.. 2020), in which the Legen- dre coefficients of the input history at time i..., u- [o. are stored and updated in the state vector model long-range dependencies in sequential
xt).. This initialization is specifically designed data. The recently proposed S4 model (Gu et al., 2022b) leverages the HiPPO initialization and ac- celerates training and inference by decomposing A into the sum of a diagonal matrix and low-rank (DPLR) structure yields a barycentric representation (Antoulas & Anderson, 1986) of the transfer function of eq. (1) that maps inputs to outputs the frequency domain, enabling fast computation in the frequency domain (Aumann & Gosea, 2023).
one. The
While the DPLR structure achieves an asymptotic speed-up up of the model, considering A to be a di- agonal matrix results in simpler structure. Compared to DPLR matrix A, diagonal SSM is not only faster to compute and easier to implement, but also allows integrating channel communica- tion via parallel scans (Smith et al., 2023), thereby improving its performance on long-range tasks. Unfortunately, the problem of diagonalizing the HiPPO framework is exponentially ill-conditioned, as n increases. Hence, while Gu et al. (2022b) shows analytic forms of the eigenvalues and eigen- vectors of HiPPO matrices, they suffer from an exponentially large variance and cannot be used in practice. So far, the most popular way of obtaining diagonal SSM is simply discard the low-rank part from the DPLR structure, leveraging a stable diagonalization algorithm for normal matrix. Discarding the low-rank component changes the underlying diagonalization problem, how- ever; and it abandons the theoretical insights about HiPPO. Still, the resulting model almost matches S4's performance, in practice. Such diagonal models are called S4D (Gu et al., 2022a) when the sys- tems are single- innpllleiipttssnlo-ooooute are multinteiinpu/muuiill--oopg
(i.e., m = p 1) and S5 (Smith et al., 2023) when the systems (i.e., m = p >>>>>>>> 1), which enables channel communication.
The issue of ill-posed diagonalization problems is not merely specific to SSMs. For example, it is known that non-normal matrices make RNNs more expressive (Kerg et al., 2019; Orhan & Pitkow, 2020). More generally, non-normality plays an important role in the training of certain neural net- works (Sengupta& & Friston, 2018; Kumar & Bouchard, 2022). While the ill-posedness of the diago- nalization problem essentially prevents accurate computation of eigenvalues and eigenvectors (i.e., we cannot have small forward error) - in fact, the true spectral information becomes meaningless' using a backward stable cigensolver, one recover the non non-normal matrix accurately (i.e., can have a small backward error) from the wrong eigenvalues and eigenvectors.
In this paper, propose a generic "perturb-then-diee tien-diagonalize" (PTD) methodology as backward stable eigensolver. PTD is based on the idea that a small random perturbation remedies the problem of the blowing up of eigenvector condition number (Davies, 2008; Davies & Hager, 2009; Banks et al., 2021), regularizing the ill-posed problem into a close but well-posed one. It is based on the pseudospectral theory of non -ormal operators (Trefethen & Embree, 2005)² and may be interpreted as the approximate diagonalization of the non-normal matrices.
Our PTD method can be used to diagonalize the highly non-normal HiPPO framework. Therefore, instead of using the eigenvalues of the normal component of the HiPPO matrix to initialize the matrix A as in the S4D and S5 models, we propose to initialize A using the eigenvalues of a perturbed HiPPO matrix (see section 4). The resulting S4-PTD and S5-PTD models are shown to be more robust than their S4D and S5 companions under certain Fourier-mode perturbations. Our method is flexible and can be used to diagonalize many SSM initialization schemes that may be invented in the future.
'f an eigenvector matrix Vi is ill-conditioned, then projecting a vector onto the eigenhasis is unstable so the cigendccomposition suffers from a large variance and docs not reveal any useful information of the matrix. "The pseudospectral theory studies the effect of perturbations on the spectrum of non-normal operator.
2

---

Published as a conference paper at ICLR 2024
Contribution. Here our main contributions: (1) We propose "perturb-then thenndiagonalizz" (PTD) methodology that solves ill-posed diagonalization problems in machine learning when only the backward error is important. (2) We provide a fine-grainn line-graneed analysis that compares the S4 and the S4D initialization. In particular, we quantify the change of the transfer function when discarding the low-rank part of HiPPO, which is done in the diagonal S4D/S5 initialization. We show that while the outputs of the S4D/S5 system on a fixed smooth input converge to those of the S4 system at linear rate as n - oc, the convergence is not uniform across all input functions (see section 3.1). (3) Based on our theoretical analysis, we observe, using the sequential CIFAR task (see section 5.2), that the S4D/S5 models are very sensitive to certain Fourier-mode imput perturbations, which impairs the robustness of the models. (4) We propose the S4-PTD and S5-PTD models that replace the normal component of the HiPPO matrix, used to initialize the S4D and S5 models, with a perturbed HiPPO matrix. Our models are robust to Fourier-mode input perturbations. We theoretically estimate the effect of the perturbation (see section 4). We propose computing the perturbation matrix by solving an optimization problem with a soft constraint. Moreover, matrix but can be applied to any initializations. (5) We provide an ablation study for the size of the perturbation in our models. We also evaluate our S4-PTD and S5-PTD models on LRA tasks, which reveals that the S4-PTD model outperforms the S4D model, while the S5-PTD model is comparable with the S5 model (see section 5.1).
method is not restricted to the HiPPO
2 PRELIMINARIES AND NOTATION
Given an LTI system in eq. (1), we say it is asymptotically stable if the eigenvalues X, of A are all contained in the left half- half-plane, i.e., if Re(A,) for all 1 S j S n. The transfer function of the LTI system is defined by
= C(s-- A
s S (CC((AA)
(2)
where I _ax= is the identity matrix and A(A) is the spectrum of A. The transfer function G rational function with n poles (counting multiplicities) at the eigenvalues of A. Assume x(0) = 0. Then the transfer function maps the imputs to the outputs of the LTI system the Laplace domain by multiplication, i.e., (s) = ((s)(Lu)ss) for all E C, where L is the Laplace transform operator (see Zhou & Doyle (1998)). Assume the LTTI system eq. (1) is asymptotically stable and the input nt)) is bounded and integrable (with respect to the Lehesgue measure) as / ranges over I.
Then the Laplace transform reduces to the Fourier transform:
(s) ((s))
E
where ÿ and ú are the Fourier transforms of y and respectively, and i is the imaginary unit. Let V E Cuxn be an invertible matrix. We can conjugate the system (A, C.D. by V, which yields (V- AV, V BB CV, D). Since the transfer function is conjugation-invariint, the two systems map the same inputs u(-) to the same outputs y-). while the states x(-) are transformed by V. If A is a normal matrix, i.c., AA' = A A, then V is unitary, in which case transforming the states by V is a well-conditioned problem and can be done without loss of information. Issues arise, however, when A is non-normal and V is ill-conditioned.
The state-space models use LTI systems to process time series inputs. Different initializations can be tailored to tasks with different natures, such as the range of dependency (Gu et al., 2023). A par- ticularly successful initialization scheme used in the S4 model is the so-called HiPPO initialization. While there exist several variants of HiPPO, the most popular HiPPO-LegS matrices are defined by
{
1 (iss) 22j IV2k I
i + k, if
(An)aa
(Bu)₃e = v -
(4)
for all S j.k I n and 1 < S m, where 1 [5>4] is the indicator that equals 1 if j > & and 0 otherwise. Such system guarantees that the Legendre coefficients of the input history (with respect to a scaled measure) are stored in the states x(t) over time (Gu et al., 2020). Since computing with the dense matrix Ah is practically inefficient, one conjugates the HiPPO system with matrix V to simplify the structure of Аи. The matrix Ah in eq. (4) has an ill-conditioned eigenvector matrix (Gu et al., 2022b); consequently, instead of solving the ill-posed problem that diagonalizes As, one exploits diagonal- plus-low low-rank (DPLR) structure:
[
1 f (-))) /2j 1/2k
j x k.
Aи = A BuBB
(A = 2 (
(5)
j = k.
3

---

Published as a conference paper at ICLR 2024
where J is a akew-symmetric matrix that can be unitarily diagonalized into A VHANVH The S4 model leverages the HiPPO matrices by initializing
ADPLR = AH 1 - B H.,
BDPLR = V
(6)
and CDPER and DDPER randomly. Such LTI system EDPER (ADPLR: BDPER: CDPER, DDPER) is conjugate via VH (A h; By, CDPLR HH: DDPLR). Hence, they share the transfer function and
the same mapping from the inputs u(-) to the outputs y(.). The S4D model further simplifies the
structure by discarding the rank-1 part from A and therefore initializes
ADiug AH, Bring 2 V
and Ding is henceforth restricted to be diagonal. While both the S4 and S4D models restrict that single-inuutsingll-oupppp (SISO), the S5 model, which
m = 1, i.e., the LTT systems also initializes ADiag = Ah and requires it to be diagonal throughout training, leverages multiple- input/multiple-ourpe (MIMO) systems by allowing m = p > 1. We provide more background information on LTI systems and state-space models in sequential modeling in Appendix B.
Throughout this paper, we use ||-|| to denote vector or matrix 2-norm. Given an invertible square matrix V. we use K( (V)= = |VI| V || to denote its condition number. Given a number 1 S p S co and measurable function f R - C, we use | for the standard I," n.r-norm of f with respect to xc)
the Lebesgue measure on R and L (P)
(f R - IIf
3 THEORY OF THE DIAGONAL INITIALIZATION OF STATE- MODELS
The S4 model proposes to initialize the SSM to store the Legendre coefficients of the input sig- nal in the states x (Gu et al., 2020). This initialization, however, has an ill-conditioned spectrum, preventing a stable diagonalization of the SSM. On the other hand, the S4D model uses a different initialization scheme that has a stable spectrum, allowing for stable diagonalization; however, such initialization lacks an interpretation of the states x. In this section, we conduct a fine-grained anal- ysis of the two initializations, which shows that: (1) for any fixed input signal u(.) with sufficient smoothness, the outputs of the two systems DDPLR and Ding converge to each other with a linear rate (of which the previous analysis is devoid) operators that map input signals to the outputs, the operators do not converge in the operator norm topology as -> (see section 3) While the first observation partially justifies the success of the S4D model, the second one allows us to observe that the diagonal initialization is unstable under certain mode input perturbations (see section 5.2). In this section, we assume m p 1, which is consistent with the S4 and S4D models. Still, our theory can be related to the S5 model, as shown in Smith al. (2023).
- Do; and (2) by viewing EDPLR r and Ding as linear
Fix an integer 1 < l S n. We assume that CDPLR = CDag = - VH. where e₂ is the éth standard basis, and DDPLR Diag- For general CDPLR =
CDaag we can decompose it onto the orthonormal basis {or VV H 1 < f S n) and study each component separately using the theory developed in this section. Let GDPER and GDing be the transfer functions of EDPLR and Ding. respectively, i.e.,
GD = CDPLR(sI-ADpp ADD-L)) + GDiag()) = CDiagl (Siaas Ding DDing- (8) Recall that by eq. (3), |GDPLR (sE) (Sig(s))) measures the difference between the outputs of the two systems given a frequency-s input. We provide a fine- fine-grained analysis of this difference in the two transfer functions in Lemma 1. The lemma is visualized in Figure 1. We
that as increases, GDiag approaches GDPLR in the low-frequency domain, i.e., when |s| is small. However, GDig develops spikes in the figh-freyuency domain. Moreover, for every n 1, zooming into the last spike located at |s| O(n²) reveals that it has constant magnitude (see the subplots on the right in Figure 1). Hence, the convergence of GDiag to GDPLR is non-uniform (see Theorem 2). Moreover, the frequency response is unstable at input frequencies near these spikes, suggesting that the S4D model is not robust to certain input perturbations (see section 5.2).
3.1 INPUT-WISE CONVERGENCE AND SYSTEM-WISE DIVERGENCE OF THE DIAGONAL INITIALIZATION
First, we present a result to show that for a fixed input signal u(-). the outputs of EDPLR and Epag converge to each other -> co. Moreover, while the previous result in Gu et al. (2022a) does not
4

---

Published as a conference paper at ICLR 2024
100
10
Gougle 10000
Gouglt 1000)
10
Gougln-
10
Gra,fom 10)
Cuess
10a
3140
magnitude
4
10rd
10
3 3.182565
318136
10
10
10
10
10
10
3,163007 3,8811177
frequency
Figure 1: The magnitude of transfer function of the S4 model, DDLLL (si),, and that of the S4D model, [si) with CDPLR CDiag = e VH and the SSM size set to different values. Note that GDPLR stays the same regardless of n. Due the limited resolution, the left panel does not correctly reveal the heights of the spikes; however, by zooming into the last spike of Giigg (si),, we see that the peak remains (1) as n + oo (see the right panels). The figure shows that GDing is os- cillatory while GDPLR is smooth; morcover, [GDia (si) does not converge to GDrLR (st) uniformly.
have a rate of convergence, we show that is linear. In fact, the rate is sharp (see Appendix F). This partially explains why the S4D model matches the performance of the S4 model in practice.
Theorem 1. Let u(.) (RR be input function with ||u|| r 1. Let YDPLR (+) and yDiag (-)
be the outputs of DDLL and Ding given the input u(.) and the initial states x(0) = 0, respectively.
For some > 1/2, suppose |ú(s] [[ss] as |s| + co. Then, we have ||YDPLR YDiag I1
+ oc, where the constant in the O-notation only depends on q and the constant in x(s) = [[ss] "). The constant does not depend on q if we restrict E [g'. oc) for fixed g > 1/2.
On V /
The proof is deferred to Appendix F. Since the Fourier transform interchanges smoothness and de- cay, what Theorem says is that under mild assumption that u(-) is sufficiently smooth, the output of the diagonal system converges linearly to that of the DPLR system as n + oo. In Section 3.2, show this smoothness assumption is needed. We know the two systems converge input- wise; natural to ask the convergence is uniform across all input signals:
Theorem 2. The function GDPLR (s) GDa (s) does not converge zero uniformly on the imag- inary axis n Co. In particular, for every 2 1, there exists an input signal u. E (RR n (IR) such that we let Yn.DPER and Yn,Ding be the outputs of EDPER and Ening of de- gree n. respectively, then have ,,,,DLL Yn.Diag |. does not converge to 0 as cc.
Hence, the answer to our question negative: combined with Theorem 1, Theorem 2 says that while sufficiently large S4D model mimics its S4 alternative on fixed smooth input, when predetermine size n, they inevitably disagree, by large amount, on some inputs. Morcover, Theorem 2, the construction of u,,() can be made explicit (see section 5.2).
3.2 SOME NUMERICAL EXAMPLES
In this section, provide some numerical examples corroborating Theorem 1. We defer the im- plication of Theorem 2 to later sections (see section 4 and section 5.2). Theorem tells us that fix a smooth input signal u(t), then the outputs YD,DPER and Yn, eventually converge each other at a linear rate as n oo. In this experiment, we fix two input functions (or more precisely, distributions)
(c( e
u 8n,
where / = s0) is the Heaviside function and 6 is the Dirac delta function at While u, (t) a very smooth function - in particular, we have [s)) = O[ss| -1) - the Dirae delta u, is very non-smooth with Fourier transform that is constantly one. We simulate both systems EpPLR
5

---

Published as a conference paper at ICLR 2024
n = 3
n 15
n 75
Input: Ue
DPLR ... Diag
DPLR
DPLR
output value
0.4
Output value
0.4
Digg
output value
0.4
Diag
0.2
0.2
0.2
5000
10000
5000
10000
0
5000
10000
20
20
20
Input: Ud
Diag DPLR
Diag
Diag
output value
15
outpul value
DPLR
output value
DPLR
10
10
10
5
-5
5000
10000
5000
10000
5000
10000
time step
time step
time step
Figure 2: Simulated outputs of the DPLR and diagonal systems with the input functions u, and ug and varying state-space dimension n. We see that for a smooth input function ug, the outputs of both systems converge rapidly as n increases, whereas the convergence does not happen for a non-smooth input function Udd
and EDaag on both (t) and (t).). More details of the simulation can be found in Appendix F. From Figure 2, we observe that given a smooth input function Ug, the output Dn, converges to yn, rapidly, but the same does not hold for non-smooth input function ug Hence, the smoothness assumption in Theorem is essential. In Figure 8 in Appendix F, we also compute the linear when the input is smooth
I,² of Yn.DPLR - Diag and verify that the convergence enough. We remark that a similar study of Fn can be found in Gu et al. (2022a), where the results appear qualitatively different from those presented in Figure 2. This does not mean either work is wrong: the key distinction is that the discretization step size of the LTT systems (see Appendix B) is fixed in Gu et al. (2022a) a priori, introducing aliasing errors and hiding the high frequencies (Tre- fethen, 2019, Ch. 4.). Consequently, when is large, the difference between GDPLR and GDiag in the high-frequency domain is overlooked. In comparison, in this paper, our theory considers the continuous-time LTT systems, which take every mode into account.
PERTURBING THE HIPPO INITIALIZATION: A NEW WAY OF DIAGONALIZING THE SPACE MODEL
In section 3, we saw the instability of the S4D transfer function at certain Fourier modes. Neverthe- less, the diagonal structure of A is preferred over the DPLR one due its training and inference efficiency and its adaptivity to the MIMO model (i.e., the S5 model) (Smith et al., 2023). To avoid instability in a diagonal model, we want to leverage the HiPPO initialization eq. (4) instead of the one in eq. (7) that discards the rank-1 part. One obvious solution is to diagonalize the HiPPO matrix A A V V-1 and conjugate (A Bir, C, D) using HH However, as shown in Gu et al. (2022a), the eigenvector matrix V is exponentially ill ill-conditioned with respect to making the spectral information meaningless. While the exact eigenvalues and cigenvectors of Ah are very ill-conditioned, since we only care about the backward error of diagonalization, propose the fol- lowing initialization scheme. let E E C be a perturbation matrix. We diagonalize the perturbed HiPPO matrix as
A AH +
E +
VHÂNVEL
We then initialize the systems using Epet = (A Bp Cpea: Dpea) (AH. 'B, C,D), where C and D are random matrices. Therefore, we approximately diagonalize the HiPPO initial- ization in the sense that although the diagonal entries in A do not approximate the eigenvalues of An. . the transfer function of Epen is an approximation of that of EDPLR (see Theorem 3). We call our model STDPTD or S5-PTD, depending the S5 model, where "PTD" stands for "perturb-ther cooen-dannnaazze.. " Since our models are only differ- ent from the $4D and the S5 models in initialization, we refer interested readers to Gu al. (2022a)
whether the model architecture is adapted from the S4D or

---

Published as a conference paper at ICLR 2024
and Smith al. (2023) for discussion of computation details and time/space complexity. Our proposed ceatur--thendiagonaliii method is not restricted to the HiPPO-LegS matrices in eq. (4). This endows our method with adaptivity any (dense) initialization scheme. This adaptivity was absent from the previous line of work on SSMs. Consider the process of diagonalizing the matrix An VH Ar и Vi that is solved by inexact algorithm. In numerical analyst's language, the forward error is the error made in computing the cigenvalues Аи and cigenvectors V H, whereas the backward error asks how close problem that we have solved exactly (i.c., Ah + E) is to the actual problem that we want to solve (i.e., A HH) As we will see in Theorem 3, it is the backward error |E] (but not the forward error) that matters our initialization because it the matrix A (but not the specific forms of V f or Ar) that is important in the transfer function.
Centered around the perturbed initialization scheme eq. (9) are two important questions: (1) What is the difference between the perturbed initialization (Apen, Bpan: Cpen, Dren) and the HiPPO ini- tialization (ADPLR: BDPER CDPLR. DDPLR)? (2) What is the condition number of V "/? The first question is important because it controls the deviation of our perturbed initialization from the suc- cessful and robust DPLR initialization. The second question is important because it shadows the numerical robustness of conjugating the LTI system by V VH. Moreover, since the state vector x(t) transformed by V r via conjugation (see section 2), a small condition number of VI shows that its singular values are more evenly distributed. Hence, the transformation VI does not significantly magnify or compress x(t) onto some particular modes. To study the first question, we define the transfer function of the perturbed system to be
Gren 5) = CPen Apat) B + Dpet-
We control the size of the transfer function perturbation by proving the following theorem.
Theorem 3. Assume Open F = CDPLR V and Dpent = DDPLR. Suppose |E| I e and we normalize the matrices so that V fr Bpen| = | fr BpPLR = |Ceen V = [[Crplr V 1.
For any s on the imaginary axis, we have
PPen (s) (D) (2 In(n) + 4)e + (Vllog()))
While our perturb-then-e -han-digoonlalizee method works for a general initialization and a bound on the transfer function error can always be established, the proof of Theorem 3 leverages the structure of HiPPO matrices to improve this bound. The error in Theorem imaginary axis. Using Hölder's inequality, for any bounded and integrable input function u(.), if Yra and YDPLR are the outputs of Epat and DDPLR, respectively, then we have
the uniform error on the
||yren VDPLRII 1.2 = llrren DDLLR 1.a = ((rem)))) - GDrLR
(10)
< [ (is) (iLL - S ||x||>>
(olg(m))))
where the first and the last steps follow from Parseval's identity. Hence, Theorem 3 gives us an upper bound on the distance between Epen and DPLL in the operator norm topology. The theorem states that the error made by the perturbation is linear in the size of the perturbation. Moreover, the error depends only logarithmically on the dimension n of the state space.
Next, we consider the conditioning of V H. which affects the accuracy of computing V Bpent and the scaling ratio of the states in x(-) (see Appendix B). The following theorem provides a deterministic estimate of the eigenvector condition number for the "best perturbation scheme."
Theorem 4 (Banks et al., 2021, Thm. I.1.)). Given any A e C and E (0, 1), there exists a matrix E E C with E < c and an eigenvector matrix V of A + E such that
x(V) < 4n 3/2 3 (1+-)(11111 1 [||])..
Theorem 4 shows the promise of finding good perturbation matrix to reduce the eigenvector con- dition number. We remark that while Theorem studies the best-case scenario. Banks et al. (2021) also contains probabilistic statement about Gaussian perturbations (see Appendix H). In this paper, we propose to compute by solving the following optimization problem with soft constraint:
minimize t (EE) = x(Vn) + + E||
A# + E VHÄV=L H Á diagonal, (11)
where 7 0 is hyperparameter that controls the trade-off between <(V WH) and ||E|]| We imple- ment a solver to this optimization problem using gradient descent. As y increases, is harder to recover the original states x(-) from the transformed states V x)-) because ke( (V,) increases, but ||a] decreases, resulting in more robust SSM that is closer to the flawless HiPPO initialization..
7

---

Published as a conference paper at ICLR 2024
Model
ListOpa
Text
Retrieva
Image
Pathfinder Path-
Avg.
Transformer
36.37
64.27
57.56
42.44
71.40
X
53.66
Luna-256
37.25
64.57
79.29
47.38
77.72
X
59.37
H-Trans.-ID
49.53
78.69
63.99
46.05
68.78
X
61.41
CCNN
43.60
84.08
X
88.90
91.51
X
68.02
$4
59.60
86.82
90.90
88.65
94.20
96.35
86.09
Liquid-S4
62.75
89.02
91.20
89.50
94.80
96.66
87.32
S4D
60.47
86.18
89.46
88.19
93.06
91.95
84.89
S4-PTD (ours)
60.65
88.32
91.07
88.27
94.79
96.39
86.58
S5
62.15
89.31
91.40
88.00
95.33
98.58
87.46
S5-PTD (ours)
62.75
89.41
91.51
87.92
95.54
98.52
87.61
Table 1: Test accuracies on LRA, where X means the model isn't outperforming random guessing. We use the boldface number to indicate the highest test accuracy among all models for each task. We use the underlined number to indicate the highest test accuracy within the comparable group.
EMPIRICAL EVALUATION AND DISCUSSION
In this section, we present empirical evaluations of our proposed S4-PTD and S5-PTD models. In section 5.1 compare the performance of our full model with the existing ones in the Long Range Arena (LRA). In section 5.2, we perform a sensitivity analysis using the CIFAR-10 dataset to provide real-world evidence that our perturbed initialization scheme is more robust than the one in the S4D/S5 model. Finally, in section 5.3, we study the relationship between the size of the perturbation matrix E and the performance of our models.
5.1 PERFORMANCE THE LONG- RANGE ARENA
The LRA benchmark comprises six tasks with sequential data (Tay et al., 2021). This collection, with its sequence lengths ranging from 1024 to 16000. designed to measure the model's capability of processing the long-range inputs. We train an S4-PTD model and S5-PTD model to learn these tasks, respectively. We adopt the same SSM architectures, and thus the same number of parameters, from the original S4D) (Gu al., 2022a) and S5 papers (Smith et al., 2023). Results reported in Table 1, along with the accuracies of other sequential models, including the Liquid-S4 model which is built upon S4 (Hasani et al., 2023). We report details of hyperparameters While the perturbation matrix E is also tunable, we restrict its size to be less than 10% of that of the HiPPO matrix Aи- promoting the worst-case robustness of our model (see section 5.2). We note that the S4-PTD model outperforms the S4D model (and even the S4 model with the DPLR structure for most tasks), while the S5-PTD model matches the performance of the S5 model.
Appendix J.
5.2 ROBUSTNESS OF OUR PERTURBED MODEL OVER THE DIAGONAL MODEL
Our discussion in section 3 suggests that the S4D initialization is not stable the S4 initialization (see Figure 1). Here, we demonstrate its practical implication regarding the robustness of the model. We train S4D model and an S4-PTD model (with ||]|||| |||/||AA) A H 10-1) =1 learn the sCIFAR task, where the images in the CIFAR-10 dataset (Krizhevsky et al., 2009) are flattened into sequences of pixels. We test the two models against two different test sets: one is taken from the original CIFAR- 10 dataset while the other one is contaminated by 10% of sinusoidal noises whose frequencies are located near the spikes of |GDia We plot the training and test accuracies of the two models in Fig- ure 3a and b. Whereas the two models both achieve high accuracies on the uncontaminated test set, the S4D model does not generalize to the noisy dataset as the S4-PTD model does. That is, the S4D model is not robust these noises. In comparison, since the ST-PTD initialization is uniformly close to the S4 initialization (see Theorem 3) when |E]] is small, the S4-PTD model is robust to noises with any mode. We also perturb the test dataset using noises at different frequencies. In Figure 4, we verify that it is indeed the spikes in |Grisg that makes the S4D initialization not robust. We make two remarks. First, the noises in Figure 3a are the "worst-case" noises and intentionally made to fail the S4D model: in practice, the distribution of sensitive modes of S4D in the frequency domain
10
In Orvieto et al. (2023), the S4D model was carefully tuned have higher accuracies. Since the model architecture does not align with those used in this work, only report the result from the original $4D paper.
8

---

Published as a conference paper at ICLR 2024
50
50
accuracy
accuracy
accuracy
60
86
50
50
30C jangro0
300 101g)00
20
Vett yodeg)
90c (9900)
20
30
40
50
10
20
30
40
50
10
10"
10t
10
10"
10"
epochs
epochs
|
A
(a) Accuracies for S4D
(b) Accuracies for S4-PTT)
(c) Ablation study
Figure 3: (a) and (b): the training and test accuracies of the S4D model and the S4-PTD model on contaminated and uncontaminated CIFAR-10 dataset (see section 5.2). (c): The effect of the perturbation size on the accuracy (shown in red) of the S4-PTD model and the cigenvector condition number (shown in blue) of the perturbed HiPPO matrix (see section 5.3).
accuracy
108
0.
GDin (is)
8
Penn (is)]
0.
04
80
8
0.6
35
02
8
04
20
8
62
10
100
200
3000
403
10
8
(a) Accuracy for S4D
(b) Accuracy for S4-PTD
Figure 4: A Fourier-mode perturbation of frequency is added to the test dataset. The blue curve, plotted the left vertical axis, shows the test accuracy the model under perturbations of different frequencies The orange curve, plotted to the right vertical axis, shows the magnitude of the transfer function of the LTI system associated with the initialization.
gets sparser as increases (see Figure which improves its "average-case" robustness. Also, to enable easy detection of frequencies at which the S4D is unstable. in this experiment, we fix the state matrix A. However, we empirically observed that training the state matrix A does not resolve the robustness issue. We provide more details about these two remarks in Appendix K.2.
5.3 ABLATION STUDY OF OUR MODEL.
As mentioned in section 4, the size of the perturbation plays a key role in the performance of our S4- PTD and S5-PTD models. When E O, the eigenvector condition number of A is exponential in n, making numerically impossible to diagonalize when n is moderately large. On the other hand, when E overshadows A H. the initialization scheme becomes random one, often leading poor performance (Gu et al., 2021). In this section, we train an S4-PTD model to learn the sequential CIFAR (sCIFAR) task. We control the size of the perturbation |E] by changing the hyperparameter in the optimization problem eq. (11). For each perturbation matrix E, we then initialize our S4- PTD model by diagonalizing A# + to different perturbation sizes. We see that our S4-PTD model achieves its best performance when the ratio between the perturbation size and the size of the HiPPO matrix between 10-2 and 1, while the accuracy drops when this ratio gets too small or too large. This aligns with our expectations. In addition, the (blue) curve the cigenvector condition number admits straight-line pattern with slope of roughly -, corroborating the factor e in Theorem 4.
Figure 3c, we plot (in red) the test accuracies with respect
CONCLUSION
In this paper, we propose perturb-then-diagonadd (PTD) methodology that can be used to di- agonalize the non-normal HiPPO matrices. Motivated by our theoretical study, we apply the PTD method to robustify the diagonal initialization used in the S4D and S5 models. While our theory focuses on initialization, some empirical evaluations suggest that the PTD method also robustifies the trained diagonal models, which is interesting future research avenue.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by the U.S. Department of Energy, Office of Science, Office of Advanced Scientific Computing Research, Scientific Discovery through Advanced Computing (SciDAC) pro- gram, under Contract Number DE-ACO2-05CHI1 1231 at Lawrence Berkeley National Laboratory. used the Lawrencium computational cluster provided by the IT Division at the Lawrence Berkeley National Laboratory (Supported by the Director, Office of Science, Office Basic Energy Sciences, of the U.S. Department of Energy) and resources of the National Energy Research Scientific Com- puting Center (NERSC, using award ASCR (SCC-EERCAP002333)) a U.S. Department of Energy Office of Science User Facility located at Lawrence Berkeley National Laboratory. both operated under Contract No. DE-AC02-05CH1 1231. NBE would also like to acknowledge NSF, under Grant No. 2319621, for providing partial support of this work. Our conclusions do not necessarily reflect the position or the policy of our sponsors, and no official endorsement should be inferred.
REFERENCES
Athanasios C. Antoulas and Brian D.O. Anderson. On the scalar rational interpolation problem. IMA Journal of Mathematical Control and Information, 3(2-3):61-88, 1986.
Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. In International Conference on Machine Learning, pp. 1120-12 1128. PMLR, 2016.
Quirin Aumann and lon Victor Gosea. Practical challenges in data-drives interpolation: dealing with noise, enforcing stability, and computing realizations. arXiv preprint 2xXi:22001.04966, 2023.
Shaojie Bai, J. Zico Kolter, and Vladlen Koltun. An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. arXiv preprint arXiv: 1803.0111.0
Jess Banks, Archit Kulkarni, Satyaki Mukherjee, and Nikhil Srivastava. Gaussian regularization of the pseudospectrum and davies' conjecture. Communications on Pure and Applied Mathematics, 74(10))2114-2333 2021.
Bo Chang, Minmin Chen, Eldad Haber, and Ed H. Chi. Antisymmetricrnn: A dynamical system view on recurrent neural networks. In International Conference on Machine Learning. 2019.
Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, et al. Rethinking attention with performers. In International Conference on Machine Learning, 2020.
Paul M. Cohn. Further algebra and applications. Peringer-Verlas London, Ltd., London, 2003. ISBN 1-85233-667-6.
E. Brian Davies. Approximate diagonalization. SIAM journal on matrix analysis and applications,, 1954)105511 2008.
E. Brian Davies and Mildred Hager. Perturbations of Jordan matrices. Journal of Approximation Theory, 156(1):82-99) 2009.
James Demmel. The componentwise distance to the nearest singular matrix. SIAM Journal on Matrix Analysis and Applications, 13(1):10-19, 1992.
N. Benjamin Erichson, Omri Azencot, Alejandro Queiruga, Liam Hodgkinson, and Michael W. Mahoney. Lipschitz recurrent neural networks. In International Conference on Learning Repre- sentations, 2021.
Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrent, convolutional, and continuous-time models with linear state space layers. Ad- vances in neural information processing systems, 34:572-585, 2021.
10

---

Published as a conference paper at ICLR 2024
Albert Gu, Karan Goel, Ankit Gupta, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35:35971- 35983, 2022a.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022b.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Ré. How to train your hippo: State space models with generalized orthogonal basis projections. International Conference on Learning Representations, 2023.
Ramin Hasani, Mathias Lechner, Tsun-Hsuan Wang, Makram Chahine, Alexander Amini, and Daniela Rus. Liquid structural state-space models. International Conference on Learning Repre- sentations, 2023.
Angelos Katharopoulos, Apoory Vyas, Nikolaos Pappas, and François Fleuret. Transformers mns: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning. pp. 5156-5165, PMLR, 2020.
Giancarlo Kerg, Kyle Goyette, Maximilian Puelma Touzel, Gauthier Gidel, Eugene Vorontsov, Yoshua Bengio, and Guillaume Lajoie. Non-normal recurrent neural network (nnrnn); learn- ing long time dependencies while improving expressivity with transient dynamics. Advances in neural information processing systems, 32, 2019.
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. Reformer: The efficient transformer. In International Conference on Machine Learning, 2020.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Ankit Kumar and Kristofer Bouchard. Non-normality in neural networks. In Al and Optical Data Sciences III, volume 12019, pp. 70-76. SPIE, 2022.
Yuqi Nie, Nam H. Nguyen, Phanwadee Sinthong, and Jayant Kalagnanam. A time series is worth 64 words: Long-term forecasting with transformers. In The Eleventh International Conference on Learning Representations, 2023.
Naoki Nonaka and Jun Seita. In-depth benchmarking of deep neural network architectures for cog diagnosis. In Machine Learning for Healthcare Conference, pp. 419- PMLR, 2021.
A. Emin Orhan and Xaq Pitkow. Improved memory in recurrent neural networks with sequential non-normal dynamics. Internation Conference on Learning Representations, 2020.
Antonio Orvieto, Samuel L. Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. arXiv preprint XXiv:2303.063449 2023.
David W. Romero, Anna Kuzina, Erik J. Bekkers, Jakub M. Tomczak, and Mark Hoogendoorn. Ck- conv: Continuous kernel convolution for sequential data. In International Conference on Machine Learning. 2022.
T. Konstantin Rusch and Siddhartha Mishra. Unicorn: A recurrent model for learning very long time dependencies. In International Conference on Machine Learning, pp. 9168-9178. PMLR, 2021.
Biswa Sengupta and Karl J. Friston. How robust are deep neural networks? arXiv preprint arXiv: 1804. 1113,, 2018.
Jimmy T.H. Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for se- quence modeling, In The Eleventh International Conference Learning Representations, 2023,
Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers. International Conference in Learning Representations, 2021.

---

Published as a conference paper at ICLR 2024
Lloyd N. Trefethen. Approximation Theory and Approximation Practice, Extended Edition. Society for Industrial and Applied Mathematics, Philadelphia, PA, 2019.
Lloyd N. Trefethen and Mark Embree. Spectra and Pseudospectra: The Behaviour of Non-normal Matrices and Operators. Springer, 2005.
Aaron Voelker, Ivana Kajié, and Chris Eliasmith. Legendre memory units: Continuous-time repre- sentation in recurrent neural networks. Advances in neural information processing systems, 32, 2019.
Pete Warden. Speech commands: dataset for limited-vocabulary speech recognition. arXiv preprint arXiv:/804.03309, 2018.
Max A. Woodbury. Inverting modified matrices. Princeton University, Princeton, N. J., 1950. Sta- tistical Research Group, Memo. Rep. no. 42,.
Cagatay Yildiz, Markus Heinonen, and Harri Lähdesmäki. Continuous-timg model-based reinforce- ment learning. International Conference on Machine Learning, pp. 12009-12018, PMLR, 2021.
Kemin Zhou and John Comstock Doyle. Essentials of robust control, volume 104. Prentice Hall, Upper Saddle River, NJ, 1998.
Tian Zhou, Ziqing Ma, Qingsong Wen, Xue Wang, Liang Sun, and Rong Jin. Fedformer: Frequency enhanced decomposed transformer for long-term series forecasting. In International Conference on Machine Learning, pp. 27268-27286. PMLR, 2022.
12

---

Published as a conference paper at ICLR 2024
APPENDIX
The Appendix is organized as follows. In Appendix A, we survey the background of ill-posed problems, including conditioning, stability, and backward and forward errors. In Appendix B, we provide more background information on LTI systems, including the derivation of the transfer func- tion, system diagonalization, and system discretization, which provides insights into our theory and models. In Appendix C, we prove result on the difference between the two transfer functions associated with the S4 initialization and the S4D initialization, respectively. Using this result, we prove Theorem 2 and 1 on the uniform divergence and the input-wise convergence in Appendix D and E, respectively. We then present Appendix the details of some numerical experiments cor- roborating these two theorems and in Appendix G an experiment the S4D initialization is not as robust as the S4/S4-PTD initialization. In Appendix H. we prove the results in section 4 that are related to perturbing the HiPPO matrix and introduce more background of perturbed diagonalization, which are then verified in Appendix by a numerical experiment. Fi- nally, in Appendix we give the details of our experiments in section 5 and in appendix K some supplementary results on the robustness of the S4D/S4-PTD model during training.
synthetic example that shows
A MORE BACKGROUND INFORMATION OF ILL-POSED PROBLEMS
To make the phrase "ill-posed" precise, we need to introduce the idea of condition numbers. The conditioning is property of problem and does not depend the algorithm that we use. Ab- stractly, we let the problem space X' and the solution space y be two normed vector spaces with the norms || IIx and ||y, respectively. Each element in X' considered as an instance of the problem and its solution in the solution space y is defined by map f X + y. For example, if we want solve a system Ax by with different matrices A and fixed vector by, then we can make X' the space of n-by-n matrices and ) the space of vectors of length n. In that case, we have ((A) = A bb. Likewise, consider the problem of finding eigenvalues. We can make X' and y both equal to the space of n-by-- matrices and f(A) A, the eigenvalue matrix of A. Now, given an instance X, we define the (absolute) condition number problem f at x to
[f(x) f(x + dx))lyy
x(1; f) lim sup
c |
|
Intuitively, a large condition number means that we perturb the problem by a little bit, then the solution may become drastically different. Hence, in general, do not expect that solution of an ill-conditioned problem can be found accurately using floating-point arithmetic because a small rounding error has a large effect on the computed solution.
Unlike the conditioning an algorithm that solves f. We particularly interested in the "backward stability" of f. That is, for any æ E X, there exists an element X so that f(x) f() and
problem æ, stability property of an algorithm. Let f ; X y be
| ø||x
Fx(2))
||~||x
small, then we say that f is backward stable. Intuitively, this is saying that our algorithm is computing the solution nearby problem é. Note that this different from saying that
f(z))> II
Ex(x)
small. This error measures how accurately solved our problem. The error Ey is called a for- ward error, while F3 is called a backward error. We can control the forward error using the backward error, and this bound is established through the condition number of the problem. (Intuitively, this says that if a problem is well-conditioned, then a small perturbation to the problem does not change the solution by too much. Hence, small backward error leads small forward error.) The advantage of studying backward stability is two-fold. First, the backward error is decoupled from
*Of course, this function is not defined at singular matrices, but since they form a Lebesgue null set, f is still defined almost everywhere.

---

Published as a conference paper at ICLR 2024
Problem Space
Solution Space
An ill-conditioned problem
A close but well-conditioned problem
A
V H
Ah
V
*
*
*
Å H
V H
Å
Figure 5: An illustration the perturbation of an ill-posed problem.
the conditioning of the problem. Hence, backward stable algorithms are much more common than forward stable algorithms, because in many cases, the ill-conditioned problems essentially prevent an algorithm from being forward stable. On the other hand, if an algorithm is backward stable, then we know that its forward error must be small on cell-conditioned problems.
In our paper, we consider the case where we are forced to solve an ill-conditioned problem æ. We propose to use backward stable algorithm of to solve it. Since the problem is ill-conditioned, we do not have that fzx is close to f(x), i.e., we cannot find the solutions accurately. However, we know that fxx is the solution to I. where x Rs T. In many machine learning applications, this is enough to guarantee an acceptable solution. (See Figure 5.)
B MORE BACKGROUND INFORMATION OF STATE-SPACE MODELS
Recall that LTI system (A, B, C) given by
x'(t) = Ax(t) + Bu(t), y() Cx(t) Du(t).
(12)
Assume the initial condition is given by x(0) and the input function u(-) is bounded and integrable. Suppose the system is asymptotically stable. i.c., A(A) is contained in the left half- haln-plane. Then, we have x()) and (-) are also bounded and integrable. By taking the Fourier transform the LTI system, have
six(s) A*(s) + Bū(s), (() C*(s) Dū(s).
(13)
Rearranging the first equation gives us
x(s) (sil - -B)
E I.
Plugging it into the second equation of eq. (13), derive the transfer function the imaginary axis:
(s) [C(sil - n(s),
s R.
G(st)
Let V E Caxn be an invertible matrix. Consider the system Σ = AV, D):
x (1) --1 AVx(() V Bu((),
(14)
y(t) CVx(t) + Du()).
2

---

Published as a conference paper at ICLR 2024
By multiplying the first equation by V, we have
Vx(t) = AVx(t) + Bu(t),
and defining the new state variable {() = Vx((), we can write Σ into
ε(I) Ag(() Bu((),
y() CE(() + Du(t).
Hence eq. (12) and (14) are equivalent with their states connected via V. We also can verify this by computing the transfer function of Σ:
((s) : CV
AV) V B C A) B B
G(s).
The LTI system Σ is continuous- time. In order to apply it to sequential input, we need to discretize the system. Given a step size Åt, there are two common ways of discretizing the system:
I+ A) B= A (I At A
Bilinear: A
A
B, (C,D) - (C,
2
ZOH : A (I)).
B = A A (exp(ALA)) I)B, (C, (C,
Then, the discrete system
x, Ax,-1 + Bu,11
(15)
Y: = Cx -Du
takes the discrete sequential input (Up, u1....). The discrete system eq. (15) mimics the continuous system eq. (12) by sampling the continuous input signal u(-) at time intervals At: (ug, u1....) = (u(OAt), u(1At).... The SSMs store the continuous LTI systems. When evaluating on a discrete input, they discretize the continuous systems using trainable step size At and either the Bilinear or the 7OH descritization.
C PROOF OF THE TRANSFER FUNCTION DEVIATION
In this section, we prove result on the difference between the transfer functions GDPLR and G The starting point is to use the Woodbury matrix identity to separate out the rank-1 part in the resolvent that appears in GDPLR- In section 3, we let C e, V и for some fixed f. Since we will reserve the letter é as an index in the proof, in the appendices, we change the notation and assume C c VH- H- While this introduces notation collision with the length of the output vector y, does not cause any confusion in the proofs.
Lemma 1. Let GDpEr and Gring be defined by eq. (8). For any E C with Re(s) 0, have
1)-- (i-s) ((te)
1 (x-
2p
(x+5
(S) GDiag (s)
(16)
1 (-1)--1 11 (j-s)
2
(++s)
The proof of Lemma I is technical. Here, we provide an intuitive explanation before diving into the proof. The idea is to expand the term (sI A B,2 Woodbury matrix identity (Woodbury, 1950), which leads to a primary term (sI - ) that gets canceled with that in GDPLr and residual terms that are expanded by Cramer's 's rule. The importance of eq. (16) is that it reduces the complicated matrix inversions to elementary operations, enabling further analysis.
the expression of GDiag using
Proof of Lemma I. For notational cleanliness, in this proof, we define A Ar, A A and B B To begin with, expand (sI A - -B using the Woodbury matrix identity (Wood-

---

Published as a conference paper at ICLR 2024
bury, 1950):
(sI H)
B (sl A BB 1 BB
= [(sI A)
+ (sI - B B
(sl - A)
B B (s1 A) A
- A
(sI - A)
B +
(sl -) B
1- B
(sI-
A)
B
2B
(sl
B
(sI A) B 1+
(sl 'B
1- B
(sl A)
2(sl A) -A---B-A 2B
(8l
A)
B -1
(sl A) A) 'B.
(81
A)-LB
Hence, when CDPLR DDig
I, the difference between GDPLR and GDiug can be written
(sI A) B
(s H
AB B (sI AH) ) -B = 2 2 (sI A) B (sT
'B. 'B.
(17)
Our next step is to study B (sI AB that appears in eq. (17). To wit, we use Hua's iden- tity (Cohn, 2003) to obtain
-A-1+ A --- ) B
B (s B
=B (-A-1 + (A(sI -))1-A))] A)) B.
It is easy to see that B A -B -1/2. Hence, we have
B (sl -
- (sI--)--
+ sB (sl - A-B
Note that when s = 0. the second term in the expression above vanishes, and therefore we already have that (A+)-1B/
A -BB. To deal with the general case when s is a purely imaginary -c//v2 because is 1//V2 times the first column of A. Hence, sB (sI - A)-1A A -BB is equal to times the first coordinate of B (sI - A)-1, which we now compute using Cramer's rule. The first coordinate of B (sl A) can be written as
number, we first note that A -B
- e, where 3 = - since is purely imaginary. By Cramer's rule, we have that
B (sl TsI ()-A)== A)
e
(sl (sl
A) e (-sl sl A+)-1B,
1
v3
V5
2n
v 3
8
v 15
33(n 1)
det
55
3
/ /5(2n - 1)
12n
e, (-ssl
A')
v3
v5
V2n 1
0
15
/3(2n 1)
v 2det
0
0
3
/ v5(2n 1)
7
Obviously, the denominator is / v2I1 (j - s). We compute the numerator by solving recurrence.
We use D, to denote this determinant. Hence, have D 1 and D2 =1 8. To compute D,,, we expand the last row and obtain
V3
V5
V2n-3
2-8
15
/3(2n-
3(2n 1)
D, D,=(-1)nt) / /2n det
3-8
V5(2n-
(5(2n11)
(n-=))D--
n 1
(2n
(2n 22n33)m (2m
4

---

Published as a conference paper at ICLR 2024
To compute the determinant of this submatrix, have
v
5
V2n -3
22n -
2
V15
3(2n -3)
32n 1)
det
3
V5(2n- -3)
55(2n- 1)
:
0
0
- 1
(2n 3) 1)
V2n
v3
v5
22n
/3(2n-1)
2-s
V5
/3 3)
---1)==2
/5(2n-11
0
3-8
V V5(2n-3
:
(2n 32
0
= --n1)=-\2n-111,ppp
Hence, combining the two equations above, obtain the following recurrence:
Dn -2n - 1)D + (n s)Dn-1 = (-n+1 Dn
It is then easy to show that
(- (--)---)--s) D, ()12
8) (-1)- s).
j=1
Putting everything together, have
(-1)a-1 11
1 (-1)- 1-1)"~111 8)
B ( 1
conj
(18)
2 (
2
2 +
Now, it remains to study the term (sl - A) in eq. (17). Since it is a vector of length n, we study component-wiss and the derivation similar to the one above. To begin with, we component p that we wish to study. Then, by Cramer's ' rule, we have
eg (sl A) - B
8 +
1
0
3
8 + 2 V15
V3
0
5 V7
8 4 3
V5
0
det
221
v
VT
0
/2n
3(2n
552n 1)
/2n
+
v3
2
V5
v 15
22et
V7
2II
v/5
22
V 332
V 5(2nn
(2p 12 1-
s+n
Clearly, we have that the denominator equal v2IT + s). To compute the numerator, we first subtract the pth column from the first column. This shows that the numerator is equal to
2
0
0
15
+ 3
0
det
V21
0
(19)
/3(2n
1)
V 152m
22
5

---

Published as a conference paper at ICLR 2024
We can then subtract / times the (p 1)th column of the submatrix from the first column, showing that the numerator is equal to
+
5
0
v
7
0
8(8 1) det
5(2m
22n
+ n
Continuing in this manner, we have that the numerator is equal to
2pp
0
0
0
T
s+p+1
0
8(8-1) ........6 (s
v2p+3
v 32
s+p++2
Lv2n-1 (2n-1)(2p-1)
(2n (2p2-+2
s+1
= =22p- 1s(s
... (s p + () + 1) + ts)
Hence, we have
e (sI - -)1-A)-11 22p 1 5)
(20)
V2 +
Note that the expression above does depend on n. Combining eq. (17), (18). (20), when CDPLR DDiag
eyy we have
(s)
GDiag (s)
=
cp
T
[sI A
A
B (sI
B
2
2
1
(-1)a-1
(j-s)
2
211 || U++)
22p 1=0 8 - 9)
2
1)
11
21 - (s + j)
(21)
2 (4t)
(-1)"
18--
22p
(J+8)
(++5)
V2 1+ (-1)a-1 (i-s)
و- (1+8)
This completes the proof of the lemma.
D PROOF OF THEOREM 2
In this section, prove Theorem 2. The idea is to locate the last spike in the figure of GDiag (see Figure 1) and control the height of its peak by lower-bounding the denominator of eq. (16).
Proof of Theorem 2. Fix 2 p. Define Sn, by
si)
8n
max
& 2 A(s) = si
is real and
(22)
+ si)
Note that this set is finite because A(s) we have that
bo; thus, its supremum is attained. Therefore,
(-1)
| 2n
sn
In + 8ni 8n
+ 8
8n
12 1= + +,,)
(j + 8,2)
n + 8
(23)

---

Published as a conference paper at ICLR 2024
In what follows, show that Sn Q(n2)s Then, combined with Lemma 1, have that as
n + 0C,
2p
1
& (8n (Snni-GGing((tt (8n1)] (DiingSnn)) 2 1p- + p sn (In (1m--,17 sn -8,) V n² + s 8n
(1))
= / nn2 + s
Sp.
n²
we can show that 8n = (2( then we have that DD (s,, Gring (8, 2(11 and does not converge to zero. To this end, first rewrite the expression into
sn
8nd
sni-j
8
f
A(sn) = 8nd
exp
- 12
arctan
Snitj
n 8n
8n
(24)
Since arctan. = (xx) as 0, if we assume, for contradiction, that 8nk = on1 2) for subsequence 8ng of on. then we must have that
arctan
arctan
- 00 as - 00.
SnE
max mxx 28
We pick some
index
7
z p large enough such that ΣØ aretan(j snn))
Σ5=1 arctan(j/ max ( 28n,)) 2
2m. Hence, as increases from Sng to max 28ne J. the
angle of the unit imaginary number
ny
i
exp
i2
arctan
changes by at least 4x whereas the angle of si/(n + si) changes by at most T/2. Hence, the winding number of the curve
arctan 2
T:8
exp
-2
5 [8nys (1g,, 28n H},
Tig + si
non-zero. That we must have an E (Sma. max { [n.. 28 }) such that the angle of A(s) equal to modulo 2x, but this is contradiction because Snt s. Hence, we have {2(2)).
E PROOF OF THEOREM 1
In this section, we prove Theorem 1. Since the proof very involved, we provide some intuition here. In Figure 1, we observe that for sufficiently large as |s| increases, the difference between the two transfer functions, GDPLR GDiag goes through three stages. In the first stage (i.e., the pre-spike stage), the large spikes have yet developed. In this stage, increases, |GDPLr Diagl decreases uniformly. the second stage (i.c., the spike stage), the spikes start to occur. This the stage in which do not get uniform convergence. However, by carefully controlling the locations and the total measure of the spikes, we can show that when the Fourier transform of a fixed input function with sufficient decay is multiplied with GDPer Ging: its integral the second stage vanishes linearly as n - co. Finally, after the last spike, we enter the third stage (i.e., the post-spike stage). In this stage, |DDPer GDaag enjoys rapid decay. In what follows, we carefully analyze the three stages separately to prove Theorem 1.
Proof of Theorem 1. Let satisfy the assumptions in Theorem Without loss of generality, we assume ú(s) vanishes on (-o0, 0J because the argument would be symmetric for a negative s. Let
1)" IT (1-))
(e-))
B2p
(++,
H₂(s) = (s) GDiag
2 1 1)--1 T1--0-
(1+5)
'We say f(n) (l(g(z)) if there exists constant C 0 such that f(n) Cg(n) for all

---

Published as a conference paper at ICLR 2024
We set (1) en, where is universal constant determined later on, and
8 2 0 A(s) = 88 (-1)" + si) is purely imaginary and Im(A(s)) 2 0
si)
n
max
(25)
By the same argument as in the proof of Theorem 2, we have 8 = O(n)). To compute the integral of H, (s)) on [0, oo), we do each of the three stages, marked by [0, () (1) (8))
and [sh") oo), respectively. Since 2p- is constant appearing unanimously in HH (s)))(s)){ for all n, absorb it into the asymptotic notations in this proof. Unless otherwise stated, the constants in the asymptotic bounds in this proof are universal constants depending only on pt in particular, or
they do not depend
Integrate on the pre-spike stage: Since
1-1)-=
8)
|s| 11 +
+
and (n)) whenever 0 S S sn the denominator of Hn is lower-bounded by constant independent of n. Hence, we have [H,() = O(n on [0, ) Using Hölder's incquality, have
(s)())(())) de < | (=)())))2=(0000))) [iss) (( O(n 7)
(26)
Integrate on the post-spike stage: For 2 sn the denominator of Hn is lower-bounded by constant independent of n. Hence, have [M O(s where the constant does depend on n. Hence, we have
(((ss)
O(s 2- 24 ds O(n 2-40)
(27)
because (2) (n²).
Integrate on the spike stage: To integrate [H, on [sn (1) (2) function by
we first define the angle
si
arctan 2
j
a(s) arg
+ 2
arctan
+2
arctan
71 + 8i
8)
8
si)
si
(mod 2xr).
si)
The importance of a(s) is that when (ss) is close to (2k + 1)TT for some integer ks, we get a spike in the figure of [. We therefore partition the oscillation stage into parts:
S (s E (" s/ (2)) (a(s)
(2k 1)xl 7/4 for some N),
S2 = 1s (2)) St.
The integral on S₂ is studied in the same way as the decay stage:
f (n(())((ss){ ds < Jo(.)) O(s 2-
O(n²)
2--24 24 = O(n -1-20).
(28)
S2
To study the spikes, we first need to derive a simplified expression of the denominator. Fix an E Sy. We let
a(s) = min (a(s) - (2k + 1)
b
and
1 n- 1 (j si)
d(s) = + si
1= (j + si)
8

---

Published as a conference paper at ICLR 2024
Since
(-1)a= (j si)
r(s) =
1i + si)
5² + n2
by the cosine law, we have (see Figure 6)
π a(s)
-d df rrs))² + 4 I (a(s)/2)
cos
2 2
4r(s) (in((()()))
a(s))
a(s))
2 rs)) + sin? [a(s)
d(s)2 = r(s) + 4sin
4r(s)) sin
2
where the last inequality follows from the fact that r(s) 1/2 for sufficiently large constant in the definition of sn 6 Therefore, we have
H (s))((())) 018-2222)) 2-22 d(s) 0( 22-24)
(29)
r(s)2 + a(s)²
where we used the fact that 2//T S sin(x) S for all S r I 1/2. Clearly, we have
52 + n2
5² + n² - 5² V + n²
n 8
r(s)²
O
(30)
+ n²
3 + n2 +
because (()).. To study a(s), we first need to compute a(s). To this end, note that since we assume (n), there exist two universal constants C1, C2 > 0, independent of n, such that
Ci j S arctan ( s C2
1 Isján-1.
Hence, we have
n
a(s) θ
By the intermediate value theorem and monotonicity of a, there are kn (()) frequencies $1.... Sh, between (1) and sn such that a(s;) = (mod 2r) for all < j I kn-
Each sy contained in a connected component s(D) (G) of S1 and S to s(1) That is, we have
(1) < Ek. Sk. Gt. Ekn 8k.-1 11 Sh. 1 < E1 < $1 S1 (
Moreover, there are two universal constants C2 C2 0, independent of or j. such that
nn2 8 S nn
Combined with eq. (30), we have
r(s)² =
14 )
SE st9), 1 i5jsk
where the constant is universal and does not depend on n or j. To integrate H, (s))ū(())) on Si, we integrate it on each of (&,, G)). To do so, we study the value of ass) using the Mean Value Theorem. First, we note that for any given 8j, we have
d a((s))=
Σ
0
=
1
k
ds
1+
F
2
52
x
= -11
R-1
Σ
k
=
n²
= -
$
83
n
TThis is our only requirement of the universal constant appearing in the definition of (1)

---

Published as a conference paper at ICLR 2024
where the constant in the O-notation does not depend on or j. Hence, fixing a $ j S kn and choosing s E (Eg.5t). by the Mean Value Theorem, we have
a(s) = [a(s) a(8;)))))))))))))))) a(s;)]= 011) n 1s 831-
This shows Cj 8j.83 - = On1² ((2) Hence, we have
H (s))(((())) ds O(1 -222 -2 -22))22220))
r(s)2 + a(s)2 ds
E
s 0012222n 1-)
s5-1 r(s)3
ds
1 als)3 /
1941 a(sjodl)
1
(31)
53
=((2113)
n²
n²
S 001)2220n 19)
8- =0 100)
14 +
44
Suppose q > 1/2 and let 0' -
1/2. Then, we have
kn Σ
An ds Σ
/ H (si)ú(s))) ds ds =
Hn (s))
1424 -2-44
(32)
< 0
where the constant in the last O-notation only depends on p. Combining eq. (28) and (32), we have that when q 1/2, it holds that
(H)())((ss) ds O(n
(33)
Put everything together: Combining eq. (26), (27), and (33) and applying Parseval's identity, obtain
[|DDLLR -Vriagll Ia
llýDeI. -FDingllz HH, (s))(((s)))
(1)
(s))((()))) ds (H((s)((s)) ds [ (n()))((s)) ds O(n
This completes the proof.
F NUMERICAL EXPERIMENTS ON THEOREM I AND
In this section, we explain details and show supplementary results for the experiment in section (see Figure 2). The first experiment examines the behaviors of the DPLR system and the diagonal system given a single Fourier mode as an input. By doing so, we observe the "numerical unstable modes" the S4D model. This corroborates Theorem 2. Then, we compare the two systems using two different input functions: exponentially decaying function and the unit impulse. We will show that the smoothness condition in Theorem is necessary and the linear convergence rate is tight.
In each of these experiments, we simulate LTI systems on some continuous input signals. It is done follows: given input signal u(t)? and an LTI system E, fix the step size to be At 10-3 3 For some final time step N, we discretize our input function to obtain a vector (0))), u(Al)..... .., u( N (N)). We then discretize the LTT system bilinearly (see Appendix B) and compute its output on the input u. We call this procedure "simulate", i.e..
y simulate(u, N.
In this section, we let EDPLR, to be the DPLR system with state size of S4 and Eiing to be the diagonal system with state size n of S4D, where we always take C e and =
Since will he scalar-valued in this section, we do not make holdface.
10

---

Published as a conference paper at ICLR 2024
Im
7(9)
e4
sin(a(6)
Figure 6: Illustration of the proof of Theorem 1.
S
200
S 322.5
$
500
x100
x10
oulput value
10
diag
DPLA
Antipu value
0.2
diag
output value
dieg DPLA
0
-0.2
16
500
1000
500
1000
1000
time step
time step
time step
Figure Simulated outputs of the DPLR and diagonal systems with cosine wave inputs of different frequencies s. Note the scale of the y-axis when s = 322.5.
F.1 THE DIAGONAL SYSTEM BEHAVES DIFFERENTLY FOR DISTINCT FOURIER MODES
Our first experiment considers the outputs of DPLR, and Diag,n when the input is cosine wave
U. (1) = cos(st).
This function has a dense Fourier mode at frequency We fix = 32 and let change. In Figure 7, we plot simulater Ug, DDFLR,32 10 and simulate(u,, DDiag.32 103) with 200, 322.5, and 500, respectively. We see that when each other at least, they are on the same order of magnitude. However, when of the diagonal system blows up. In fact, this value of is exactly where the spike in the plot of ||G occurs when n = 32. Hence, we visualize the counter-example that shows the divergence Theorem 2.
200 or 500, the outputs of the two systems
close to 322.5, the output
F.2 THE DPLR AND DIAGONAL SYSTEMS CONVERGE THE EXPONENTIALLY DECAYING FUNCTION
To test the function-wist convergence of the diagonal system to the DPLR system (see Theorem 1), we consider the following exponentially decaying function:
ue (1) = H()),
where H = [0,00) is the Heaviside function. The Fourier transform of this function is
0.(3
is
Hence, it is function that satisfies the assumptions of Theorem In the left panel of Fig- show the difference between the simulated outputs ||s |Isimlaate(e.. ate(ue, EDPLR 104)

---

Published as a conference paper at ICLR 2024
simulate(ue: DDigg, 104)| increases. We see that as increases, simulate Epiag. converges to simul ( EDPLR. 10 Moreover, the slope of the curve is roughly -1, indi- cating a linearly convergence rate as n + oo. This matches the theoretical statement in Theorem 1.
Exponentially Decaying Input
Unit Impulse Input
1
0.49
simulation error
simulation error
0.48
0.47
0.46
10-3
0.45
25 50 100
5
10
25
50
100
n
Figure 8:
The difference between the outputs [|simulate(u. EDPER, 10t) simulate (u, EDiag. 10*)|l for difference values of n when is the exponentially decay- ing function ue (left) and the unit impulse signal 6g (right).
In Figure 2, we show the behaviors of the two simulated outputs as n increases. For the exponentially decaying imput function uee the outputs demonstrate a clear pattern of convergence as n OC.
F.3 THE DPLR AND DIAGONAL SYSTEMS DIVERGE THE UNIT IMPULSE
Our experiment with the exponentially decaying input shows that the DPLR and diagonal systems converge on a sufficiently smooth input function. One may wonder, however, the smoothness condition is necessary. To show that mild one is indeed required, consider the Dirac delta function .g. well-known that the Fourier transform of it constantly one:
(s()) 1, E R.
In that sense, &n highly non-smooth its Fourier transform does not decay all. Since dg is distribution rather than a classical function, we cannot sample it directly. However, we can mimic by setting the discrete input to be the unit impulse (1, 0, 0). In the right panel Figure 8, we see that I|simulas e EDPLR. 10*) s simulate increases. We can take closer look in Figure 2, where we plot the two output functions with different state-space dimensions n.. In particular,
EDag... 10*)l does not decay as
see that as n increases, the output of the DPLR system remains the same, whereas the output of the diagonal system becomes more oscillatory. We do not have convergence. The oscillatory behavior be explained by observation in Figure 1: the larger the n, the later the spike emerges. This means that for larger n, the outputs of two systems differ higher frequency (i.e., oscillatory mode).
G IMPLICATION OF THE THEORY: NON- OF THE DIAGONAL INITIALIZATION
The analysis of GDpLR and GDiag suggests the following caveat: while the S4 and the S4D models tend to pertain similar behaviors as n gets large, the diagonal initialization scheme used by the S4D model less robust perturbations the frequency domain (see Figure 1). particular, by eq. (3), for fixed state size input signals with frequency modes dense at the spikes of the plot of G are harder to process for the SSM. In turn, the S4D model is unstable near these modes. This does not happen with the S4 model. Our observation suggests that instead of replacing the ill-posed diagonalization problem with a well-conditioned but distinct one (i.e., the S4D initialization), which creates a large backward error (GDPLR(s) a backward stable algorithm, even if the forward enror (i.e., the miscalculation eigenvalues and eigenvectors) will be large (see section 4).
G GDag (s)l, one should solve the ill ill-posed problem using
12

---

Published as a conference paper at ICLR 2024
True Amplitude
S4 Prediction
S4D Prediction
S4-PTD Prediction
Interpolation
amplitude
0.6
0.5
0.6
0.6
0.4
0.4
0.4
0.4
4.5
0.5
0.5
0.5
0.2
0.2
0.2
20
40
60
80
40
60
80
20
40
60
80
20
40
60
80
Extrapolation
amplitude
0.8
0.8
0.5
100
0.6
0.6
0.5
0.5
0
0.4
0.4
0.4
4.5
-05
0.5
..5
0.2
20 40 80 80
20 40 60 80
20 40 60 80
20
40
60
80
frequency
frequency
frequency
frequency
Figure 9: An SSM trained to learn the amplitude of a signal A sin()(() given its samples. The model is tested using signals {A, sin(s)t)), where Aj and 8j sampled on uniform grid on [0, 1] [0, 100]. For each Aj and 5j. the model predicts an amplitude Aj. which we show in the contour plots. The top row shows the interpolation result, where the functions in the training datasets have frequencies only in s E [0, 40] [60, 100). The bottom row shows the extrapolation result, where the functions in the training datasets have frequencies only in E 10). The figure shows the interpolation and extrapolation results for the S4 model, the S4D model, and the S4-PTI) model (see section 4). We observe that our S4-PTT> model interpolates and extrapolates better than the S4D model. In particular, the S4D model is not stable around s = 80, where the predicted amplitude decreases to -4 when the true value increases from 0 to 1. More quantitative results for the interpolation and extrapolation errors can be found in Appendix G.
We demonstrate on synthetic example that the S4D model, regardless of its size n, is not robust under input perturbation of certain frequency modes (which depend on n.). Our training set contains sinusoidal signals parameterized by a frequency and amplitude A:
(t) = Aj sin(s,t),
where A, E [0, 1] and 8j E Sinterp == [0, 40) [60, 100] for an interpolation problem or 8j E Sextmp := [0, 80] for an extrapolation problem. We sample each input function (z(-)) uniformly on t E [0, 10 and train an S4 model and an S4D model with = 32, respectively. Our goal is to learn s and A from the sequential input. In Figure 9, we plot the model prediction of the amplitude A over test set of signals for which and A are on uniform grid on [0, 100] [0, 1].
Figure shows that while both the S4 and S4D models predict well on sampled domains (i.e., Simerp and Sexmm)) the $4 model is significantly better at interpolating and extrapolating on the unsampled domains. In particular, the S4D model suffers from an extrapolation disaster: for 8 >> 80, as the true amplitude of the signal increases from 0 to the predicted amplitude decreases monotonically with a minimum value less than -4. This happens because |GDiag has a spike around |s] = 83, making the information of [0, 80) impossible transfer to [80, 100). Hence, while the S4D initialization stabilizes the diagonalization process, making the computation more efficient, its underlying state-space model is unstable near certain Fourier modes, impairing its robustness (see also section 5.2).
We quantitatively evaluate this point in Figure 10. More specifically, define four domains by
Sinterp 00, 40] U [60, 100],
Sextrap = [0, 80],
Unterp [o, 100] Smeerp (40, 60), Uextra = [0, 100] Sextrap (80, 100).
We are given training samples only with Simrr (resp. s E Sextra) and we test the sequential model on the entire domain [0, 100). Given a test set, we measure the mean-squared error of our model. We uniformly sample the test set from (resp. s E Uextap) to evaluate our models' performance on generalization to unseen data in the seen domain, and their performance performs even better on the seen domain (i.e., Sierp or Seutag)). ). but its interpolation and extrapolation capabilities are much worse than those of the S4 and our S4-PTD models.
Simerp (resp. E Sextrap) and E Unterp
interpolation (resp. extrapolation). We see that the S4D model
13

---

Published as a conference paper at ICLR 2024
Interpolation Errors
Extrapolation Errors
100
10
error
10.
S4, St
error
10
S4, &
S4D, 8
S4D. St
S4D.
SAD. U.
S4-PTD, S
S4-PID, S.
10
S4-PTD.
10
S4.PTD, Ug
10
10
100
10
epochs
epochs
Figure 10: The interpolation and extrapolation errors of predicting the amplitude sinusoidal signal made by the S4, S4D, and S4-PTD models. Each curve shows the mean-squared test error of one model on either the seen domain, S or Se, or the unseen domain, U; or Ue- The yellow curve for the S4-PTD model and the red curve for the S4 model almost overlap in the extrapolation problem.
H PROOF OF RESULTS IN SECTION 4
In this section, present the proof of Theorem 3. In addition, we introduce a probabilistic statement of the eigenvector condition number of a matrix perturbed by a random Gaussian matrix.
The proof of Theorem 3 is classical forward error analysis, but to maintain the best result, we need to explicitly compute the resolvent of A
Proof of Theorem 3. For notational cleanliness, in this proof, we define A
Ah, B
Bи, and
C CDPLR V We have
(S GDPLR(())
B A) - - A E -C-
B (sI -)-1-1-6 (sI A E) C
I B ( A) (s
E) |C
where, by result Demmel (1992), we have
(sxI A)- (sI
E) 1 S E (sl 1)- + O
(sI
A) (34)
We set
C1.1
0
0
8
0
0
C2
0,22
0
V3
2 8
0
0
63,1
C3,2
C3,3
sI A)
V5
V15
3 3-8
LCn,1
Gr,
Ca.
/2n 33 1)
55(2n 1)
n-8
Then, fixing column and a row j 2 we have
CEAVV 1 1 1 -))
-s)
=0,
(35)
Ck.i 2j + +1+ + =0.
(36)
Multiplying (35) by / 2j 1/ /2j 1, we have
Σ
ck. 2
v2j - Gj - -s) ( =0
22 +1
(37)
Am
14

---

Published as a conference paper at ICLR 2024
Subtracting eq. (37) from eq. (35), have
C Cj + +1 - C₁ v2j -2 (j - 8) + +0+1,(+++++ +1 = 0.
After simplifying, get the recurrence relation
1
22i 1 +1
C i-8 s
C1+1, =
8) (i)+)
(j + 1-+8-110- 1j
C/+1,
Cj j2it1.
(j
+ 1) 22
Solving this recurrence relation gives
Cke = (-1)) 2i V2 /2k
I (e+
k2i+1.
s) 8) 22i +
l (e
Since s is purely imaginary, we have
c+
8-8 1.
Therefore, we can control the size of Cj.s by
22 T Tv2k
| + li + 8)
-2
[Ck, j s| i 1 - s| k - - s| |k -s| - - s) 1k s sl²
k21+2.
Clearly, this value is maximized when = 0. Hence, we have
(2i 1)(2k 1)(2 1)
4i
|ck.i 2
(k - 1)2222 - (1: - 1)2k
Note that this inequality holds also for the case when k = -it 1. Now, we have
n k-1
4i
n
[ A- (sI A) -
(k 1)2k
7
k=2
2(k - 1)k
+2 2 1l +
1)2k
The result follows from eq. (34).
In section 4, we show the effect of a "best-case" perturbation scheme on the eigenvector condition number. In this appendix, we present a probabilistic statement of the eigenvector condition number in the "average case." Our result is heavily based on Banks al. (2021, Thm. 1.5). To simplify our statement, Given square matrix M, we define its eigenvector condition number to be
Keig (M) inf kVV),
where ranges over all invertible matrices such that M = VAV for some diagonal
Theorem 5. Given any matrix A E C* perturbation size E (0, 1), and spectral radius R > 0. Let Gn E C xm be the Ginibre matrix and let Ω be the event that the spectrum of + eG,, is contained in Dn, the disk centered zero of radius R. Then, we have
R
E Weig A + <G < ||A| PP(()))
Proof. By Banks et al. (2021, Thm. 1.5), we have that
R2n 2
E
<(X)
< |||A||
With slight abuse of notation, the letter i here stands for real-valued index instead of the imaginary unit.
15

---

Published as a conference paper at ICLR 2024
where Ai.... An eigenvalues of A cGn, and k(X,) is defined in Banks et al. (2021). When dg E DR(0) for all S S n2, we have
Kiegg (A + EG 5 x(A.)².
1=1
Hence, this shows
Rtn2 6
- E Keig (A + +
P(2) + Σ
(AED (1
22 P(N)) s ||al1
n
We are done.
Comparing Theorem 5 to Theorem 4, we note that the bound in Theorem 4 is slightly better than that in Theorem 5. However, the Gaussian perturbation in Theorem 5 is problem-independent and be generically implemented, whereas it is not necessarily easy to identify the perturbation in Theorem 4.
I NUMERICAL EXPERIMENTS ON THEOREM 3 AND 4
1.1 THE RELATIONSHIP BETWEEN E|| AND THE TRANSFER FUNCTION PERTURBATION
In the proof of Theorem 3, we used the inequality between the matrix spectral norm and the Frobe- nius norm:
[|(II A)
(s1 (11-A) (sI-A)-
In practice, this estimate is rarely sharp, given that s1 - A is a dense matrix. Another non-sharpness the average case comes from the inequality
(( - A)- (sI
E) C) s | (sI A) (sI E) -1 1|-1
To understand the average-cast perturbation of the transfer function, we conduct a simulation, where sample E and C randomly but restrict that || 0.1 and |||| = 1. We then compute the maximum error between the perturbed transfer function and the unperturbed one. In Figure 11, we observe that increases, instead of increasing logarithmically, the maximum error G Grealle loc averaged over all trials decays quadratically, i.e., ||GDplR - Grealloe = O 2) Hence, in the average case, obtain better empirical error estimate compared to the worst-case error estimate in Theorem 3. We remark, however, that ||DPplR Greullso not a relative error because while fix [E|| S G, the norm of the state matrix ||A #1 increases
1.2 THE RELATIONSHIP BETWEEN E|| AND KEIG (A
The performance of our perturbed model is heavily based on two things: the perturbation size |E] and the condition number of V H. The former value controls the difference between our initialization the nnown-to-be-gook HiPPO initialization, whereas the latter one controls the unfairness when transforming the states via V H. By Theorem the condition number x(V,, H should depend linearly on || and depend sub-quadratically on the state space size.
In this section, we present numerical experiment that investigates the relationship between these three values. To do so, we solve the optimization problem in (11) with different state space dimensions n and values 7. We then record the size of the perturbation and the eigenvector con- dition number of the perturbed matrix. In Figure 12, we see that the eigenvector condition number Reig AAr) depends polynomially on both the state space dimension and the relative perturbation size ||]|||| A/l Numerical values are reported in Table and 3. Using the data, one can compute where p 0.87. Hence, we are doing slightly better than the theory of Theorem 4. Another surprising observation that be made with little bit computation is that we normalize A to have a spectral norm of 1, then the eigenvector con- dition number Keig (A,r) does not depend on n at all. This is much better than the bound proposed in Theorem 4.
that we have Keig (Ai = O([ [[][]]] [|
16

---

Published as a conference paper at ICLR 2024
2
||GOLLR - Gren oc
10
10"4
10
14
25
45
61
82
110
149
The perturbation error ||GDpLr - Greall for different state-space dimension n. For each n, the matrix Aи is perturbed by 100 randomly sampled matrices E, respectively. The red curve shows the average error among the 100 trials. The horizontal axis is on the logarithmic scale and the green reference line has slope of 22
Figure 11:
In(iiill ( (A#))
6
4 66 88 |n(|/ Anll))
2
4.5
3
In(n)
Figure 12:
The relationship among the state space dimension the relative perturbation size ||||/||A A []. and the eigenvector condition number Keig (Aig(AA))
DETAILS OF EXPERIMENTS IN SECTION 5
In this section, we provide the details of the experiments presented in section 5.
?
10
102
103
104
105
106
107
n
8
4.40e0
8.62e0
1.73e
3.5le1
7.12e1
1.45e2
2.96e2
16
6.59e0
1.32e1
2.69e1
5.53el
1.14e2
2.35e2
4.86e2
32
9.98e0
2.02e1
4.16e1
8.63e1
1.79e2
3.72e2
7.75e2
64
1.52e1
3.12e1
6.45e 1
1.34e2
2.80e2
5.84e2
1.22e3
128
2.34e1
4.82e1
1.00e2
2.09e2
4.37e2
9.14e2
1.91e3
Table The eigenvector condition number Keig (An) when the optimization problem eq. (11)
solved with different values of n and 7
17

---

Published as a conference paper at ICLR 2024
10
10a
103
10a
10s
106
107
8
2.81e0
1.16e0
4.78e-1
1.98e-1
8.24e-2
3.45e-2
1.45e-2
16
6.77e0
2.86e0
1.22e0
5.18e-1
2.22e-1
9.50e-2
4.09e-2
32
1.62e1
6.96e0
3.00e0
1.30e0
5.62e-1
2.45e-1
1.07e-1
64
3.89e1
1.68e1
7.32e0
3.19e0
1.39e0
6.11e-1
2.69e-1
128
9.37e1
4.07e1
1.78e1
7.80e0
3.42e0
1.51e0
6.65e-1
Table 3: The perturbation size
|B]] when the optimization problem eq. (11) is solved with different
values of n and 7-
Task
Depth
#Features
Norm
Prenorm
DO
LR
BS
Epochs
WD
A Range
ListOps Text Retrieval
256
BN
False
0.
0.002
50
80
0.05
(le-3,1e0)
256
BN
True
0.01
16
80
0.05
(1e-3,le-1)
128
BN
True
0.
0.004
64
40
0.03
(1e-3,1e-1)
Image
128
LN
False
0.1
0.01
128
2000
0.01
(le-3,le-1)
Pathfinder
512
BN
True
0.
0.004
64
300
0.03
(le-2,1e0)
Path-X
128
BN
True
0.001
20
100
0.03
(Ie-4,le-1)
Speech
128
BN
Truc
0.
0.01
16
40
0.05
(1e-3,1c-1)
Table 4: Configurations of the S4-PTD model, where DO, LR, BS, and WD stand for dropout rate,
learning rate, batch size, and weight decay, respectively.
J.1 DETAILS OF THE EVALUATION OF OUR MODEL IN THE LONG RANGE ARENA
To compare our perturbed models with the diagonal S4D and S5 models, adopt the same model parameters used in Gu et al. (2022a) and Smith et al. (2023) but possibly change the training parame- ters, such the learning rate, number of epochs, batch size, and weight decay rate. For choosing the perturbation matrix, we again solve the optimization problem in eq. (11). Instead of allowing to be an unbounded positive tuning parameter, we require that is large enough so that |E|] | I 0.1. This improves the worst-case robustness of our model (see section 5.2). We provide the detailed con- figuration of our S4-PTD model in Table 4 and that of our S5-PTD model in Table In particular, note that the first two columns of Table 4 are almost the same those in Gu et al. (2022a)² and the first four columns Table match those in Smith et al. (2023) - these are model parameters. The only remaining non-trivial thing is that in the Path-X task, we start with batch size of 32. We half the batch size after epoch 30 and epoch 60. By making the batch size smaller, we improve the generalization power of our model.
Task
Depth
H
P
DO
LR
SSM LR
BS
Epochs
WD
ListOps
128
0.003
0.001
0.05
Text
256
192
12
0.1
0.004
0.001
50
40
0.07
Retrieval
128
256
16
0.002
0.001
32
20
0.05
Image
512
384
0.1
0.0055
0.001
50
250
0.07
Pathfinder
192
256
8
0.05
0.0045
0.0009
64
230
0.05
Path-X
128
256
16
0.0018
0.0006
32
0.06
Speech
96
128
16
0.1
0.008
0.002
16
40
0.04
Table 5: Configurations of the S5-PTD model. parameter labels.
See Smith al. (2023) for the meaning of the
The only exception is that in the Path-X task, we half the number of features in order to reduce the compu- tation time. This only simplifies our perturhed model.
18

---

Published as a conference paper at ICLR 2024
Model
16kHz
8kHz
InceptionNet (Nonaka & Seita, 2021)
61.24
05.18
ResNet-1 (Nonaka & Seita, 2021)
77.86
08.74
XResNet-50 (Nonaka & Seita, 2021)
83.01
07.72
ConvNet (Nonaka & Seita, 2021)
95.51
07.26
S4 (Gu et al., 2022b)
96.08
91.32
Liquid-S4 (Hasani et al., 2023)
96.78
90.00
S4D (Gu et al., 2022a)
95.83
91.08
S4-PTD (ours)
96.04
91.53
S5 (Smith al., 2023)
96.52
94.53
S5-PTD (ours)
96.87
94.49
Table 6: Test accuracies on Speech Commands classification. indicate the highest test accuracy among all models for each task. indicate the highest test accuracy within the comparable group.
We use the boldface number to
We use the underlined number to
J.2 DETAILS OF THE ROBUSTNESS TEST OF THE DIAGONAL MODEL AND OUR MODEL.
In the robustness test presented in section 5.2, we train both an S4D model and an S4-PTD model. Our models have 4 layers, 128 channels, and each layer contains an SSM with 32 states. The perturbation matrix in the S4-PTD model computed by setting y = 0.03 in eq. (11). From Fig- ure 3c, it can be seen that the perturbation thence computed has magnitude of roughly 10% of the magnitude of A We fix a universal discretization step At (see Appendix B) for all channels. We leave the training dataset and the validation dataset unchanged, but we add 10% of noises in the form of cos(325.41) 4l) to the test dataset. The frequency 325.4 is chosen at one of the sensitivity regions of the diagonal SSM when n = 32. We train both models for 50 epochs and report the evolution of the training accuracy, the test accuracy on uncontaminated data, and that on noisy data. As mentioned in section 5.2, the matrix A is not trained (but see Appendix K for an example where A is trained).
J.3 DETAILS OF THE ABLATION STUDY OF OUR MODEL
In section 5.3, we train models with different perturbation sizes to solve the SCIFAR task (Krizhevsky et al., 2009). Our models have the same architecture those in the sensitivity test (see section 5.2). We set the batch size to be 64 and the learning rate to be 0.001 for SSM parameters and 0.01 for other model parameters. These are common setups that are adapted from the original S4 and S4D papers. We use the parameter in eq. (11) to control the size of the per- turbation E||. We set ? 10-4 10-3 10 and train the S4-PTD model for 200 epochs to lear classifier. These correspond to the first 14 points in Figure 3c, where we report both the test accuracy of the model and the eigenvector condition number at initialization. Since setting y small does not help redueing fieig all the way down to 1, the smallest condition number possible, to obtain the rightmost point, we perturb A by random symmetric matrix with large norm.
K SUPPLEMENTARY RESULTS OF THE EXPERIMENTS
K.1 FULL RESULTS ON ADDITIONAL DATASETS
In addition to the LRA dataset, we test our S4-PTD and S5-PTD models on the 35-way Speech Commands classification task (Warden, 2018), asking the model to identify a word in recording from a pool of 35 words. The quantitative results are shown in Table 6.
K.2 SUPPLEMENTARY RESULTS OF THE ROBUSTNESS TEST
In the robustness test presented in section 5.2, the state matrix A is fixed throughout training, and the non-robustness of the S4D initialization is subsequently observed at a particular frequency. In this section, present supplementary results that aim to clarify the following two points:
19

---

Published as a conference paper at ICLR 2024
S4D Transfer Function
S4-PTD Transfer Function
15
Accuracy = 10%
200
10
150
GsLPD (si)
[G (si)
100
5
50
0
0
100
102
100
102
15
Accuracy = 40 %
200
M
10
150
(Gs-rD(ss))
[GsD (si)
100
5
50
0
0
10
102
100
102
15
Accuracy = 55 %
M
200
[Gs+rD [
10
150
|Gs (si)
100
5
50
0
0
00
102
100
102
15
Accuracy = 70%
200
M
150
[GsaPPD (si)
10
[GsaD(st) (st)
100
m
5
50
0
0
0
102
10
1
Figure 13: The evolution of the transfer functions in the S4D model (left) and the S4-PTD model (right). We observe that the spikes in the S4D model are preserved during the training while the S4-PTD model never develops any spike.
20

---

Published as a conference paper at ICLR 2024
1. As discussed in section 5.2, the perturbation used is the "worst-case" perturbation that makes S4D fail. In Figure 4, however, we showed that not all Fourier-modk perturbations have essential effect. In fact, we revealed a strong correlation between the effect of the perturbation at particular frequency and the location of the spikes in the transfer function of the S4D initialization. We provide more details of this experiment in this section.
2. While the state matrix A is not trained in the robustness test in section 5.2, present different experiment where we train the matrix A as usual. We then observe the evolution of the transfer function. We will empirically see that training does not solve the robustness issue of the S4D model.
In the first experiment (see Figure 4), we use the same setting as the one presented in Appendix J.2 but we contaminate the test set with a varying frequency s. From Figure 4(a), we observe strong correlation between the effect the perturbation and the location of the spikes the graph Griag In particular, frequency located such spike drives the test accuracy to below 30%. This cor- roborates our discussion that the spikes in the graph of Diag impairs the worst-case robustness the S4D initialization. We also show the analogous results for the S4-PTD model in Figure 4(b), where the size of the perturbation is chosen to be the same as the one used for LRA tasks (see appendix I.1). In this case, the performance of the S4-PTD model is relatively stable across all different modes of perturbation because the transfer function Geen is relatively smooth.
Next, we empirically study the behavior of the LTI systems during training. To do so, we again train a 4-layer SSM to learn the SCIFAR task, where all layers are either S4D layers or S4-PTD layers. Instead of fixing the state matrix A, we also train this matrix. In particular, we set the learning rate of the state matrix A to be 0.004 and the learning rate of the other model parameters to be 0.01. We also apply a weight decay rate of 0.05. This is standard practice training an SSM. We then watch the evolution of the transfer functions of the LTTI systems as the test accuracy increases. In Figure 13, we plot the evolution of both models as the test accuracy increases from 10% (i.e., random guessing) to 70%. While there are many copies of LTI systems in a 4-layer, multi-channel SSM, we randomly select one and show its transfer function.
remark that other randomly selected LTI systems all that while the heights of the spikes in the transfer
show similar behaviors. From Figure 13, function of the S4D LTI system decay slightly throughout training, they never vanish. In fact, the rate of decay gets smaller throughout training. The locations of the spikes also change, albeit by small amount. Combining Figure 13 with Figure 4, we can conclude that training does not completely fix the robustness issue of the S4D model. In comparison, the transfer function of the S4-PTD model never develops any spike during the training process.
21


Reviews:
[{"review_id": "Y4qExAvEc6", "reviewer": "Reviewer_Xrm4", "summary": "The authors study the State-space models (SSMs), which have shown promising results on long range sequence tasks. To speed up SSMs, diagonal approximation has been considered, which works well in practice but lacks theoretical guarantees. In this work, the authors propose a generic, backward stable perturb-then-diagonalize (PTD) strategy. The theoretical analysis shows convergence guarantees and proves the stronger robustness of the proposed method. The empirical evaluation also shows strong performance of the proposed method with the compared baselines.", "questions": "1. Can the author elaborate more on why the proposed method can outperform the ones without diagonalization? This seems to be quite interesting.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The proposed method has theoretical guarantees for convergence and robustness.\n2. The empirical result shows solid performance of the proposed method.", "weaknesses": "1. The writing is quite dense and can be improved if the authors make it more self-contained.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ou5URGZr9H", "reviewer": "Reviewer_eMiw", "summary": "This paper presents a deep theoretical analysis of various aspects of S4 models.\n\nFirst, it provides a much more fine-grained analysis of the difference between the S4-DPLR and S4-Diag models, the two most well-known instantiations of S4 models (more precisely with the \"HIPPO LegS\" initializations), than prior work, including\n- Exact analysis of the difference in transfer functions\n- This allows providing non-asymptotic analysis of the convergence between the DPLR HIPPO init and its diagonal approximation by dropping the rank-1 term. As opposed to prior work which only showed it asymptotically in the limit of state size $n$\n- This also leads to better understanding of what type of input perturbations lead to divergence between the approximation\n\nThe technical machinery leads to a very simple but well-motivated solution to the stability issues of S4D, by perturbing the $\\mathbf{A}$ matrix before diagonalizing. This method presents modest empirical improvements over the original methods.", "questions": "End of Sec 3.1:\n> Moreover, for every $n \\ge 1$, zooming into the last spike...\n\n1. is the constant magnitude a conjecture, or easy to show from Lemma 1? How were the final spike locations derived? Seems like it should be possible to find the exact location for every $n$ and calculate the magnitude of the spikes from Lemma 1, although maybe the calculations are actually difficult.\n2. I'm confused how to interpret left side of the plot in relation to the right: are the final yellow/orange/blue spikes supposed to actually go up to magnitude $\\approx 1$ (same as the purple spike), but the sampling of the plot x-axis is too coarse grained to get the exact location of the spike so is not rendered precisely?\n3. The text claims the last spike is located at $|s| = \\Theta(n^2)$ but it seems like larger $n$ correspond to earlier spikes (this seems important as it is related to a later question)\n4. Is there a limiting result of the nature of $G_{\\text{Diag}}(n \\to \\infty)$? These plots seem to suggest that the width of the spikes also decreases, but the rate is not clear. Depending on the nature of this limit, it's not actually clear to me that $|G_{Diag}(si)|$ does not converge uniformly: the domain is bounded from below ($s \\ge 0$) so the location of the spike can't move too fast; if the spikes also have width that does not decrease too fast, then it seems like it actually could converge uniformly? On the other hand if the domain is transformed logarithmically as in Figure 1, then I would believe that $G_{Diag}(e^s i)$ does not converge uniformly, so maybe that's what you mean? (Also, is there a reason Figure 1 only plots frequencies $\\ge 1$?)\n\nIn Section 3.3:\n> the outputs of the two systems diverge given the unit impulse (i.e., the Dirac delta function) as the input\n\n5. I'm not sure of the accuracy of this statement, or at least am confused how it relates to results from the S4D paper. If I understand correctly, the outputs of the two systems given the unit impulse should just be the impulse response or the \"continuous convolution kernel\" $K(t) = C e^{tA} B$. However the S4D paper indicates that this does converge, at least pointwise. (See S4D Figure 2, and the corresponding reproduction: https://github.com/HazyResearch/state-spaces/blob/main/notebooks/ssm_kernels.ipynb) To my understanding those results are generated in the same setting as Appendix F.3, Figure 9 in this submission, but the results look quite different.\n6. Additionally the associated discussion says \"The oscillatory behavior can be explained by our observation in Figure 1: the larger the $n$, the later the spike emerges. This means that for a larger $n$, the outputs of two systems differ at a higher frequency (i.e., a more oscillatory mode)\". But Figure 1 seems to say that larger $n$ corresponds to spikes at lower frequencies (which is related to another question above)\n\n7. Section 3.4: I think the description of the experiment does not state which $n$ is used; based on Fig 1, it seems like it is fixed to $n=10000$?\n\n8. Theorem 3: how tight is the upper bound empirically? This result says that as $n$ grows, even if the total norm of the error matrix is controlled (hence the entries decrease), the total output deviation still increases with larger $n$. In practice, does the output deviation increase with state size?\n\n9. How is the hyperparameter $\\gamma$ set for each dataset? It doesn't appear in the hyperparameter table. From Figure 3(c) it seems as if the results can be somewhat sensitive to this parameter.\n\nOverall I think this paper is a very interesting and technical deep-dive into the theory of S4 models.\nI am happy to increase my soundness score and overall score after discussion of some of these questions.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. Extremely clearly written technical background and presentation of technical results.\n2. The discussion of the relationship of S4-DPLR and S4-Diag from a functional analysis and transfer function perspective is strong and substantially contributes to the theoretical understanding of these popular models. The theoretical derivation of their properties and validation through synthetic experiments is well-done.\n3. The main motivation of the perturb-then-diagonalize (PTD) being justified because only backwards error (instead of forwards error) matters is a nice idea.\n4. The empirical results are strong and ablations are well-motivated.\n\nOverall, the paper has high originality, quality, and clarity.", "weaknesses": "The only potential weakness of this paper is the practical significance to the machine learning community. To my knowledge, in practice most applications of S4 use the diagonal variants to no noticeable detriment, and also often use the simpler S4D-Linear initialization to which the present theory has less practical implications for. Put another way, while technically strong, it is not clear whether the ideas can lead to truly new capabilities for machine learning models.\n\nOn the other hand, this paper does provide a new direction toward revitalizing interest in the HIPPO theory and other (perhaps yet undiscovered) flavors of SSMs, which may still provide practical benefits in more specific regimes (e.g. limited data or when a strong inductive bias towards the memorization interpretation of HIPPO is desired). And as a mainly theoretical paper with strong contributions towards the theory of these popular models, in the context of this paper this is not a major weakness.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "eCPIgilYTe", "reviewer": "Reviewer_faKE", "summary": "This paper presents a theoretical and empirical study of the frequency-domain implications of various approximations common across deep state space models.  The authors make the connection between the structure of approximation, non-normality of operators, and the non-uniformity of the approximation across frequencies.  This pathology is tested across several synthetic examples. This insight is used to propose a new initialization scheme, which appears to ameliorate the degeneracy.", "questions": "**Q.1.**  This is a very general question, but something that I cannot wrap my head around.  I invite the authors to try and help me understand.  (The remedy to this will also affect my perception of the weaknesses) \n\nThroughout the text this is an analysis of (e.g.) HiPPO initialization, but the results presented are for trained models.  I am surprised that the pathology predicted by the initialization is preserved through training quite so neatly.  For instance, in Figure 3, how do you obtain the frequencies that the model should be perturbed at?  Do you inspect the trained model, or the a-priori specified HiPPO matrix?  Each channel in the S4D model should have different frequencies, and so are you corrupting them all?  Or just one? \n\nMy follow-up is why does the approximate diagonalization on initialization lead to more robust trained models?  Presumably $E$ isn’t learned?  Once the model is diagonalized, the initialization is essentially thrown away.  As a result, there is a real disconnect in Figure 3.  It looks like S4-PTD might be _slightly_ worse than S4, but the training curves are ostensibly identical.  But the robustness to perturbations is clear.  Why do these pathological frequencies survive training in S4?  \n\nI think there might be two contributions here that are currently wrapped up in one-another.  The first contribution is an analysis of the frequency properties of a hypothetical (single) SSM.  The second contribution is an initialization scheme for systems of deep SSMs that is derived from some of this insight, but doesn’t necessarily directly solve the initial problem.  \n\nFor instance, I would like to see an experiment sweeping over the perturbation frequency (e.g. combining Figures 1 & 3), and showing that there are certain choice frequencies that destroy performance, as opposed to just a model that is less robust to any perturbations.  I think this is an important distinction that is somewhat predicted by the theory, but is unverified.  I think this is what Figure 6 might be driving at, but to my eye, it looks like the diagonal system is just unilaterally worse than DPLR.  This would also be an empirical validation of Figure 1.  I would like to see the results for an untrained system, a single trained DPLR system, and also for an S4D model (where there are multiple S4 systems in parallel, each with their own characteristics).  \n\nIn that vein, I'd also like to see this graph traced out over training and across different input frequencies.  E.g. If the training data has a lot of power in a frequency where there is a peak in the transfer function, then how does/does the model learn to move that peak away from that frequency?  \n\nI invite the authors to comment on this.  I realize my thoughts on this are a little jumbled, but I think this is maybe emblematic of a shortcoming in the presentation or cross-linking of the ideas.  If the authors want more clarification, I will happily try and restate my concerns.  \n\n**Q.2:**  Why is it relevant that the final peaks in Figure 1 are all (roughly) the same height?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The paper itself is pretty well written, with figures and tables prepared excellently.  I believe the analysis is both a novel and useful insight into this new class of model. The theoretical analysis itself appears correct (although I did not go through the proofs in detail).  The experiments conducted also back-up the theoretical claims (modulo my comments below).", "weaknesses": "I think the paper is pretty sound.  My main concerns are about some of the presentation choices, and the “secret sauce” in the contribution.  I have tried to outline some concrete weaknesses here that, if remedied, would improve the paper.  I have then tried to write out my slightly more high-level question in Questions.\n\n**W.1.: The core method**:  Please can the authors clarify what the method actually entails.  As far as i can tell, the “method” essentially specifies an initialization for S4D, given in (10) followed by diagonalization.  This is used instead of the diagonalization of the Normal component of the HiPPO matrix in S4D.  If this is true, this should be absolutely, concretely clearly stated somewhere.  I think right now it is tied up in the introduction with spectral analysis, PTD, backwards/forwards approximations etc.  \n\n**W.2.: Choice of what the paper emphasizes**:  The authors have clearly done _a lot_ of theoretical analysis.  However, I think the net result is an incredibly simple result – adding a random perturbation makes models more robust.  I really think this result, and any usable empirical analysis of this result, should be highlighted.  Some of the theoretical points should be relegated to the appendix.  If the results are to be believed, then the banner result should be “ignore DPLR/diagonal state spaces and use random perturbations”.  This core message is not highlighted.  \n\nSimilarly, there are some nice experiments in the supplement demonstrating the pathology experimentally, e.g. fig 7-10.  These should be brought up to the main, as they help flesh out the arguments far more than a detailed introduction to the intricacies of the conjugation of HiPPO or some of the theorems.  I think this will help improve the reach and impact of the paper.  \n\n**W.3.: Figure 2 results**:  I find the results in Figure 2 very hard to parse out.  The core point (to my understanding) is that S4-PTD is a faithful approximation of S4.  However, I do not necessarily see why S4Ds predictions _should_ be reminiscent of S4 – it’s a different model! – or that it should be worse in the way predicted.  I think it is coincidental that S4D is so much worse than S4(-PTD) on this task.  I ask the authors to clarify the intuition and results of this experiment.  (See also Q.1.)  \n\nI think it is also somewhat coincidental that S4-PTDs predictions look like S4 – surely different random seeds should result in different predictions? \n\n**W.4.: Missing definitions and exposition**: Lots of terms are basically undefined, e.g.:\n- The introduction of “backward stability” is deferred to the appendix, but is really the core of the method.  There is no discussion in the main really of why forward/backwards even relevant, or why existing methods fall down here.\n- Pseudospectral theory is undefined, nor why adding a random perturbation ($E$) helps.  \n- Why spectral information with a poor condition number is “meaningless”.\n\nThis makes the contribution a little less self-contained.  I would rather drop some of the theoretical analysis from the main to the supplement and bring up more of the core definitions.  This would dramatically improve the impact of the work.  \n\n**W.5.: No conclusion or discussion**:  Not really a weakness, but there is no conclusion, discussion or outlined future work.  It would be nice to see some critical self-reflection on the work.  It would also help you reinforce the core messages you are trying to convey in the paper.  \n\n## Minor weaknesses\n- (e.g.) “section 3.4” should be capitalized.\n- The explanation of Figure 2 on Page 6 is very poor.  \n- use \\citet and \\citep where appropriate.\n- Author style in the bibliography is inconsistent.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DLsrz2WnJu", "reviewer": "Reviewer_X2Rr", "summary": "Explores the diagonal state matrix used by recent linear SSMs such as S4D and S5 which were proposed as an approximation to the original diagonal plus low rank conjugation of the HiPPO matrix explored in the S4 paper. The work analyzes the differences between the S4 and S4D initialization through the transfer function and finds the convergence of S4D to S4 is not uniform. Experiments show that S4D models are sensitive to input perturbations which degrades the model's robustness. A proposed fix is adding a perturbation to the HiPPO matrix to allow for stable diagonalization. The proposed initialization is evaluated on the LRA tasks.", "questions": "Please see weaknesses above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The theoretical analysis seems sound and leads to several interesting insights including implications on the robustness of the initialization of the diagonal models. \n\n- The empirical robustness experiments (synthetic and sCIFAR) help to illustrate the implications of the theoretical divergence \n\n- The proposed fix is relatively straightforward and simple to implement, yet seems to lead to improved robustness results in the synthetic and sCIFAR tasks", "weaknesses": "- I am not convinced by the claim the S4-PTD model outperforms the S4D models on LRA. The LRU paper (https://arxiv.org/abs/2303.06349) reports results for S4D that are much better than the original reported S4D paper results. In addition, the appendix of the current paper under review states that mild hyperparameter tuning was performed for the S4-PTD models. Would mild hyperparameter tuning improve the S4D results as well?\n  - To be clear, I am not claiming that S4-PTD must outperform S4D on all tasks to be a valuable contribution, simply that I think the claim about improving performance is too strong, and I would expect to see more results in a fair comparison to support this claim\n\n- More experiment results on a broader range of robustness tasks would help strengthen the paper. At a minimum, perhaps modified versions of the speech and BIDMC tasks from previous papers in this line of work could have allowed for further evaluation of this method on some form of real world data.\n\n- There seems to be a missing piece related to the training dynamics that I think would strengthen the paper. The theory is related to the dynamics matrix $A$ at initialization, but the $A$ matrices are trained during the experiments. How do the trajectories of the eigenvalues differ during training such that the PTD models are more robust during the perturbation experiments? This should be relatively to simple to track during training.\n\n- While it is true the proposed perturbation method could be applied to any dynamics matrix, no other matrices are explored. Perhaps some of the other diagonal matrices proposed in the S4D paper could have been experimented with. Would a similar approach be beneficial if applied to initialize the effective dynamics matrix of the Liquid-S4 model?\n\n- The paper doesn't really seem to support the claim that the S5 models suffer from the same perturbation \nissue. I believe this is likely the case, but no empirical results are presented that support this. Can an example of S5 and S5-PTD be included in the perturbation experiments also?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)