Autocorrelation Matters: Understanding the Role of Initialization Schemes for State Space Models
Authors: Fusheng Liu
Qianxiao Li
Venue: ICLR 2025 Poster (2025)
Keywords: State Space Model; Optimization;
Primary area: other topics in machine learning (i.e., none of the above)

Abstract: Current methods for initializing state space model (SSM) parameters primarily rely on the HiPPO framework \citep{gu2023how}, which is based on online function approximation with the SSM kernel basis. 
However, the HiPPO framework does not explicitly account for the effects of the temporal structures of input sequences on the optimization of SSMs.
In this paper, we take a further step to investigate the roles of SSM initialization schemes by considering the autocorrelation of input sequences. 
Specifically, we: (1) rigorously characterize the dependency of the SSM timescale on sequence length based on sequence autocorrelation; (2) find that with a proper timescale, allowing a zero real part for the eigenvalues of the SSM state matrix mitigates the curse of memory while still maintaining stability at initialization; (3) show that the imaginary part of the eigenvalues of the SSM state matrix determines the conditioning of SSM optimization problems, and uncover an approximation-estimation tradeoff when training SSMs with a specific class of target functions.

Published as a conference paper at ICLR 2025
AUTOCORRELATION MATTERS: THE ROLE OF INITIALIZATION SCHEMES FOR STATE SPACE MODELS
UNDERSTANDING
Fusheng Liu, Qianxiao Li
National University of Singapore, Singapore fushengêu. wt..auuuusecceee edu, qianxiao@nus edu.
ABSTRACT
Current methods for initializing state space model (SSM) parameters mainly rely on the HiPPO framework (Gu et al., 2023), which is based on online function approximation with the SSM kernel basis. However, the HiPPO framework does not explicitly account for the effects of temporal structures of input sequences on the optimization SSMs. In this paper, we take a further step to investigate the roles of SSM initialization schemes by considering the autocorrelation of input sequences. We: (1) rigorously characterize the dependency of the SSM timescale sequence length based sequence autocorrelation; (2) find that with proper timescale, allowing a zero real part for the eigenvalues of the SSM state matrix mitigates the curse of memory while still maintaining stability at initialization; (3) show that the imaginary part of the eigenvalues of the SSM state matrix determines the conditioning of SSM optimization problems, and uncover an approximation- estimation tradeoff when training SSMs with specific class of target functions.
INTRODUCTION
The state space model (SSM) is a sequence model that has recently shown great potential in long sequence modeling across various applications, including computer vision (Zhu et al., 2024; Liu et al., 2024). time series forecasting (Rangapuram et al., 2018: Zhang et al., 2023) and natural lan- guage processing (Gu & Dao, 2023; Dao & Gu, 2024). a continuous-tine ordinary differential equation h (t) = Wh(() + Ba((), v(t) = Ch(()) + Dz((), where W, W,B, B, C, D are trainable parameters, x(t) is the input sequence, and yt) is the output se- quence. For discrete input sequences, timescale A > O will be introduced to discretize the model. Different from the attention mechanism ( (Vaswani et al., 2017). SSMs are recurrent-basse architectures that treat the input sequence token by token, yet can achieve first-order time complexity on the sequence length through parallelization (Gu et al., 2022b). There are two well known issues for training recurrent-based architectures, the vanishing and the exploding gradi- ent problems (Pascanu et al., 2013). By introducing complex-valued initialization schemes, proper parameterization methods and regularization techniques, recent works demonstrate that SSMs can achieve performance comparable to attention-bat based architectures in terms of both computational cost and sample efficiency (Gu Dao, 2023; Dao & Gu, 2024; Zhu et al., 2024; Yu et al., 2024; Wang & Li, 2024; Liu & Li, 2024; Yu et al., 2024; Bick et al., 2024; Hwang et al., 2024; Wang et al., 2024; Waleffe et al., 2024). However, the theoretical understanding on the roles of the initialization schemes is still lacking and needs to further explored. In this paper, we particularly look into the timescale A and the state matrix W, and we aim to study the following fundamental question
mathematics, SSM layer is defined by
hyperparameter
Given sequential dataset with length L, how should the timescale A depend on L and what is the role of W on training SSMs?
Based on the analysis of continuous-time SSMs, previous works (Gu et al., 2022b;c; 2023) propose the HiPPO framework where W. B are initialized such that the SSM basis kernels {e, e Wt ] orthogonal in [0 oo) with some measure w()), and the timescale A scales as 1/L to capture long range dependencies of sequences with length L. Common HiPPO-basee initialization methods such S4D-Legs and S4D-Lin typically presume that the measure (1)) is exponential decay and

---

Published as a conference paper at ICLR 2025
the discrete input sequences have inherent timescale A that is shared with the model. However, these assumptions are restrictive because exponential decay measures weaken the effects of temporal dependencies in input sequences, and in practice, usually lack prior information about the data's timescale. To address this concern, we take an initial step towards understanding the relationship between the autocorrelation of input sequences and the SSM initialization schemes. Specifically, we focus on the diagonal SSM' (Gu et al., 2022c) where the state matrix W is a complex-valued diagonal matrix. By studying the stability condition for given input sequences E IR,, we find that the connection of the timescale A and the sequence length L is highly related with the spectrum of the data autocorrelation matrix E[xx"] ]. Different temporal dependencies in the imput sequences can cause significant variations in the spectrum of the autocorrelation matrix. For example, when ar is sampled from a standard normal distribution, æ has zero temporal dependencies, and the auto- correlation matrix becomes an identity matrix. On the other hand, if x consists of constant values, the input sequence exhibits full temporal dependencies, and the autocorrelation matrix is low rank. For the state matrix W, our stability amalysis shows that even with a zero real part, i... R(W) O, the diagonal SSM can still be stable at initialization if A is properly set. One benefit for setting the real part to zero is that the learned SSM kernel functions initialization do not exponentially decay, which helps to mitigate the curse of memory (Li et al., 2022). Our convergence analysis indicates that the imaginary part (:W)) plays crucial role in the convergence rate and explains the benefits for complex- valued SSMs compared to realevalued SSMs in terms of the optimization. In particular, the more separated the imaginary parts (ww) are, the faster the convergence. When considering both approximation and optimization, we characterize an approximation -estimatioot tradeoff when the target function has closely spaced dominant frequencies. Then well separated ((w) values lead to fast convergence, while achieving a good approximation requires close imaginary parts.
To summarize, the main goal of this paper is to provide a theoretical understanding on the effects of three specific hyperparameters: the model timescale A, the real part of RRW). and the imaginary part of the state vector (W).. These components are connected as a data-dependent initializa- tion scheme for SSMs. First, for any given sequential dataset. we can estimate its autocorrelation. Using this information, we can apply Theorem 3.1 to initialize A, taking into account both data autocorrelation and sequence length. Second, if the true input-output mapping is represented by an underlying linear functional, often referred to as a memory function, that exhibits a long memory pattern, our second theory, detailed in Section 3.2, suggests that initializing with zero real part can help mitigate the challenges posed by long sequences. Finally, the third theory introduced in Sec- tion 3.3 discusses an approximation-estiaatee
tradeoff' that arises when the true memory function p" features closely spaced frequencies. If we can accurately recover p" from the sequential data, we can then initialize the imaginary part based on the dominant frequencies of p", thereby finding an optimal balance informed by theoretical insights. Accordingly, our contributions are as follows:
In section 3.1, characterize the dependency between the timescale A and the sequence length L by taking into account the autocorrelation of the input sequences. Even if the eigenvalues of the state matrix W have zero real part, the stability condition on the magnitude of the output value at initialization can still hold with an appropriate setting of A.
section 3.2, we show that the real part of the eigenvalues of the state matrix W determines the decay rate of the SSM kernel functions. Allowing the eigenvalues of W to have zero real part at initialization can significantly increase the model's effective memory and help mitigate the curse of memory for fixed-length tasks that require long-term memory.
In section 3.3, we prove that the conditioning of SSM optimization problems is determined by the separation distance of the imaginary parts of the eigenvalues of the state matrix. Well-separated imaginary parts induce faster convergence, whereas closely spaced ones lead to slower conver- gence. This explains the benefits of complex-valued SSMs over real-valued SSMs. Furthermore, it uncovers an frequencies in the frequency domain.
tradeoff when the target function has close dominant
2 PRELIMINARIES
In this section, we briefly introduce the diagonal SSM and the problem setting we consider through- out this paper. Specifically, we consider the following single-input single-output (SISO) diagonal-
To simplify the analysis, we omit the skip connection by letting 1 0.
2

---

Published as a conference paper at ICLR 2025
SSM built in the complex number field C and then cast into the real number field R,
d dt M(1 Wh(() + bur(1), v(() Ric A()): L 12 O,
(1)
where R(()) represents the real part; xxt)) is input sequence from an input space" x Co (R20. R); v(() IK is the output sequence at time 1; k(t) E C is the hidden state with k((0) = 0: W Cm b.c C.C. E CDhh are trainable parameters. In particular, the state matrix W diag(w),- (11 Wmm
a diagonal matrix. To simplify the analysis, we omit the skip connection matrix D. Following the training setup in Gu et al. (2022c), the read-out vector follows standard normal distribution and the read-in vector in (1) is fixed an all-one vector at initialization without training. Under these settings, the input-output relation in (1) is explicitly given by the integral
y(t) Lf RCcc (ee"))((( - s)ds,
where E C is the state vector that contains all the diagonal entries of the state matrix W, and the function R(c² exx)) is called the memory function or the kernel function.
Discretization. To handle discrete sequences, we follow (Gu et al., 2022c) to use the zero-order (ZOH) hold method for discretization. Then given a timescale A > 0 and any discrete sequence (ro..... 11) R with length L. for = 1,2,.................. L, the ZOH method induces a model output
m
Amy
Aws
Y1 = R
oye (w)
2o + +
Tt-1-
(3)
Wj
Remark 2.1. Here we focus on the SISO model with ZOH discretization. It is also possible to extend to the multiple-inpu multiple-output (MIMO) case, by noticing that the MIMO output is essentially a linear combination of several single-input single-output (SISO) models. As a result, we can extend results to the MIMO scenario by examining each SISO model individually along with its respec- tive input-output mapping. However, our theory are not directly applicable to other discretization methods (e.g. bilinear method), which involve different matrix forms for the model's 's input-output mapping (see Appendix D), and require different techniques to yield theoretical insights.
In the following section, we tackle the problems related to the initialization schemes of SSMs that were introduced in the Introduction. Specifically, we will explore the following questions:
1. Timescale Initialization: How should we initialize the model timescale A for fixed-length tasks to enhance the training of SSMs? Is the previously used scaling A 1/L. a universal approach?
2. Real Part of the State Vector: What role does R(w) play? Can we initialize R(w) to be zero, and what benefits might arise from a zero real part?
3. Imaginary Part of the State Vector: What role does 3ww) play? What advantages do complex- valued SSMs offer compared to real-valued SSMs?
3 MAIN RESULTS
In this section, we present our main results by focusing on three initialization parameters A, R(w) and (ww) respectively. Specifically, in section 3. we rigorously characterize the relationship be- tween the timescale A and the sequence length T. in terms of training stability at initialization by taking into account data autocorrelation. In section 3.2, we demonstrate that allowing the state vector's real part to zero can prevent exponential decay in the SSM kernel function, thereby miti- gating the curse of memory in certain scenarios. In section 3.3, we explore the relationship between the convergence rate and the separation distance the state vector's imaginary part. In particular, uncover paproiimation-ssiiaatoo tradeoff for class of target functions.
3.1 RELATIONSHIP BETWEEN A AND L
In this subsection, we derive a stability condition for the ZOH-discretized diagonal SSM (3) when the state vector's real part R(w) is non-positive. From both theoretical and numerical perspectives,
A linear space of continuous functions from Rzo IR that vanishes at infinity.
3

---

Published as a conference paper at ICLR 2025
Test loss for copying task i.d.d data
Arad [xx
] for different processes
1 A
for resized Image datasets
10a
3
101
rat
10
10
rand
2
10"
2
10-3
109
2'
21
26
28
215
212
2
2
Figure 1: (Left) Training a diagonal SSM (3) on copying task using i..d. data with dimension of 128. We vary the minimal timescale Amin 1/L.1/v and the maximal timescale Amax 1/L, 1/VL, 0.1 w.r.t sequence length L. (Middle) The maximal eigenvalue of the autocorrelation matrix 16[x ] on different random processes of z. (Right) The maximal eigenvalue of E B|xx ] on sequential image datasets sMNIST and SCIFAR 10 with different resize rates varied from 0.5 to 4.
rand,
0.5
ou, 9i(w)
10P
109
100
10-1
10
10-1
A=L
d-1-1
10
d=1-- 1-1
d=1-1
10-1
10
2P
221
2
2
221
2
2
2
2
2
2
2
2
Figure 2: The expected magnitude of the SSM output value synthetic sequences with different autocorrelation. The real part R(w) = -0.5 follows the common practice and we consider four dependencies between the timescale A and the sequence length L.
we demonstrate that the dependency of the model timescale A on the sequence length L is strongly influenced by the data autocorrelation. To start with, we prove the following theorem that provides an upper bound on the magnitude of the model output value.
Theorem 3.1. Consider ZOH discretized SSM (3) with timescale A > O and R(wz) S for 1..... m. Suppose that the input sequence [20....,IL is sampled from a unknown distribution in IR,,, and the read-out vector is from i.i.d. standard normal distribution. Then we have
Be (wi) 3 Amax (E|x:x T))
where Xmmax() (-) represents the maximal eigenvalue.
The proof is provided in Section D. In practice, the hidden state size m is often much smaller than the sequence length L (Gu et al., 2023). Given this, we focus on fixing the hidden size m and investigating the relationship between the model timescale A and the sequence length L. We see that Theorem 3.1 connects the model timescale A with the sequence length L in terms of the data autocorrelation matrix E [xx If we have normalized the sequences such that ||||2|||]]
1, then simple observation is that S Amax Aaaxx((222 [xxx"" I) < L because Tr(E,ar [rv 'I) = L. This indicates that the have different dependencies
I. based on the temporal dependencies. For example, when the elements in the sequence are uncorrelated with each other, exhibits zero temporal dependencies, and the autocorrelation matrix is identity matrix with Amax (Bfe:: [] 1. In this case, A should scale I/VL to ensure training stability. On the other hand, when æ is a constant sequence (1, 1,..,, The autocorrelation matrix then becomes rank-1 1 matrix with Amax Amaxx(EEaa (E[rz }) = L, implying that A should scale as 1/ L. Additionally, this upper bound is applicable for all cases where R(w) < 0. As the real part R((w) approaches zero, the exponential decay rate of the SSM kernel slows, resulting in a tighter bound. Specifically, when the real part is zero and the input data lacks temporal dependency, the bound becomes tight up to constant factor. This occurs because the data autocorrelation matrix E|xx 'I simplifies to diagonal matrix under these conditions. Given specific task with an input sequential dataset, we then initialize the model timescale A 0(1 VT>mmax (Efx:)) JD).
maximal eigenvalue of the autocorrelation matrix
, 1), then x exhibits full temporal dependencies.
4

---

Published as a conference paper at ICLR 2025
R(w)
rand,
R(()
rbt, R(() 0.0
10a
10'
10a
10"
107
10/
104
10
10a
10P
10f
2-11-
10
10a4
109
1011
10-1
10"2
6-2-km
20
23
2
2
231
214
23
28
2
2
20
24
2
Figure 3: The expected magnitude of the SSM output value on synthetic sequences with different autocorrelation and different dependencies between A and L. The real part R(w) is set to be zero.
SMNIST, Rlw) 0.5
SMNIST, N[W)
SCIFARIO, S8/w) 0.5
SCIFARIO, 58/M0
107
-433
10a
Am
10
10'
L-133
10t
10a
10
A-1
A-1
10-2
100
10-1
10
10P
2
27
21
2
27
21
27
211
21
27
2.1
21
Figure 4: The expected magnitude of the SSM output value on sequential image datasets with dif- ferent resize rates (ranging from 0.5 to 4) and different dependencies between A and L.
Remark 3.2. Theorem 3.1 applies for the final output mode, while in practice, there are also some other output modes and the analysis on the stability condition is case by case. For example, for the pooling mode y = +. Σ v². Also, in this paper we consider fixed-length tasks, i.e., all the sequences the varied-length case, we may first cluster the sequences into several groups, and then increase the model feature dimension to manage varied sequence lengths separately. For example, if the sequence length alternates between L, and L2, then we can double the feature dimension and initialize the model separately for the first and second halves, allowing the model accommodate two fixed- length datasets simultaneously.
we can use Theorem 3.1 to get same upper bound for E,,> [y"]-
with the same length. For
Numerical experiments on Amax (B(rx: and E(vī)- To validate our theory, we conduct ex- periments on the exact values of the magnitude of the model output and B[xx Specifically, we consider both synthetic and real sequential datasets in both negative and zero real part cases. For synthetic datasets, we consider Gaussian process with mean ( and autocovariance function B[x,,,]] = K (i. By restricting ii, as the autocorrelation matrix. In this paper, we choose 4 Gaussian processes with different autoco- variance functions and plot their maximal eigenvalues. The autocovariance functions for "ou, iid, rbf" are (ii,j) exp(-1) (i/l), exp(-li- -ip---i-l-- for "rand" is given by Σ where Σ is a random matrix with i.idd uniform distributed entries in [o, 1]. As Figure (Middle) shows, different processes have varying dependencies of Amax (EE) [x. on L ranging from ((1) to (L)). For the i..d. case, Amaax (E(2:x 'I) is not always 1 in Figure (Middle), which is because we use the sample autocorrelation matrix to replace the expected auto- correlation matrix. For real sequential datasets, we choose resize the MNIST dataset (LeCun al., 2010) and the gray CIFARIO dataset (Tay et al., 2021) with resize rates [0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4] and the flatten the images to sequences. More experiment details are provided in Appendix B. We record Amax [B] [xx based on the entire training dataset. As shown in Figure (Left). the maximal eigenvalue scales (almost) linearly with sequence length across the resize rate for both sequential MNIST (sMNIST) and sequential CIFAR10 10 (SCIFAR10) 10) datsets. Additionally, we plot the relation- ship between the magnitude the model output value and sequence length by varying the timescale A A= [1- 00.75 -U33 L -0.25 In Figures 2 and 4, when R(w) -0.5 (following the setup in Gu al. (2022c; 2023)). the magnitude E[y?. remains stable for both synthetic and resized image datasets for all decay rates of A. When R(w) = 0. Figures 3 and 4 demonstrate that for the 'rand' process, A = L is stable. For the 'iid,' 'ou,' and 'rbf' processes, A = L- 0.75 stable. This indicates that bound in Theorem 3.1 effectively characterizes the relation between A and I. for
then the autocovariance matrix exactly the same
j12 /E) respectively. The autocovariance matrix
5

---

Published as a conference paper at ICLR 2025
i for different real parts
Test loss for different real parts
Decorrelated sMNIST
1.0
Ri/i 0.5
R(() 0.5
Ri/w]
[[()
90
0.5
80
Wiww - 0.5
0.0 v wy Mow
= 0
50
100
200
400
50
100
1
Epoch
Epoch
Figure 5: (Left) Training a diagonal SSM (3) on task that requires long-term memory. The learned memory function " effectively captures the spike in long-range dependencies. However, it struggles do so when the real part is negative. (Middle) Test loss on the long-term memory task when initializing R(w) = 0 and R(w) = -0.5. (Right) Test accuracy for training a diagonal SSM on decorrelated sequential MNIST dataset with different real parts at initialization.
R(w) O. Moreover, as shown in Figure 4, for the sequential image datasets, A should scale as 1/L to ensure stability when R(w)= = O: otherwise, the magnitude increases with sequence length. This finding aligns with the empirical results in Gu et al. (2022c) that A should scale as 1/LL to ef- fectively capture the range of dependencies for length L. But their theoretical reasons are based on Fourier analysis of continuous-time SSMs and do not explicitly account for the data autocorrelation.
Experiments on copying task with different timescales. We tested the performance of the diag- onal SSM (3) on a copying task with various dependencies of A on L. It is worth noting that, as discussed in Jelassi et al. (2024), SSMs struggle with the copying task because the model's state dimension needs to scale linearly with the sequence length memorize all the input tokens. How- ever, the limitation highlighted in Jelassi et al. (2024) pertains to the length generalization training an SSM with short sequences and then testing on longer sequences will fail if the hidden size does not grow linearly with L. Here, we focus on fixed-length task, where both training and test sequences have the same length. We find that, with an appropriately initialized timescale, SSMs can effectively handle the copying task even with small state size. In this paper, use diagonal SSM with fixed state size m = 32 to learn a copying task on i.i.d. data with a dimension 128, and the timescale A R 128 We vary the minimal and maximal timescales (Amins Amax) with different dependencies on L. From Figure (Left), we see that the combination (Amin- Amax) = (1/L, 0.1), which is commonly used in practice (Gu al., 2022c; 2023) to train real datasets, consistently performs worse than setting Amin 1/ /L. This stable scaling is in line with theoretical suggestions for i.idd data. Therefore, the data autocorrelation is very crucial for us to get a good initialization scale on the timescale. More experiment details are provided in Appendix B.
3.2 BENEFITS OF ZERO REAL PART
In this subsection, we investigate the henefits of initializing R(w) = 0 for tasks that require long- term memory. In previous works (Li al., 2021: 2022), shown that recurrent-based models suffer from the curse of memory in both approximation and optimization when there is long-term memory in the target. For example, we consider using a diagonal SSM (3) to learn input-out relationship given by a real-valued target function p² such that
ye = =p + P L= 1,2,
The objective function is given by the squared difference between the model output yr. and the corresponding label vī. Then in special case when the input sequences have zero temporal depen- dencies with E[rx II,, the expected mean squared error is given by
El v2l21 -|ø 1=
where 7 is vector
A
Wi - Aw,0
R E AA - (,)_,yydu,(L-11)) that rep- w,
resents the model's memory, and p² - = (pö. ......,,i-)) 1)_-)). Therefore, a well-trained SSM means that the model memory function matches with the target function, i...
Am;
R
c, Amge
= = e= =0 L....., 1.
W'i
6

---

Published as a conference paper at ICLR 2025
Recovered memory functions for decorrelated SMNIST
2.5
2
2.5
2.5
0.0
0.0
0
-2.5
-2
-2.5
500
500
500
L
2
2.5
0
0.0
-1
-2.5
-2
-2
500
500
500
500
L
Figure 6: Recovering the memory function on the decorrelated sequential MNIST dataset by solving a linear equation X p = Y. where X E N is the collected sequence matrix, Y KI0 is the one-hot label matrix, and * is the convolution operator. Then E R _2x10 has 10 channels and we plot the scaled function VIp each channel to show the underlying memory patterns.
Then we can see that the curse of memory happens when the target function p² has a sudden spike in a very long distance. For instance, consider a shifting task that requires mapping an input sequence (Fo....,,L-11. to shifted sequence (0, ................................ O, 2o). In this task, the target p² is (0, ................................ 0, 1). which is challenging for an exponentially decaying SSM kernel p to capture long-term memory when R(w) < 0. However, if we allow the real part to be zero at initialization, then ء does not un- dergo exponential decay. As a result, we can potentially avoid the curse of memory, even for long sequences, in this scenario. It is worth noting that in this paper,
do not consider stable pa- rameterization to ensure R(w)= < O strictly during training. This approach means R(w) is likely to become positive during training. Our goal is to allow the model to learn directly from the data without introducing new variables, such as reparameterization methods, which could complicate the analysis. Otherwise, it would be unclear whether the improvements are due the zero real part or the introduced reparameterization method. Our experiments demonstrate that initializing with a zero real part still helps enhance training. even without a stable parameterization. This suggests that, despite the potential optimization stability challenges during training, a zero real part can be beneficial for training on certain tasks.
Experiments on the benefits of zero real part. To validate the effectiveness of having zero real part, we conduct experiments on both synthetic and real datasets that require long-term memory. For the synthetic task. we use i..d. sequential data to easily visualize the expected error via the memory function. The goal is to learn an input-output mapping from (z0....,I14
to zo + IL-1, which requires the model to memorize both the first and last token. The target memory function p" is (1,0,... 0, 1). In our setting, the sequence length L is 128, and the hidden state size m is 32. As shown in Figure 5 (Left) and (Middle), the SSM with zero real part outperforms the case with a negative real part. It is evident that by initializing R(w) = O, the learned memory function is able to capture long range dependencies. For the real-world task, we utilize the sequential MNIST (sMNIST) dataset. Before training, we preprocess the entire dataset with a linear transformation to decorrelate the training sequences, resulting in an autocorrelation matrix that an identity matrix. We recover the underlying target memory function by solving a least square problem min, |X x I| where X E IR 5000000784 is the collected sequence matrix, Y E |550000 10 is the one-hot label matrix, and denotes the convolution operator. The recovered target memory function p E R 784x10 has 10 channels. To illustrate the underlying memory patterns, we plot vLp for each channel in Figure 6. We observe that for the decorrelated SMNIST dataset, the underlying memory function exhibits a sudden spike at a long distance, implying the curse of memory when R(w) < O. This observation is confirmed in Figure 5 (Right), which shows that initializing R(w) = 0 outperforms the case with a negative real part. We also apply our methods to the Long Range Arena (LRA) benchmark (Tay et al., 2021). which features six diverse tasks ranging from text to image processing. Given that we lack precise knowledge of the memory function for each LRA task, we opt to initialize fraction of the real part as zero and compare this setup to the default S4D model (Gu et al., 2022c). In particular, for each single layer of an L-layer S4D model, which has a feature dimension of and state size of m, there are d state vectors w E C At initialization, fraction p E [0, 1] of
7

---

Published as a conference paper at ICLR 2025
ListOps
Text
Retrieval
Image
Parhfinder
PathX
Avg
Baseline
60.47
86.18
89,46
88.19
93.06
91.95
84,89
Initialize 10% of R R)) be 0
61.44
88.05
90.73
89.11
95.58
97.55
87.08
Ratio for R(w) 2 0 after training
1.29%
1.99%
2.58%
4.31%
4.31%
4.04%
3.09%
Table 1:
Test accuracy for training S4D on the LRA benchmark with different fractions of zero real part at initialization.
these state vectors randomly selected to have their real parts set zero. When p the training proceeds following the baseline setup. As increases, the model starts with more zero real parts. To ensure credible results, we exclude any reparameterization method on the zero real part, allowing the model to adapt from the data during training. This approach isolates the impact of the zero real part on performance without confounding variables introduced by reparameterization. All models were trained with 6-layer architecture, maintaining the original S4D training conditions as specified in the work by Gu et al. (2022c). We present both the test accuracy and the ratio of non-negative real part parameters to the total Ldm real part parameters upon completion of training in Table 1. We can see that, initializing an appropriate fraction of state vectors with zero real parts enables the model to outperform the default S4D configuration. Importantly, even post-training, some non-negative real parts persist, suggesting that the model retains stability and effectively adapts to the data. We provide more experimental details in Appendix B. In Appendix B.2, we add ablation studies on the gray-sCIFAR dataset with varied fractions of zero real part at initialization.
3.3
IMAGINARY PART INDUCES AN APPROXIMATION-ESTTMAATTT
TRADEOFF
In the previous subsection, we show that the real part R(w) is related with the long-term memory when training SSMs. In this subsection, we focus on the imaginary part (w). We will demonstrate how (w)) influences the conditioning of the SSM optimization problem within convex framework. Additionally, from an approximation standpoint, we reveal an approximation-estinn estimmttin tradeoff that arises when training SSMs with particular class of target functions.
Convergence analysis. Here we consider the continuous-time SSM (2) and assume that the read-out vector e is Rm This real-value assumption necessary the current version to get an theoretical estimate for the spectrum of the induced Gram matrix because use the Gershgorin circle theorem to prove Theorem 3.4. This theorem is applicable only when the matrix has dominant diagonal en- tries. If the vector e is complex- complex-valued, the resulting Gram matrix would not be diagonal-dominane rendering the Gershgorin circle theorem ineffective. Suppose the true input-output relation is given by some real-valued target function p² (s) E [0 oo) A L oo) with (t) = J p" (s))xxtt - s)ds. We use the squared difference between the SSM output y(t) and the target output y (t) at some terminal time T > 0 averaged over input distributions, which can be written
a(c,a) := E, (v(T)) v(T))
(4)
To make the theoretical analysis amenable, we assume that ar(t) is sampled from white noise, i.e., a(T s)ds = dW, where W, is the canonical real-valued Wiener process. Then by Itô's isometry (Proposition C.2), the expected risk (4) can be rewritten as [(c, w) Jsd (c R (e"*) p² (s)) ds. In the practical training, the sequence length very long and thus we take T - oo investigate the effect of long-term memory. To study the effects of the state vector initialization, we consider the following convex optimization problem where w is fixed.
arg min Ce cERm
Σ Re - (s((s ds.
(5)
From the perspective of function approximation, the HiPPO framework (Gu et al., 2020) initializes w such that the SSM basis kernel functions {) are orthogonal in 1 oo) w.r.t. some measure wx). In this paper, we discover the effects of the state initialization on the optimization problem (5). Let c* be one of the solution of the convex problem (5), then c* is stationary point (pp((s) (s)ds where G E mm is Gram matrix with
that satisfies Ge = =oo "(e"")pp
[G /
R(e") Re wK w
(6)
8

---

Published as a conference paper at ICLR 2025
K(G) for different initializations
k(G) for S4D-Lin
k(G) and (mx)(M) for 54D-Lin
2
10
K(G)
10"1
23
100
10a
(M)
107
m 1000
100
S4D Real S4D-Lin
10
2
10-2
2 Hidden dimenslon m
2a
2a
2-1 2a
2a
2
2
2
Scaling factor for 3(w)
Ratio víε
[G] for S4D-Real and S4D-Lin with different hid- den size (Middle) x(G)) for SiD-Lin with different m by varying scaling factors of the imaginary part ((a). (Right) x(G) and approximation measure Vmax(M) (in the approximation-estinnn tradeoff part) for S4D-Lin by different ratios of model frequencies and target frequencies &.
Figure 7: (Left) Condition number x(G) := (G)
Therefore, the spectrum of the Gram matrix G determines the numerical stability and convergence rate of optimization algorithms for solving the convex problem (5). We show in the following proposition that when R and all wy are distinet, when C and all the imaginary parts (ww) are non-zero and distinct, then G is positive definite.
Proposition 3.3. Let w. = Uj + i. with aj 12y E IR. j= m. afll vj = 0, ie., w E T then G is positive definite given that all aj are distinct. If U2 are all non-zero, ie., E Cm, then G is positive definite given that all v/ are distinct.
The proof is based on the argument of Vandermonde matrix, and we provide details in Appendix E. ready to study its spectrum. In the follow-
Given that the gram matrix G is positive-definite ing theorem, we show that for complex-vallee valuue SSMs, the gram matrix G can be well-conditioned provided that the imaginary parts 3ww) are well separated.
Theorem 3.4. Let Amin (Gin(G), Amaxx (G) be the extreme eigenvalues of G defined in (6), and let coth(x) = 2. Suppose that wj = -0.5 + i Uj for U, E I. and we define the separation distance == minz#* |01 0z/. Then o >> O. we have
3π
1T (5)
3x
1.19
coth
< Amin S Amax
coth
(
46
12 48
The proof is based on the Gershgorin circle theorem, with details provided in Appendix F. The setup Wy 00.5 + Fy follows the configurations in Gu et al. (2022a;e; 2023). This theorem shows that the Gram matrix can be well-conditioned when the separation distance 5 large. One example is that for the commonly used S4D-Lin initialization (Gu et al., 2022c), Uj- = Ā. Then the separation distance 8 n. Numerical calculations show that 0.2 < Ami S (G)
v2, meaning that G is well-conditioned for any hidden size m, and its condition number has a uniform upper bound w.r.t. m. Note that x coth(z) 2 and is increasing on [0, oo), which implies that the bound 3m for Amin(( (G) non-trivial when 48 coth (7) < 1.19. By numerically solving this inequality, is sufficient to have 8 > 2.3. However, Proposition 3.3 suggests that as long as 5 > O, the positive- definiteness of G is guaranteed. This indicates gap between the lower bound and the actual minimal eigenvalue, which we leave for future research.
Real vs complex. We can now compare real-valued SSMs and complex-valued SSMs in terms the conditioning of the convex optimization problem (5), which is determined by the condition number G. For real-valued SSMs with the S4D-Real initialization (Gu et al., 2022c), where wy = -j. we have Gj, J In this case, G a Hilbert matrix, whose condition number grows exponentially with respect to its size m (Todd, 1953). For complex-value SSMs with w, = -0.5 + iv,, Theorem 3.4 indicates that if the separation distance 8 remains uniformly large with respect to then G can be well-conditioned even for larger values of m. For S4D-Lin initialization, we already know that 0.2 < Amiu (G) & Amax (G) < 22 by the above argument. Therefore, unlike real-valued SSMs, the condition number of G in the complex-valuee case can remain well-conditioned even for large m, given that the imaginary parts are well separated. This difference illustrated in Figure 7 (Left), where compare the exact condition numbers for S4D-Real and S4D-Lin. As the scaling factor

---

Published as a conference paper at ICLR 2025
of the imaginary part increases, the separation distance also increases. Figure 7 (Middle) shows that the Gram matrix G for S4D-Lin becomes better conditioned, validating Theorem 3.4.
tradeoff. Despite the fact that complex-valued SSMs with adequately separated imaginary parts (ww) enhance the conditioning of G. we cannot simply initialize w with widely separated (w). This is because 3 (w)) determines the frequencies that the SSM can capture, and misaligned frequencies relative to the target p lead to a large approximation error Le For example, suppose that the target memory function p" (s) = e == cos(fs) with č E R". Let = 005 + iv for R then we have
peprooimmation-sstmmaane
Mc, fee = [e0=-1000000 (s) (= --1 10'1/ (f,e c cos( (s)ds) = = Mc,
where M E <Rmx is given by
e cos(Es) cos(Es) T do (%.0 00 === cos(us) do) 0-1 (cooss)))) cos(Es) ds
We can see that the maximum singular value max (M) of M determines the approximation error. Now, let's consider limiting case when V/ = μj with μ + oc. According to Lemma C.5, we know that = I.mm a scaled identity matrix, possesses the best possible conditioning. Furthermore, & is finite, then as f - oc, fo e cos(u;s) cos(Ex8) ds = 0, indicating that the worst approxi- mation error Jo p² (s) ds. On the other hand, if we aim to minimize the approximation error, might align the frequencies such that e = However, when the target function comprises closely spaced frequencies Ei. Sm. such alignment may cause G to have large condition number (as per Theorem 3.4). Balancing these two aspects reveals an approximation-estiane crucial when selecting an SSM initialization. Numerical evidence for this tradeoff is illustrated in Figure 7 (Right). In this figure, we set ε = 0.1TT with relatively small separation distance = 0.1 and we vary the ratio UE/Ez from 2 to 25. As the ratio increases, the optimization is expected to improve, while the approximation deteriorates. This trend is shown in Figure 7 (Right), where the induced Gram matrix G becomes conter-conditione, whereas the approximation measure m increases. In practice for specific task, if we manage accurately recover the memory function from the sequential data, we can then apply a Fourier transform to identify the dominant frequencies of the memory function. Given state size m, we can greedily select m frequencies that exhibit the largest separation distance from these dominant frequencies to initialize the imaginary part. Our theory suggests that this approach achieves an optimal balance between the tradeoff of ap- proximation and estimation. However, practically, recovering the memory function accurately from the data is challenging, and hyperparameter tuning might be needed to find the optimal balance.
tradeoff, which
CONCLUSION
In this paper, we study the question proposed in the Introduction section, focusing on two initializa- tion schemes for state space models (SSMs): the timescale A and the state matrix W. Regarding the timescale A, we investigate from the perspective of training stability at initialization. Our findings indicate that its dependency on sequence length is determined by data autocorrelation. By analyzing data autocorrelation, we can initialize A to enhance SSM training for tasks involving fixed-length sequences. For the state matrix W. we differentiate between the real part R(W) and the imagi- nary part (W)). The real part R((W) is crucial for capturing long-term memory in temporal data. Allowing for a zero real part can effectively mitigate the curse of memory while maintaining train- ing stability at initialization, provided the timescale is appropriately initialized. The imaginary part ((W) affects the conditioning of the SSM optimization problem. well-separated 3(W) leads to well-conditioned Gram matrix, improving the convergence rate. However, from an approxima- tion standpoint, excessively increasing the separation distance can result in a frequency mismatch between the SSM and the target function, leading to approximation-estimale three components are intricately linked as a data-dependent initialization scheme for SSMs. There several potential future interesting directions. For instance, we have not discussed the effects of gating (Mehta et al., 2023) and model depth on the approximation and optimization of SSMs, which leave for future research.
tradeoff. These
10

---

Published as a conference paper at ICLR 2025
5 ACKNOWLEDGEMENT
We would like to thank the anonymous reviewers for their constructive comments. Q. Li supported by the National Research Foundation, Singapore, under the NRF fellowship (project No. NRF- NRFF13-2021-0005),
REFERENCES
Aviv Bick, Kevin Y Li, Eric P Xing. Zico Kolter, and Albert Gu. Transformers to ssms: Distilling quadratic knowledge to subquadratic models. arXiv preprint arXiv:2408. 10189, 2024.
Nicola Muca Cirone, Antonio Orvieto, Benjamin Walker, Cristopher Salvi, and Terry Lyons. Theo- retical foundations of deep selective state-space models. arXiv preprint arXiv:2402. 19047, 2024.
Tri Dao and Albert Gu. Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. In Forty-first International Conference Machine Learn- ing, 2024.
Albert Gu and Tri Dao. Mamba: dinear-time sequence modeling with selective state spaces. arXiv preprint arX1v:2312.00752, 2023.
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, 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, Ankit Gupta, Karan Goel, and Christopher Ré. On the parameterization and initialization of diagonal state space models. Advances in Neural Information Processing Systems, 35, 2022c.
Albert Gu, Isys Johnson, Aman Timalsina, Atri Rudra, and Christopher Re. How to train your HIPPO: State space models with generalized orthogonal basis projections. In International Con- ference on Learning Representations, 2023.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces effective structured state spaces. In Advances in Neural Information Processing Systems, 2022.
Dan Hendrycks and Kevin Gimpel. Gaussian error linear units (gelus). arXiv preprint 100.v:1606.08415. 2016.
Sukjun Hwang, Aakash Lahoti, Tri Dao, and Albert Gu. Hydra: Bidirectional state space models through generalized matrix mixers. arXiv preprint (rXiv:2407.09941 2024.
Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and eran malach. Repeat after me: Trans- formers are better than state space models at copying. In Forty-first International Conference on Machine Learning, 2024.
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6900 2014.
Yann LeCun, Corinna Cortes, and CJ Burges. Mnist handwritten digit database. ATT Labs [Online].. 2, 2010.
Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. On the curse of memory in recurrent neural networks: Approximation and optimization analysis. In International Conference on Learning Representations, 2021.

---

Published as a conference paper at ICLR 2025
Zhong Li, Jiequn Han, E Weinan, and Qianxiao Li. Approximation and optimization theory for linear continuous-tine time recurrent neural networks. Journal of Machine Learning Research, 23 (42):1-85, 2022.
Fusheng Liu and Qianxiao Li. From generalization analysis to optimization designs for state space models. In Forty-first International Conference on Machine Learning, 2024.
Yue Liu, Yunjie Tian, Yuzhong Zhao, Hongtian Yu, Lingxi Xie, Yaowei Wang, Qixiang Ye, and Yunfan Liu. Vmamba: Visual state space model. arXiv preprint arXiv:2401.10166 2024.
Harsh Mehta, Ankit Gupta, Ashok Cutkosky, and Behnam Neyshabur. Long range language model- ing via gated state spaces. In The Eleventh International Conference on Learning Representations, 2023.
Antonio Orvieto, Samuel Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pas- canu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pp. 18670-266688 PMLR, 2023.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In Proceedings of the 30th International Conference on Machine Learning, 2013.
Syama Sundar Rangapuram, Matthias W Seeger, Jan Gasthaus, Lorenzo Stella, Yuyang Wang, and Tim Januschowski. Deep state space models for time series forecasting. In Advances in Neural Information Processing Systems. Curran Associates, Inc., 2018
Mark Rudelson and Roman Vershynin. Hanson- inequality and uub-gaussian concentration. 2013.
Jerome Sieber, Carmen Amo Alonso, Alexandre Didier, Melanie N Zeilinger, and Antonio Orvieto. Understanding the differences in foundation models: Attention, state space models, and recurrent neural networks. arXiv preprint 17511 2024.
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. In International Conference on Learning Representations, 2021.
J. Todd. The condition of the finite segments of the Hilbert matrix. National Bureau of Standards Applied Mathematics Series, 1953.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, et al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406. 078887, 2024.
Junxiong Wang, Daniele Paliotta, Avner May, Alexander M Rush, and Tri Dao. The mamba in the llama: Distilling and accelerating hybrid models. arXiv preprint arXiv:2408. 15237, 2024.
Shida Wang and Qianxiao Li. StableSSM: Alleviating the curse of memory in state-space models through stable reparameterization. In Forty-first International Conference on Machine Learning, 2024.
Annan Yu, Michael W Mahoney, and N Benjamin Erichson. There is hope to avoid hippos for long-memory state space models. arXiv preprint arXiv:2405.13975, 2024.
Michael Zhang, Khaled Kamal Saab, Michael Poli, Tri Dao, Karan Goel, and Christopher Re. Ef- fectively modeling time series with simple discrete state spaces. In The Eleventh International Conference on Learning Representations, 2023.
Lianghui Zhu, Bencheng Liao, Qian Zhang, Xinlong Wang, Wenyu Liu, and Xinggang Wang. Vision mamba: Efficient visual representation learmin with bidirectional state space model. In Forty-first International Conference Machine Learning, 2024.
Nicolas Zucchet and Antonio Orvicto. Recurrent neural networks: vanishing and exploding gradi- ents not the end of the story. arXiv preprini arXiv:2405.22 20064 2024.
12

---

Published as a conference paper at ICLR 2025
A RELATED WORKS
Optimization of SSMs. Recurrent- becurrent-based architectures are known for two issues: training stability and computational cost (Pascanu et al., 2013). To mitigate these challenges and capture long range dependencies more effectively in sequence modeling, the S4 model was introduced with novel pa- rameterization, initialization, and discretization techniques (Gu et al., 2022b). Recent updates to the S4 model have further simplified the hidden state matrix by using a diagonal matrix, thereby improving computational efficiency (Gu et al., 2022c; Gupta al., 2022: Orvieto al., 2023). Ad- ditionally, regularization methods such as dropout, weight decay, and data-dependent regularizers (Liu & Li, 2024) are employed with SSMs to prevent overfitting, In this study, we explore how temporal dependencies in input sequences impact initialization schemes in terms of opúmization, with a particular focus on the timescale and state matrix.
Curse of memory in SSMs. The "curse of memory" is a newly introduced concept that highlights the difficulty recurrent- bassde models face in capturing long-term memory (Li al., 2021; 2022), and has been discussed in recent works (Cirone et al., 2024; Sieber et al., 2024; Zucchet & Orvieto, 2024). This issue arises due to the exponential decay property of the model's kernel basis functions. A common strategy to parameterize the real part of the state matrix's 's eigenvalues involves stable parameterization (Gu et al., 2022c; Wang & Li, 2024), ensuring stable training dynamics even if the input sequence is infinitely long. However, this stable parameterization constrains the real part of the state matrix's 's eigenvalues to be strictly negative, thereby limiting the model's ability to capture long-term memory. In this paper, we argue that if input sequences have fixed lengths, it is reasonable to set the real part of the eigenvalues to zero by appropriately setting the timescale. This relaxation allows the model capture long-term memory while still maintaining training stability.
B EXPERIMENTS DETAILS
In this section, we provide more experiment details that produce Figure 1, 3, 2, 5, 6 and Table in section 3.
Figure 1 (Left). The synthetic dataset that we use to produce Figure 1 (Left) is i.i.d. sampled from standard normal distribution with dimension 128, i.e., each input sequence of shape (1, I,, 128) where L is its sequence length. We use ZOH discretized diagonal SSM layer (3) with hidden size m = 32, model dimension = 128 to handle the 128 dimensional dataset. We initialize the state vector w by S4D-Lin with real part -0.5. The read-out vector i initialized as i...d. standard normal distribution. We vary Amin and Amax in the SSM layer and use the Adam optimizer (Kingma, 2014) to train the hyperparmeters A, R(w., 3ww), C. without weight decay. The learning rate for A, R(w,, ((w) is 0.001 and the learning rate for c is 0.1.
Figure (Middle), 2, 3. The synthetic datasets that we use to produce these figures Gaus- sian processes with mean zero and varied autocovariance functions E6 B[z;a,] = (i,jj) for i, = 1,2,..., L. Specifically, the 'iid' dataset refers to K Kiū,j) K(i,)) exp(-li- j1/2); the 'rbf" dataset refers to K K(i,j) òزاي; the 'ou' dataset refers to exp( pop(-Tli- 513); and the au- tocovariance matrix for the 'rand' dataset is given by ΣΣ /L where E R T. x. is a random matrix with i..dd entries sampled from a uniform distribution U00 v3.. For all the four synthetic datasets, we have ii, 1. The plot for Figure (Middle) records the maximal eigenvalue of the sample matrix that we fix the data size to be 1000 and vary the sequence length L as plotted. So we can see some deviations between theory and practice. For Figure 2 & 3, we also use the 1-dimensional SSM layer (3) with S4D-Lin initialization on (w) and vary the real part R(w) to be -0. or 0.
Figure (Right), 4, 5 (Right), 6. For the resized sequential image datasets, we choose resize the original images with resize rates [0.5, 1, 1.5, 2, 2.5, 3, 3.5, 4]. Then we standardize the whole images and flatten them into 1-d sequence. For sequential MNIST (sMNIST) dataset, the sequence length is 784,2 and for sequential CIFAR10 (SCIFARIO), the sequence length is 1024r² where is the resize rate. The plot for the maximal eigenvalue of the autocorrelation matrix and the output value are based on the whole training set. We use the 1-dimensional SSM layer (3) with S4D-Lin initialization and vary the real part R(w to be -0.5 or 0 to calculate the output value magnitude. For the decorrelated sMNIST dataset, we choose the original MNIST dataset and the decorrelation transformation is given by centered matrix with a whitening matrix after flattening images. The centered matrix is the mean of the sequential data along the batch dimension, and the whitening
13

---

Published as a conference paper at ICLR 2025
R(w) for synthetic task
R(w) for decorrelated SMNIST
0.5
600
0.0
Value
Count
400
-0.5
200
-1.0
10
20
30
-1.0
0.0
0.8
Index
Value
Figure 8: Behavior of the real part R(w) after training on the synthetic task and the decorrelated sMNIST dataset.
matrix has shape I. x L. The whitening matrix can be obtained by SVD on the data matrix. To train the decorrelated sMNIST dataset, we use 128-dimensional SSM layer (3) with m 32 and GELU activation (Hendrycks & Gimpel, 2016) on the model output, and also apply gated linear unit after the GELU activation. The experiment for R(w) = -0.5 follows the default training setup in Gu et al. (2022c), and for the experiment with R(w) without any reparameterization method. We use dropout with rate 0.1 and apply a decoder layer for classification. We use Adam optimizer with learning rate 0.001 on A, R( (ww) 3(w) and AdamW optimizer with weight decay 0.01 on the rest hyperparameters. For the plot of the memory function in Figure 6, we solve least square problem by taking the pseudo inverse of the sequence matrix X E IR 50000x784 and then get the recovered memory function p.
initialize all the real part R(w) to be
Figure (Left), (Middle). The comparisons on zero real part and negative real part in Figure (Left) & (Middle) are conducted on I-dimensional synthetic dataset. We sample the training and test dataset from i.i.d. standard normal distribution with length 128. The training sample size and the test sample size are both 1000. We use the SSM layer (3) with m 32, S4D-Lin initialization on 3 (w) and initialize the timescale A = 1/v128. We use Adam optimizer with learning rate 0.001 on A, (wr), (ww) and learning rate 0.01 on c.
Table 1. Compared with the default training setup in Gu et al. (2022c) for S4D models, the only difference is that we randomly select a fraction p E [0, 1] over the feature dimension such that the selected state vectors are initialized with zero real part. For these selected vectors, their corre- sponding timescale Ap Re.d is initialized to a constant. As suggested by Theorem 3.1, to ensure the stability. Åo can be chosen to be O(L/ LAmax [E[rz /LL//LAmaxx ID)). Here we do not conduct a prior numerical check on the spectrum of the data autocorrelation,
we simply take an upper bound of Amax (E [ }), which is given by (L) if we have normalized the sequences such that ||||==|||] In that case, Ao (I)L)) In the default training setup (Gu al., 2022c), the model timescale A is sampled from a uniform distribution U U'[Amin: Amax] with Amin ~ (l) Hence, for the LRA benchmark simply initialize Ao as a constant Amin as Gu al. (2022c).
1.
We also summarize the in other hyperparameters for training S4D on the LRA benchmark in Table 2. And we provide the behavior of the real part R(w) after training on the synthetic task (ref Figure 5 (Left), (Middle)) and the decorrelated sMNIST dataset (ref Figure 5 (Right)) in Figure 8. We can see that with zero real part initialization, is possible that there remains some non-negative real part after training if no reparameterization method is introduced. But the improvement on the experiments indicates that the model can learn from the data effectively even without constraining the real part values.
B. ADDITIONAL EXPERIMENTS FOR S4D-LEGS INITIALIZATION
In this subsection, we include more experiment results in Figure 9, 10, 11, 12 for SSMs with S4D- Legs (Gu et al., 2022c) initialization on the imaginary part 3(w). The S4D-Legs initialization is an approximation on the original S4-Legs initialization (Gu et al., 2022b) by taking diagonal part of the diagonal plus low-rank HiPPO-Legs matrix. In Figure 9, 10, 11, plot the magnitude of
14

---

Published as a conference paper at ICLR 2025
D
H
N
Dropout 0
Learning rate
Batch size
Epochs
Weight decay
ListOps
6
256
0.01
32
40
0.05
Text
6
256
4
0
0.01
16
32
0.05
Retrieval
256
0.01
0.05
Image
6
512
64
0.1
0.01
50
200
0.05
Pathfinder
6
256
64
0.0
0.004
64
200
0.05
PathX
6 256
64
0.0
0.0005
16
50
0.05
Table 2: List of the S4D model hyperparameters for the LRA benchmark, where D, H,N denote depth, feature dimension and hidden state space dimension respectively.
lid, RR/) 0.5
rand, (R(W)
0.5
O0, R(())
0.5
rbr. RRW)
0.5
10
10a
10
10a
10a
sty")
10a
10-2
6m
10
A=1-21
Amt
6 100
10-2
A
A-L-113
A
A-1-115
10
2235
1915
10"2
A-1-115
2
21
2
2
2
2
210
21
2
2
2
Figure 9: The expected magnitude of the SSM output value on synthetic sequences with S4D-Legs initialization and different autocorrelation. The real part R(w) 0.5 follows the common practice and we consider four dependencies between the timescale A and the sequence length L.
ou, 31(w)
10
10a
10a
A-1-335
10a
10a
107
10a
10
d=i-1
109
10
10-1
10a4
10a2
221
2.4
2
221
2
2
2
2
24
Figure 10: The expected magnitude of the SSM output value synthetic sequences with S4D-Legs initialization and different autocorrelation and different dependencies between A and The real part R(w) is set to be zero.
sMNIST, RR/W) - 0.5
SMNIST, Rl/w) 0
SCIFAR1O, 3R() - 0.5
SCIFAR1O, R[w) -
10a
10¹
A
10a
10?
Aiit
101
10-1
10a
Am
A-
A
Ami-
10a
10P
Aw/
d- L-335
2
Figure 11: The expected magnitude the SSM output value for S4D-Legs initialization on sequen- tial image datasets with different resize rates (ranging from 0.5 to 4) and different dependencies between A and
15

---

Published as a conference paper at ICLR 2025
i for different real parts
Test loss for different real parts
Decorrelated sMNIST
1.0
91/w) 0.5
Sk/w) 0.5
90
91/w)
9R(w)
80
0.5
70
R(w) 0.5
0.0
9(w)
100
200
400
100
Epoch
Epoch
Figure 12: (Left) Training a diagonal SSM (3) with S4D-Legs initialization on task that requires long-term memory. The learned memory function " effectively captures the spike in long-range dependencies. However, it struggles to do so when the real part is negative. (Middle) Test loss on the long-term memory task when initializing R(w) = and R(w) -0.5. (Right) Test accuracy for training diagonal SSM with S4D-Legs initialization on decorrelated sequential MNIST dataset with different real parts initialization.
Initialize fruction of (ww) to 0
P-U U
0.1
p = 0.2
p = 0.3
p= 0.4
p == 0.5
Accuracy
84.09 (0.47)
84.60
(0.38)
84.23 (0.49)
(0.46)
83.77
83.50
(042)
83.19 (0.39)
Ratio for R[w) > Dafter training
0%
3.90%
7.62%
10.85%
14.34%
17.25%
Table 3: Test accuracy for training 4-layer S4D model in sCIFAR dataset with varied fractions of zero real part at initialization.
the SSM output value given the S4D-Legs initialization for both zero real part and negative real part cases. The experiment settings follow the guidelines we introduce before with only change on the initialization of (w)). We can see that for S4D-Legs initialization, our conclusion still holds in the sense that negative real part is stable initialization for all all the scaling that considered in this paper, while for zero real part, the dependencies of A on L varies for different sequence autocorrelation. We also compare the effects of real parts optimization with S4D-Legs initialization. The results are shown in Figure 12 and we obtain consistent results as the S4D-L.in initialization. One interesting finding is that on the decorrelated sMNIST dataset, the comparison between Figure 5 (Right) and Figure 12 (Right) shows that the S4D-Lin initialization outperforms the S4D-Legs initialization in both zero real part and negative real part cases.
B.2 ABLATION STUDIES ON THE EFFECTS OF FRACTIONS OF ZERO REAL PART AT INITIALIZATION
In this subsection, we conduct ablation studies on the gray-sCIFAR dataset (with sequence length 1024) to evaluate the benefit of initializing the real part zero in multi-layer S4D models, while making minimal modifications. Based on our theory, zeroing the real part can alleviate the curse of memory in scenarios where the memory function exhibits long memory pattern. But since do not have a precise knowledge of the memory function for the gray-sCIFAR dataset, we do ablation studies on the effects of zero real part by varying the fraction p as specified in Section 3.2. For the selected state vectors with zero real part, their corresponding timescale Do E IFd is initialized as a constant 0.001. We use a 4-layer S4D model with feature dimension 128 and vary p across [0, 0,1.1. 0.2, 0.3, 0.4, 0.5). We present both the test accuracy and the ratio of non-negative real part parameters to the total 4 x 32 x 128 = 16384 real part parameters upon completion of training in Table 3. We can see that initializing enables the model to outperform the default S4D configuration. Importantly, even post- traininn some non-negative real parts persist, suggesting that the model retains stability and effectively adapts to the data.
appropriate fraction of state vectors with zero real parts
C AUXILIARY LEMMAS
In this section, we provide the description for Itô's isometry and few auxiliary lemmas that we will need for the proofs of Theorem 3.1, Proposition 3.3 and Theorem 3.4.
16

---

Published as a conference paper at ICLR 2025
Lemma C.1. If R(2) S
then
e
S 1.
2
Proof. Notice that
e 1
So e'ds
So e=
[ds)
So²
ex((2)
ds
So
[ss
1,
|2|
|2|
|2]
which finishes the proof.
Lemma C.2 (Itô's isometry). Let W [0, x 12 + R denote the canonical real-valued Wiener
process defined up to time T O, and let X [0. x - be a stochastic process that is adapted
to the natural filtration of the Wiener process. Then
X₁ aw, E [1%
)
=B
x² dt
where E denotes expectation with respect to classical Wiener measure.
Lemma C.3 (Gershgorin circle theorem). Let be a complex i E ,I...., n.. let R2 be the sum of the absolute value of the non-diagonal entries in the i-th row: R3 = Ejjs asgl Let D(au, R&) S be a closed disc centered at on with radius Rd Then every eigenvalue of lies within at least one of the discs D(asi Ra).
matrix, with entries azy- For
Lemma C.4. For any E R.
+
coth(Tt).
+ 12
212
2t
Proof. This is side result of the Basel problem. The related proof can be found in the Wiki page. We omit it here.
Lemma C.5. For any Uy Uk E we have
(
e ee' cos(u)s) = 2 1+ (ny ux)² + (o, + )2
Proof. Notice that
coos cos
Jo
1
- / E- cos((1) + vk)s)ds
1
e cos((u; vs)s)ds
2
J
2 2
1
000
1
L. 00
R (exp (-s + (u) (2)ss) ds
2 2
R Reexp (-spp t (e) + vx)s)) ds
2
0
1
2 R
1 i. (1/ - vx)) +
1 1 i. (ui + vb)
1
2
1 (v, )2)²
1 (1) + m2
Lemma C.6 (Hanson-War Wright inequality). Let (X,, ...., Xa) E R be a random vector with
independent components X; which satisfy E.X, 0 and ||x 11ol < K. Let A be an matrix. Then, for every 2 O,
12
K2|A )
P P[xT AX TXT T > 1) S 2exp
-emin
K+ P
17

---

Published as a conference paper at ICLR 2025
where is positive absolute constant and the subgaussian norm || 1l is defined
I| = supp
p21
In particular, if is a standard normal distribution, then ||Ell = v8/3.
Proof. We refer the proof to Rudelson & Vershynin (2013).
D PROOF OF THEOREM 3.
In this section, we prove the upper bound on the second moment of the model output value in Theorem 3.1.
Proof. First, may express the model output yz in matrix form. To do so, rewrite that contains the real and imaginary part of
2m x vector (f(c))..... R(cm). 3..)).... and let V be 2m x L Vandermonde-like matrix
R
Aw -1 Aw))
R
Aw -1 g
Awi H
- -1 e eamr
Awi
Aw1
3wm
11 c
Awnl
H
Awm
Amm! 1
R
Asm -1
Ammm A.-m(111))
Am
Am
Am
V
Awl auw1
Aw1
Awil
Aw1 9
Awl
Awi (L-1)
Awi
Awn
Awm Aww
Ammm -1 e Aawnn Awm
-3
Amm -1 e @dmm(b-1) Awm
Awm
Then yr. can be written in matrix form
c VJz,
where J E BLXL is row reversed identity matrix, i.c.
0
J
0/
Furthermore, we may connect V with standard Vandermonde matrix Vi,, by noticing that ΦV = DV,,
where Vy. a 2m x L complex Vandermonde matrix with 2m nodes eAw, eAu. Aum e Awm
Am
(t-11
A m
A@m(L (L-1)
V1
=
Aw
Aww (L-1)
E 2mx
Aum
@wwm(L- (L-
scaled unitary matrix
0
E 2m22m
0
0
1
18

---

Published as a conference paper at ICLR 2025
with $6 =b" 21gm, and D is diagonal matrix
-1
Aw1
30m -1
D
mm
C 22mx2m
Awr
Awm
Hence, we have V = 2 1 DVr. Notice that both R R(ury) and R R(Aū;) are non-positive, then
by Lemma C.I we have || 1. Now combining it with V 1 DV2 and the fact that the exchange matrix an orthogonal matrix, then when R(w,) I 0 for all j, we have
E...Lvf. = Ee,x (A C.c..JJ VJ+) I1
A [c VIE fr [JVJ
Be[llc|| /mll (E/ 1 VH)
A mm
Amax (E(xr'
m A A mAmax [E(ex
(AAR(1)) +... +( (cAR(1Z)) 1-1) )
f=1
= A 3 LAuax [E(xx) 1),
which finishes the proof.
It is also possible derive high-probability bound using Lemma C.6 (the Hanson- Wright in- equality). It's important to note that we do not make any assumptions about the input sequential data; instead, we only assume that the read-out vector e is i..d. Gaussian, as stated in Theorem 3.1. This allows us to apply the high-probability bound to the expression c' V'VEE [x "JJV" c, where V VEJE [xx "JJV
Hanson-Wright
deterministic matrix. By applying the Hanson- Wright inequality in Lemma C.6, we take a 0, and let V (xxx JJVT, Hanson-Wright 1ox(2/4) we have (2/(2/ ||||] Then for small enough 6, i.e., for large enough
v8/3, then by solving
KI
2
K=AAT Therefore, by Lemma C.6 we get with probability least 1- 6, we have K*DA
8A 2 log(2/0) ||||]
Ex iv2 < A 22m² 2 LA ( Fx
3c
4A 2 log (2/6)
< 2mm 2mmLLXu LA Fx 1)
Amaax
F Tr(V2
3c
log(1/6)
mL (1[[xT 1)
+
where X hides a positive absolute constant.
E PROOF OF PROPOSITION 3.3
In this section, we show the proof for Proposition 3.3.
Proof. Since Gj,A- f R(e (R("))))(("))> Ro(c")))dss then for any & E R. have
E T = / Σ R( ds 2 0.
19

---

Published as a conference paper at ICLR 2025
Hence, G' is positive semi- for both real-valued w and complex- valued w. Let TTGE O, then ンØ= ")))"">) = 0 for $ 2 O.
When a E IR"h, we take the discrete time points و = 0, 1,.... m to form m equations. Note that R(e")*) u If wj are distinct, then the Vandermode matrix given by W1,..... Wm is invertible, indicating that the only solution for Σド= Ej ()) = 0 is E = 0 for j I....,,.. m. Thus, G is positive definite in that case.
When E C with distinct imaginary parts, we can always find scaling factor > 0 such that e ,...,, e ,...., e are distinct, where w is the conjugate of w. Then by the argument of Vandermonde matrix, the only solution of the equation Σj= Sye" + = Eye for s 2 for j 1,..., m. Since 2*(e")**) e em /ten Σ ()
0 is that ε E
only has zero solution.
Combining these two cases we finish the proof.
F PROOF OF THEOREM 3.4
In this section, we prove Theorem 3.4 based on the Gershgorin circle theorem (Lemma C.3).
Proof. First, we need to bound both the diagonal entry and the off-diagonal sum. The diagonal entry Gja [1 + 14403) ) which can be bounded as
1 1+ 1+ 403
1
SGnSS j= I...., m.
2
For the off- dif--iaaonal sum, we have Vjj = 1....,,,
2Rg = 2 Σ GGja
1
1 (u, 0x)
+ (ey + 22)2
kfj
2
1
1
+
+
1+ 82 1
1+ (1) 0x)2
1 1 4u²
k=1
2
1
1+ 1
1 22 /
1/
1 4u
+
Σ
1+
1
+ 8212
1 40²
k=0
2
1
82
1/62
+
k2
8 k=
(1 02)82 +1 +
[1
1 40
k=1
where the first inequality is due to the fact that the minimal separation distance minjæa |03 Ux 2 5. and the last inequality is because by > 0 and reordering {2H Faz1 does not affect the result for
Ex= 14
Then by Lemma C.4, we have
2
122 + 1
Σ
1
3
Σ
1
82
1/82
82
01+ k=1 (1
22
82
1/82
k2
k=
k=
3
82
n
1 coth
(7)
82
2
3
3m
T coth
+
2
25
20

---

Published as a conference paper at ICLR 2025
Hence have,
1
3
3m
G
>
2
1
1
4v²
2
2
28
coth
5
=
2
1
+
1
+
40
1
2
3x
4
2
max
1
x2
1
4x
48
coth
3π
1.19
coth
46
Under the same argument, we get
3
3m
(0))
1
Gjs+ Rg Rg
< <1+
coth
+
max
28
2
1 [1 0z
+ 407
1
3m
1
1
4
+
46
coth
+
max
1+
uz ]
1 4u
3m
coth
12
48
Combining the two bounds and Lemma C.3, we finish the proof.
21


Reviews:
[{"review_id": "BBEAvg5Oxe", "reviewer": "Reviewer_YHUj", "summary": "The paper studies the initialization strategy used in SSM from various point of view, namely the time scale, the real and imaginary part of the state matrix. It demonstrates the dependency of the SSM timescale on sequence length based on sequence autocorrelation. Further it is shown that having a zero real part for the eigenvalues of the SSM state matrix mitigates the curse of memory. Finally, the paper demonstrates that the imaginary part of the eigenvalues of the SSM state matrix determine the conditioning of SSM optimization problems, and present a approximation-estimation tradeoff when training SSMs with a specific class of target functions.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper has presented some very novel an interesting insights into the initialization scheme for SSM.\n2. The analysis is well grounded and thorough.\n3. The experiments are well though out, extensive and results are convincing of the claims made in the paper.", "weaknesses": "While the experimental results are convincing on the datasets considered in the paper, how are the expected to hold in larger scale experiments? Can authors comment on that.\n\nWhat do the authors think of extending these insights or draw parallel to setting(s) where the sequence length can be varied?", "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": "SoduU0ibf7", "reviewer": "Reviewer_vKbW", "summary": "An analysis of the effects of some parameters of the so-called SSM layer is provided. Specifically, the authors analyze 1) the relation between the timescale $\\Delta$ and the sequence length $L$, 2) the effect of the (zero) real part of the time evolution matrix, and 3) the effect of the imaginary part of the time evolution matrix, both theoretically and empirically.", "questions": "No questions that may change my evaluation.", "limitations": "", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The claims are clear, and most of them are supported both theoretically and empirically.\n- The topic is indeed important as the SSM-based models are increasingly used in various problems.\n- Although the analyses may not be surprising per se in terms of dynamical systems theory, their implication especially in using SSM architecture sounds new and useful.", "weaknesses": "The empirical evaluation is limited to synthetic data or somewhat simple benchmark data. This fact does not much diminish the value of the paper, but reporting the applicability of the proposed theory to more real-world datasets would certainly be appreciated.\n\nA thing remained a bit unclear to me is that how the discussions are necessarily relevant only to the initialization. Doesn't it make sense to constrain the parameter values (e.g., $\\Re(W)$) as guided in the given theory, not only at the initialization but also during training iterations? As the paper's title emphasizes the role of *initialization*, a bit more discussion in this regard might be helpful.", "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": "gNIozP0jw5", "reviewer": "Reviewer_BxH3", "summary": "This paper provides a theoretical analysis of a state-space model from three perspectives: discretisation size ($\\Delta$), the real components of the diagonal weights, and their imaginary counterparts. It discusses the scaling laws of $\\Delta$, highlights how zero real parts contribute to long memory, and introduces an approximation-estimation tradeoff influenced by the imaginary components.", "questions": "**(I)** The analysis focuses on the single-input single-output case. How would these results extend to the multiple-input multiple-output (S5) case? High-level insights would suffice.\n\n**(II)** The paper assumes ZOH discretisation. How would the results differ with bilinear discretisation?\n\n**(III)** Several questions regarding Theorem 4.1:\n\n   (a) The upper bound applies to the expected output. Is there a corresponding probabilistic bound?\n\n   (b) Can the upper bound be shown to be tight?\n\n   (c) How does the real part of $w_j$ affect this bound? Specifically, how would increasing or decreasing it impact the upper bound?\n\n   (d) The analysis focuses on the final output. How would the results change if the sum of outputs $\\sum_{\\ell = 1}^L y_\\ell^2$  were considered instead? This pooling is adopted by many SSMs.\n\n**(IV)** The paper explores various settings for how $\\Delta$ should scale with $L$. Given a specific task, what method would you use to determine the appropriate $\\Delta$?\n\n**(V)** When the real part of $w$ is set to zero at initialization, is it still trained in practice? If so, do you enforce non-positivity? In this case, how do you handle reparameterization, especially since you can’t use the logarithm of it? Could you also show the behavior of the real part of $w$ after training? If it becomes negative, what is the benefit of zero initialization?\n\n**(VI)** On line 376, the assumption is made that $c$ is real-valued. However, many SSM implementations use a complex $c$. What necessitates this assumption?\n\n**(VII)** Given a specific task, how would you determine the appropriate balance between approximation and estimation tradeoff? Would this be driven by hyperparameter tuning or derived from theoretical insights?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**(I)** The analyses presented are rigorous, with clear explanations connecting the results to different initialization strategies, which is crucial for understanding SSMs.\n\n**(II)** The paper is well-structured, with illustrative examples that help to clarify the theoretical concepts.", "weaknesses": "**(I)** The paper primarily focuses on theoretical aspects, with limited empirical evaluation beyond basic synthetic examples. It would be beneficial to include practical applications that demonstrate the utility of the proposed methods.\n\n**(II)** The explanation of the three theoretical components sounds a bit separated. While they all fall under the umbrella of initializations, a stronger connection would give the paper a more cohesive narrative.\n\n**(III)** Although the theoretical results are clearly presented, they are dense. More emphasis on distilling key messages and offering practical guidelines would enhance the paper's accessibility.", "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 (Poster)