On Feature Learning in Structured State Space Models
Authors: Leena Chennuru Vankadara
Jin Xu
Moritz Haas
Volkan Cevher
Venue: NeurIPS 2024 poster (2024)
Keywords: State space models
feature learning
hyperparameter transfer
scaling theory
deep learning
Primary area: deep_learning_architectures

Abstract: This paper studies the scaling behavior of state-space models (SSMs) and their structured variants, such as Mamba, that have recently arisen in popularity as alternatives to transformer-based neural network architectures. Specifically, we focus on the capability of SSMs to learn features as their network width approaches infinity. Our findings reveal that established scaling rules, such as the Maximal Update Parameterization, fail to support feature learning as these models cannot be represented in the form of Tensor Programs. Additionally, we demonstrate that spectral scaling conditions, shown to be effective for feature learning in a host of other architectures, do not hold the same implications for SSMs. Through a detailed signal propagation analysis in SSMs, both forward and backward, we identify the appropriate scaling necessary for non-trivial feature evolution in the infinite-width limit. Our proposed scaling shows behavior akin to the Maximal Update Parameterization, such as improved stability, better generalization, and transferability of optimal hyper-parameters from small to large scale SSMs.

On Feature Learning in Structured State Space Models
Leena Chennuru Vankadara
Jin Xut2
Moritz Haas
Volkan Cevher
'GII Foundations, Amazon
2 University of Oxford" LIONS, EPFL
University of Tübingen, Tübingen Al Center?
Abstract
This paper studies the scaling behavior of state-space models (SSMs) and their structured variants, such as Mamba, that have recently arisen in popularity as alternatives to transformer-basee neural network architectures, Specifically, we focus on the capability of SSMs to learn features as their network width approaches infinity. Our findings reveal that established scaling rules, such the Maximal Update Parameterization, fail to support feature learning as these models cannot be represented in the form of Tensor Programs. Additionally, we demonstrate that spectral scaling conditions, shown to be effective for feature learning a host of other architectures, do not hold the same implications for SSMs. Through detailed signal propagation analysis in SSMs, both forward and backward, we identify the appropriate scaling necessary for non-trivial feature evolution in the infinite-width limit. Our proposed scaling shows behavior akin to the Maximal Update Parameterization, such as improved stability, better generalization, and transferability of optimal hyyer-parameten from small to large scale SSMs.
1 Introduction
State-space models (SSMs), such as Mamba (Gu and Dao, 2023), have become popular in deep learn- ing alternatives to transformers like GPT and BERT series (Radford al., 2019, Brown al., 2020, Achiam et al., 2023, Devlin et al., 2018, Touvron et al., 2023, Chowdhery et al., 2023, Gemini Team al., 2023). SSMs integrate elements from RNNs, CNNs, and control models, excelling inference and handling long contexts (Gu et al., 2021, Gupta et al., 2022, Gu et al., 2022, Smith et al., 2022).
The success of foundation models based on transformers and SSMs alike is largely attributed to their bothe-both in terms of data and model size. However, this increased scale often introduces challenges, such precision issues due instability or the vanishing/exploding gradient problems. Additionally, the sequential nature of state-space models (SSMs) makes them notoriously difficult to train. Therefore, developing rigorous understanding of how SSMs scale as their dimensions increase and identifying optimal scaling rules is crucial.
In this vein, infinite-width asymptotics, such Neural tangent kernel (NTK) analyses have been central tool providing key insights in DL theory. However, key limitation of the NTK analysis is the lack of feature learning in the infinite width limit (Yang and Littwin, 2023). Feature learning addresses the more realistic training setting where we have unrestricted movement of the neural network parameters (Yang and Littwin, 2023).
Intriguingly, recent work by Yang and Hu (2021) showed that under an expanded space of hyper- parameters, which includes layer-wise scaling of the learning rates, one can find a unique parame- terization called Maximal Update pPrameterization (μP) that induces non-trivial feature evolution
* denotes equal contribution. This work conducted during Jin's, Moritz' and Volkan's 's time Amazon. Correspondence saron.jin.xu gmmil.com
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Ingul [[0ae" -/26g
la
Ingul [iAba
Iiall v/N
Sogul |AAiil v/
v.Ngg
2.5%
v/.6g
[Heuntet))..
JPP-SS
v/N
//C
V.N
SSM Letent States, ZOH
55M Output States, 2011
5SM Latert States, Eutlr
SSH Output States, Buler
Figure 1: Under our derived scaling JP- Mamba achieves feature learning in all three SSM layers. In contrast, both Standard Parametrization (SP) and μP (heuristic) lead to instability or vanishing updates for either the latent states xi
(8) the output signal y: or both in each SSM layer.
The figures above illustrate the scalings when I = 1, but they exhibit the same trend across recurrence steps. We simultaneously scale up N₂ and Nu. We run each experiment 10 times, and the shaded areas indicate the standard deviation of these runs. Both Zero-Order Hold (ZOH) and Euler discretization of B; are studied and indicated in the subtitle.
even in the infinite width limit for host of key architectures. The results are obtained via the mathematical framework of Tensor Programs, which provides the foundation to study the effects of parameterizations on the learning dynamics in the limit. To this end, it is natural ask:
When do SSMs admit feature learning in the infinite-width limit?
Our research tackles this question by investigating the behavior of SSMs as network width approaches infinity. We examine how SSMs learn and evolve features in this context and assess the adequacy of recent scaling rules such as jP and spectral scaling conditions.
Our contributions:
We provide a detailed scaling analysis of forward and backward signal propagation in SSMs as the width approaches infinity, identifying that standard scalings lead to unbounded signals.
We demonstrate theoretically and empirically that popular scaling rules like μP do not yield correct scaling for SSMs due to-as we prove- their non-representability Tensor Programs (cf., Figure I).
We derive unique correction that ensures correctly balanced signals in both the forward and backward passes, stabilizing the training process and enhancing model performance.
Empirically, we validate that our proposed scaling facilitates hyper-parameter transfer from small-scale large-scale SSMs, similar the effects observed in MLPs and transformers.
2 On the lack of feature learning in SSMs at infinite width
2.1 Feature learning in sequence models
Sequence models are built by combining sequential layers that transform input sequences into output sequences. sequence layer can be represented as yI.z notation for sequence uI...., uz, with u RV,., y: E R.N, Note that fseq can be easily generalized to accept and output multiple sequences in residual connections and multi-branch architectures (see Appendix 1). Here Nr and N,, are treated as widths of the sequence model, which can be scaled up to create larger sequence models. To denote the backward pass of model concisely, let L be the overall training loss, and we write
(u1:)) where u1:z compact
8C Instead of studying the scaling of each element in a vector
de
E R-N., we study the norm ||u||2. When elements in are in @(1), we have ||u||2 E ((VN.). Throughout this work, we will use to represent the asymptotic order the limit (N, -> co in this case), but omit the precise notation for readability.

---

In layerwise sequence models, where the k-th layer is denoted as U1:n Sx [x(u)
(A)
(uf with u
RN,, we have the following definition for feature learning:
Definition 2.1 (Feature Learning in Layerwise Sequence Models), A layerwise sequence model is in the feature learning regime if for any k E [K]. the features uj./. and the updates of features Auj. (k) after one gradient update have the following scaling:
I.f.
B E [4] ||u (At 1|2 E @( (x)) (Stability at initialization)
(INIT)
E/E E [4], || 122 E @(v NN) (Non-trivial feature updates)
(A)
Note that the first condition is a stability condition that demands the activations to remain coordinate- wise (1) in the forward pass, whereas the second condition ensures nonvanishing and nonexploding activation updates. This definition can be easily generalized to general cases for arbitrary sequence models in Appendix B.
For particular sequence layer in the model, if we assume the inputs to the layer are asymptotically independent and identically distributed (i.i.d), correctly scaled, and the gradients backpropagated into the sequence layer have the correct scaling (as in Assumption 3.1), and this sequence layer satisfies the two conditions (INIT) and (A), we say this sequence layer admits feature learning. The corrected scaling pP-SSM we propose in Section satisfies an even stronger condition: The updates of all trainable weights should have nonvanishing and nonexploding effect on the activation and output function updates. We then say that the sequence model effectively feature learning. This is similar in spirit to the requirements for μP from Yang and Hu (2021) that ensure maximal stable updates of all trainable weights in standard architectures like ML.Ps, but requires different initialization and learning rate scaling rules for the trainable weights in SSMs, we show in Section 3.
2.2 Tensor Programs, spectral scaling, and the Maximal Update Parmeterization
Tensor Programs. The framework of Tensor Programs (TP) (Yang, 2019) was initially developed to understand the behavior of wide neural networks both at initialization and during training. While there are many versions of Tensor Programs, the most general version (that subsumes all the previous versions) is referred to as the NEgoR T program (Yang and Littwin, 2023). A NE&OR T program constitutes a sequence of vectors in IR and a sequence of scalars in R inductively generated from an initial set of random scalars, vectors, and matrices following specified set of instructions: matrix multiplications, non-linear outer products, and vector averages. Yang and Littwin (2023) show that for any architecture whose forward pass can be represented NEλORT T program (which includes many modern architectures used in practice including transformers and convolutional networks), there exists unique scaling rule called Maximal Update Parametrization (μP) under which features in every single layer evolve in a width-independee independent fashion with scale. Interestingly, it has been shown that unlike standard parameterizations, JPP allows transferability of optimal hyper-parameee from small to large scale models.
The key idea of the "Tensor Program Ansatz" (Yang, 2019) is that pre-activations and their gradients arising training most neural network architectures via standard update rules (such as SGD and ADAM) can be represented as vectors in a NESOR T program. The Ansatz suggests that vectors in TP become asymptotically independent as well as identically distributed both at initialization as well as during training. Leveraging this, the framework allows for mechanistic tracking of the behaviour of such vectors by assigning a random variable Z" to represent the (asymptotically identical) distribution of the coordinates of the vector This principle underlies the development of the key thcoretical result of the Tensor Program machinery - the Master Theorem - which can be viewed as the compositional, non-linear generalization of the law of large numbers. It allows the theoretical computation of infinite width limits of different quantities such as (pre-)activations or the output of neural network.
While the following proposition also holds for other structured SSMs such as MAMBA, for clarity, we first introduce the simpler S4 model.
The S4 recurrent layer. The S4 recurrent layer y fsu (u w (B, CJ) can be viewed as discretization of a continuous- time SSM given by
d x;
Ax, + Bu, and y: Re[Cx,]
(1)
dt
3

---

for 1 = ...... L with X = 0 and where Re(-) gives the real part of a complex vector, and where we let A (a) with a C.M,, B C.2, No, Ce E C xN= We have the discretized sequence to sequence mapping Ssa as follows:
X 1'x-1 + BBu and y1 = =e(C'xi
(2)
where we follow the Zero-Order-Hold (ZOH) discretization method,
A' exp(T
B' = (A' - I)A and C'=
In practice, simple Euler discretization can also be used for B, i.e., r To model long-range dependency, S4 advocates for HiPPO theory (Gu et al., 2020). There are many possible Hippo-based initializations possible for the structured transition matrices A (e.g., Hippo-Leg-S, S4D-Inv, and S4D-Real) and our result holds for all the parameterizations.
While the training of most common architectures including transformers or convolutional networks via update rules such as SGD or ADAM are representable as a NESOR T program, Proposition 2.2 shows that this does not hold in general for structured state space models.
Proposition 2.2 (Structured SSMs are not generally representable as NEgOR programs). The forward and backward signal propagation of structured SSMs including S4 (2) and (3) and MAMBA (Section 3. 1) trained via standard algorithms such as SGD or ADAM are not representable as NESORT T program. Indeed this holds for all existing Hinpo-parrmeterizations of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegS-D, S4D-Im; S4D-Lin, and S4D-Real.
Proof sketch. The formal proofs provided in Appendix C. The key architectural component that is not representable in a NESORT T program is the Hippo-based structured transition matrix A. Due to the rapid decay of the diagonal entries of A -1 the hidden states Xg are not even asymptotically identically distributed, already at initialization. Therefore coordinates of X cannot be generated by sequence of NEgORT T computations, neither initialization nor over the course of training. As we will discuss in Section 6, developing a TP-like framework to cover SSMs requires a substantial gen- eralization of the existing TP framework and is beyond the scope of the current paper. Nevertheless, in Section 6, we outline the key challenges and potential paths toward such generalization.
Spectral scaling condition. Yang et al. (2023a) showed that many practical architectures admit feature learning if certain spectral scaling conditions are satisfied. Let W, and AW, E N denote the Ith layer weight matrices and their gradient updates. The condition states that for feature learning to hold, the spectral norms of the matrices should satisfy:
(V
N₂ Nr-1
N₂ Ni-1
Well. E
and Awwi.. E (V
for all E [4].
Under jdP, architectures that are representable as a NEgORT program satisfy the spectral scaling condition (Yang et al., 2023a). Here, we show that for structured SSMs, architectures that satisfy spectral scaling conditions do not in general satisfy conditions for feature learning.
Proposition 2.3 (Spectral scaling does not generally imply feature learning in SSMs). Structured SSMs including S4 and Mamba trained via standard algorithms such as SGD or ADAM that satisfy spectral scaling conditions (*) do not satisfy conditions for feature learning given in condition (A). Indeed this holds for all well-know Hiipo-parameterizatte HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegS-D, S4D-Inv, S4D-Lin, and -eel.
of the structured matrix A including
Proof sketch. The formal proofs are provided in Appendix C. To gather some intuition for this result in simplified setting, consider the scale of the hidden states for token index O. which is already wrong at initialization. By definition, ||x1|| 112 [B'u, ll2 |ABu, where A = (A' I)A
Now, observe that ||Bu, is a sum of independent random variables that satisfy the Kolmogorov condition, so that the sum behaves according the strong law of large numbers. For spectral scaling conditions to yield the right scaling of the initialization variance, it is crucial that the following condition holds:
B'u 12 θ ([ ||u|ll))
(4)
B'ull, (B bba)
However, using standard tools from random matrix theory one can show that which clearly violates (4).
E @ (())))

---

70H Discretization
Fuler Discretization
SP
J.P (Heuristic)
μP-SSM (Ours)
RR
min
d0
VNu
VNu
min [1,
NN.N.
1/a
N
Nu
VN
1
Ne
N
VM.
Na
N₂ N
|x<
(1) ((vN.)
(min 1,
Nx
0 N₂ min (1,
@(v
ly
(v..)
((vN.N.)
+ (v min (1,
N2. min
(v
Ax
(v
N..
min
0
No ver min [1,
N.
((v
|Ayilla
@(N.)
N.,
Tr, min(1, N.
Nav VN,
min(1, x;
@(VN.)
Table Overview of the different parameterizations and their corresponding scaling for latent states, outputs and their updates. Results for ZOH and Euler discretization are separated by
Identifying the unique scaling for effective feature learning in SSMs
In this section, analyze the forward and backward signal propagation in structured SSMs. Due to the generality of the architecture, we consider MAMBA as the basis for the analysis. Similar arguments apply for other SSMs such S4, S5, H3, DSS. In the appendix, we also provide detailed analysis of signal propagation in S4 and identify the correct scaling conditions for maximal stable weight updates.
3.1 Selective State Space Models
Selective SSMs the Mamba architecture YI.z = J'aamaa (ui.). can be written as follows:
B = Liny, (ug; W B, bB) E R.V.
(5)
C LinN, ( We: bc) E RR==
(6)
7 Softplus(Ta + Broadcast,, (Lin W,, b,),, To. 77 E R
(7)
For i = I,.... Ny, we have N₂ ID SSMs:
A (6) Diag(- exp(a (1) )), alog E
(
(1) 1- A x + ui (6
(i)
B(() un E x (0) ER",
x
A (0) B (0) ZOH(T( (9) A Bt), (6) ER,
y ( Tx (0) vi E..
where X (i) 0, log To E RN. and all linear layer weights WB. bg, Wc, bc, W,, b, trainable parameters. The Zero-Order-Hold (ZOH) discretization procedure A B(()) ZOH(T,A (()) can be written as:
A = exp(T(") A (Q))
= - 1)A B
(8)
In nutshell, Mamba can be seen as N. SSMs, one for each input channel. Weights for these SSMs shared and depend on the input that recurrent step. Because of the dependency inputs, Mamba can model non-stationary sequences. Weight matrices are initialized as follows:
WB ~ Ni 1) E N.N.N. We ~ ~N(,00811) E RNaXNo W, ~ oi E RIN..
and zero initialization for all biases. For each SSM, A is still initialized according to the HiPPO theory. When (i) is real-valued, we let [ [J] log(j + 1) and To can be seen as bias term initialized to Softplus (2ttpplus-1(((00 (27 0.1)). We will assume that is not trained, as this does not have any effect on the scale of the different quantities under consideration. This is a minor technical assumption and our results would also hold if we considered 7 as trainable parameter.
5

---

Mr R
Discretization
ERR
Per Channel Linear Recurrance A 1
exp(y 40)
(A
1 A.B. or
gao B
diag( 61
al R
B
Wi e Re by R We R.N-NN R.M. RRIK To
B: R C) 8
Selaction
B: Wan + by
R
Wert bo
Broailentt (Wau, Rt) Soliplaa
Figure 2: Illustration of the Mamba S6 Layer. The computation is modularized into three components: selection, discretization, and per-channel linear recurrence. Mamba introduces selection mechanism where weight matrices B, C₁ depend on the inputs u. These weight matrices are then separated into per-channel parameters, and discretized using either the ZOH or Euler methods. The discretized, per-channel weights are then applied in linear recurrence, allowing each channel to perform computations in parallel. Trainable parameters are shown in blue.
3.2 Forward signal propagation through S6 Mamba layer
To derive the correct choice of initialization scalings о and oe, we begin by analyzing the scale of activations (i.e., hidden states and outputs) in a S6 Mamba layer in the first forward pass as N,, + oo then N₂ + oo. We believe that our results would also hold in the proportional limit where N, oo and N₂
with N- No O11). However deriving the results in this setting would incur significant technical overhead and defer this analysis future work. We briefly discuss this in Section As is common in practice, we assume that the SSM layer is embedded into a neural network architecture containing standard architectural blocks such as MLPs, normalization layers, residual connections. All proofs are provided in Appendix C.2.
Assumption 3.1. Assume that the forward pass of all the components of the network except the SSM layer are expressible as NEΦOR programs and are parameterized according to pP.
This assumption ensures that the inputs to the SSM layer are asymptotically i.i.d and correctly scaled. It also ensures that gradients into the SSM layer have the correct scaling. All results in this section are stated under Assumption 3.1. Through the forward signal propagation analysis, we identify the correct scale of initialization for weight matrices W n and W c. We show that both standard parameterization as well as spectral scaling conditions do not yield the correct scale of initialization for the weight matrices. The key results summarized in Table 1.
For simplicity of exposition, first consider the scale of x = (A) - I)A (2) B before generalizing to arbitrary ( [L].
Proposition 3.2 (Scale of hidden states x in Mamba at initialization). Under the 70H dis- cretization procedure, as N₂ then Nº approach infinity. for any E [Nu]. the squared 12-norm of the hidden states ||x{ (2)11 [|3 is a.s. scaled as |xx" 11 e (c(2) o | (u (u()))). where ((2) denotes the
Riemann zeta function at 2.
Following condition (INIT), the ID SSM admits stability initialization under the following conditions:
If u 1 e @(1) then ||x] (0 || No) and 00 O(1).

---

Therefore, for stability at initialization, the initialization should scale as oB E @(v N.. Note that, under both standard parameterization (e.g., Kaiming LeCun initialization) and spectral scaling conditions, dn is initialized as Ø( N.) which leads to vanishing hidden states according to Proposition 3.2. We empirically verify this fact in Figure I.
Proposition 3.3 provides the scale of the output of a Mamba layer.
Proposition 3.3 (Scale of outputs yi" of Mamba layer at initialization). Under the ZOH discretization procedure as Nu then N₂ approach infinity, for any E [Nu). the output " converges in distribution to Gaussian with mean 0 and standard deviation Ca || 21 for some width- independent constant C
Accordingly, imposing the conditions for stability of o according to condition (INIT) implies the initialization scaling condition oc E @(v N.N.). Note again that standard parameterization
suggests initializing oc E @(v No under which the outputs would diverge with scale. Under spectral scaling, oc is initialized much larger. However, since the hidden states vanish with width, the outputs of the SSM admit the correct scaling here as demonstrated in Figure 1.
Generalizing arbitrary 1 E [L]. The corrected scalings of n and oe derived above generalize to the entire sequence, as sum over the sequence usually does not cancel out the scaling. More formally, for all / [L.], we have
L-1 [AA
"B
69) i-m "im
m=0
First, observe that the operator (A, (A)))) A (6) Dia ,'v_) with d', e br(" (cos(T)
(1)) TM does not change the width-scaling, To see this note that since Tn) + (isin(((") xn)), we have that, for all
complex vectors v E C. it holds that [I(A, ((0)
vvll emo? ///2/111 | for any m E [Z.. Now since, E @(VN2) each term the summation is u perfectly cancels out with the terms before E @(vN.). The same argument can used to
N N for any u, setting aB E @
yields
(1)
m
of order (v)N.). Unless, for every I, the term B
to affect the width scaling, we have that x (
show the stability of yi (0) Concluding this argument, we have derived the correct scaling of the initialization variances dB and oc for feature stability in a Mamba layer summarized below.
Conditions for stability of a S6 Mamba layer at initialization. The features of a S6 Mamba recurrent layer Vi.t- f mamba (uizzi w) stable initialization in the infinite-width limit under the following scaling conditions:
N N,
and ac E (V
N₂
3.3 Backward signal propagation in S6 Mamba Layer
In this section, we provide the correct scaling of the learning rates 7a, 7/. and VJc by a detailed analysis of the backward signal propagation in the limit of N, analyze the scale of the activation updates for both hidden states and outputs in the first backward pass through Mamba layer. We also show that both standard parameterization as well as spectral scaling conditions do not yield the correct scale of learning rates for the weight matrices. The key results are summarized in Table
then N2 + cc. Specifically, we
Proposition 3,4 (Scale of the updates of hidden states Ax after 1 step of SGD). Under the ZOH discretization procedure, as N₂ then N., approach infinity, for every ID SSM, the squared (4) l| of the hidden states after one step of SGD is a.s. scaled as |x ( l|2 E (1I VN. oc |n 1 <(4) where ((4) denotes the Riemann zeta function at 4.
bz-oorm of the updates
||Ax
7

---

The Riemann zeta function at 4 evaluates to a width-inddepedent constant and oe E @ N.N. due to Proposition 3.3 for stability of outputs. Therefore, for the scale of the hidden state updates to be D(yNa), the correct scaling of the learning rate 1/8 is given by θ (y-)). On the other hand, spectral scaling suggests that 7/8 must scale as O N₂ ). Under this scaling, however, updates of the N.
hidden states would vanish with width and therefore the first block of the SSM is in the lazy regime. This is corroborated by our experiments in Figure
Next, to derive the correct scaling of the learning rate nc, we consider the scale of output updates.
Proposition 3.5 (Scale of the updates of outputs Ay after 1 step of SGD). Under the ZOH discretization procedure, as Na then N₂ approach infinity, for every 1D SSM, the squared ly-norm of the updates of the hidden states after one step of SGD scales
Ay 1 E θ (noo BV N. [
The result suggests that for the correct scaling of the updates, the learning rate `C must scale as 1 @(,, N₂ V.,. ). Under spectral scaling, Ωc scales much larger as O ) However, since the updates of the hidden states vanish under spectral scaling, this larger incorrect scaling of 1Jc: downward corrects the scale of the updates in the outputs as shown in Table and empirically verified in Figure 1.
On the correct scaling of na- It turns out that the scaling of the learning rate ?h does not play a role in either stability at initialization or for non-trivial updates with scale. However, if Ta is not scaled correctly, then the transition matrix A is not updated. In particular, as shown in Appendix C.2, fe needs to scale as (N.). Below we summarize the correct scaling conditions to achieve non-trivial feature updates in the infinite-widd width limit of Mamba layer.
In the same vein as the discussion in Section 3.2, it is straightforward to verify that our results hold for arbitrary 1 [7.) and more gradient steps as soon as we assume that the updates of the weight matrices and activations do not perfectly cancel out the corresponding initial quantities.
Conditions for non-trivial feature updates in a S6 Mamba Layer. The updates S6 Mamba recurrent layer y1.z = fmamba [U1.:; w) evolve non-trivially in the infinite-width limit under the following conditions:
oc E E 1 ) na E ONN.,, TB E O and nJe E O 1 N.
оB E Φ(1
N..N
NaVN..
4 jPPSSM implies stability and feature learning in Mamba
Empirical verification of different scalings in Table 1. First, we verify that our derived JPPSSM scaling (see Table 1) for Mamba indeed leads to feature learning, i.c., ensures stability at initialization as defined in condition (INIT) and non-trivial feature updates during training as defined in condition (A). We scale up the SSM latent state size N₂ and the SSM output dimension No simultancously and track the scaling of both features and feature updates. Due to the linear decay in the eigenvalues of the transition matrix A we typically observe a strong finite sample effect at small N₂. Constrained by computational resources, we opt for much smaller N., (N, = N2 than is usually employed in practice. This adjustment enables us scale up N₂ effectively, thus mitigating the finite-sample effect and to more clearly demonstrate the scaling behavior in the asymptotic limit in Figure 1. For all experiments this section, we train Mamba with SSM blocks for language modelling on the wikitext dataset (Merity et al., 2016) and use plain Stochastic Gradient Descent (SGD) to perform gradient updates. We use the huggingface (Wolf et al., 2019) Mamba implementation and the juP package (Yang et al., 2022) for scaling in our experiments.
As shown in Figure 1. under Standard Parametrization (SP), both the SSM latent states and the outputs explode at initialization, and their updates also explode, leading to instability both at initialization and during training. Under the spectral scaling parameterization prescribed in Yang et al. (2023a), other layers except for the SSM layers have the correct scaling by design. However, in the SSM layer, when using Zero-Order Hold (ZOH) discretization for B,, the latent states at initialization and their updates vanish when scaling up the width, while the output signals still have the right scaling as predicted by theory. On the other hand, when Fuler discretization used for B,, the latent states will have the correct scaling, but the output signals will explode. The results clearly highlight the importance of
8

---

correcting the scaling of the pP parameterization. When using the corrected scaling which we call uP-SSM, both the latent states and the output signals have the right scaling at initialization, and their updates have the same correct scaling. This holds true for both discretization schemes. Note that the slight shift scaling when the width is small is due to finite sample effects. stabilizes once the width is sufficiently large.
Stability, generalization, and hyper-parameet transfer. In Figure we employ Mamba generative model on the wikitext- 103 dataset and conduct single-epoch training for 20 iterations. We plot the test loss against the learning rate on a logarithmic scale and compare the results across different model widths (both N2 and N₂). In this experiment, we use the standard setting where Nu > N₂ (N, 1 in this case). Using μP-SSM scaling or jPP (heuristic) significantly improves test performance compared to standard parameterization of Mamba for this task. For larger learning rate, pP-SSM shows better stability compared to pP (heuristic), highlighting the importance of deriving the correct scaling for SSMs rather than heuristically adopting jP (heuristic) without investigation. Furthermore, previously unreported, we observe stable HP transfer from small to large widths and monotonically improving performance with increasing model widths in structured SSMs. In contrast, we observe completely non-monotonic behavior under standard scaling of SSMs.
Note that optimal learning rate also appears to transfer under spectral scaling. The reasoning behind why optimal parameparseserss transfer across scales is not completely understood. For instance, the optimal learning rate has been empirically shown to transfer across depth in transformers under appropriate
scaling of the residual branches (Bordelon al., 2023). However, from a theoretical standpoint, layers within each residual block of a transformer are in the lazy regime in the limit (Yang et al., 2023b). This is very similar to Mamba under spectral scaling. Under the ZOH discretization, the first block of the SSM is in the lazy regime but the outputs themselves are updated non-trivially. This suggests that more thorough understanding of both necessary and sufficient conditions of the transferrability of hyper-parameters is warranted.
SP
JPP (Heuristic)
PPP-SMM (Ours)
4.5
512
4.0
1924
Test loss
3.5
3.0
3.5
loga (Ir)
loga (Ir)
loga (Ir)
Figure 3: Test loss against learning rate on Mamba with varying widths (N, and N,). Using μP-SSM scaling leads to substantially improved test performance compared to the SP scaling. Compared to IPP (heuristic), JP-SSM scaling provides greater stability when utilizing large learning rates. Notably, we observe stable learning rate transfer from small to large model widths. Performance improves monotonically across widths in structured SSMs under JPPSSM scaling, as opposed to standard scaling where performance actually drops with scale after certain width.
5 Related work
Signal propagation. Our work can be seen as scaling theory or signal propagation theory with the goal of preventing both vanishing and exploding signals in forward and backward passes. In this sense, we build on a rich literature,, often restricted to an analysis at or close to initialization (Schoenholz et al., 2016, Poole et al., 2016, Hanin and Rolnick, 2018, Xiao et al., 2020). Towards understanding infinite width limits of neural networks, kernel-based approaches (Neal, 1996, Jacot et al., 2018) and applications of feeld-field theory (Mei et al., 2018) have yielded valuable insights.
Tensor Programs. Most promisingly, the Tensor Programs framework (Yang, 2019, Yang and Hu, 2021, Yang and Littwin, 2023, Yang et al., 2022, 2023b) covers many modern deep learning archi- tectures, optimization algorithms and arbitrary ale-parameterizations. Each praaeteetttiiitioonn essentially defined by layerwise scaling of initialization variance and learning rate as function of

---

network width. Seminal work by Yang and Hu (2021) shows that there exists unique maximal update parameterization (μP) that attains stable feature learning fhiiite-widtth limit. This parameterization has since been shown to be a good model for understanding the properties of large models (Vyas et al., 2024), and has been extended to infinite width and depth limits of ResNets (Hayou et al., 2021, Li et al., 2021, Bordelon et al., 2023, Yang et al., 2023b) and Transformers (Noci et al., 2022, 2024).
Structured SSMs. Our analysis focuses on structured state space models (SSMs). The S4 model (Gu et al., 2021) is inspired by continuous-time linear SSMs, which are well-studied in control systems, and its specific initialization is motivated by the HiPPO theory (Gu et al., 2020). S4 and its variants e.g. DSS(Gu et al., 2022). S4D(Gupta et al., 2022), S5(Smith et al., 2022), etc., demonstrate impressive long-range dependency and overcome the quadratic computational cost of transformer models (Vaswani et al., 2017) w.r.t. sequence length. However, these models are less effective at modeling text, or even perform simple tasks such as selective copying (Gu and Dao, 2023). Mamba (Gu and Dao, 2023) is proposed to address such issues with selection mechanism. This line of work has also inspired revisiting Recurrent Neural Networks (RNNs) (Orvieto et al., 2023, De et al., 2024, Beck et al., 2024), leading to the growing interest in RNN-based sequence models.
6 Discussion
In this work, we study the scaling behavior of forward and backward signal propagation in structured state space models a promising class of recent architectures. We show that existing scaling rules such as standard parameterization, μPP or spectral scaling conditions do not yield desirable properties such as feature learning in SSMs at scale. Through our analysis, we propose the correct scaling of state space models under which we empirically observe feature learning and transferability of preerpprramters from small to large scale models.
On Generalizing Tensor Programs. While our proposed scaling has been derived by a thorough analysis of signal propagation in SSM layers, our results are still limited the N2 then Nr tend to infinity setting. A completely rigorous analysis of SSMs in the proportional limit where N₂ and N₂ approach infinity with Ns held roughly constant requires us to carefully track how the different N. activations and the updates are correlated with each other. For most standard architectures, the Tensor Program (TP) machinery provides the appropriate tools to do precisely this. Therefore, it is of considerable interest to generalize the TP framework. The key assumption that requires relaxation that the different vectors in a TP (such activations or updates) are asymptotically identically distributed. As discussed earlier, this assumption crucially underlies the key theoretical results of TP called Master Theorems. However, potential path toward generalizing TP may be found by noting that the Master Theorems can be viewed as a non-linear compositional form of the law of large numbers or central limit theorem. Accordingly, may be possible to relax the assumption of being identically distributed in the limit and instead ask that the entries of the vectors in a TP asymptotically satisfy weaker conditions such as Lindenberg or Kolmogorov conditions. Note however, that any such generalization is highly technical and is beyond the scope of the current work.
10

---

References
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877/ 2023.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp. Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory. arXiv preprint Stt.::24000045117., 2024.
Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer residual networks: Dynamics and scaling limit. arXiv:2399.16620, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in neural information processing systems, 33:1877-1901, 2020.
Aakanksha Chowdhery, Sharan Narang. Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240):1-113, 2023.
Soham De, Samuel I. Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Alhert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint vrXiv:2402.19427 2024.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv:1810.04805. 2018.
Gemini Team, Rohan Anil, Sebastian Borgeaud, Yonghui Wu, Jean-Baptiste Alayrac, Jiahui Yu, Radu Soricut, Johan Schalkwyk, Andrew M Dai, Anja Hauth, et al. Gemini: family of highly capable multimodal models. arXiv preprint arXiv:2312.1 1805, 2023.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, arXiv preprint ar xrii::2312.00752, 2023.
Albert Gu, Tri Dao. Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. In Advances in Neural Information Processing Systems, volume 33, pages 1474-1487, 2020.
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2021.
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, 2022.
Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are effective as structured state spaces. Advances in Neural Information Processing Systems, 35:22982-22994, 2022.
Boris Hanin and David Rolnick. How to start training: The effect of initialization and architecture. Advances in neural information processing systems, 31, 2018
Soufiane Hayou, Eugenio Clerico, Bobby He, George Deligiannidis, Arnaud Doucet, and Judith Rousseau. Stable resnet. In International Conference on Artificial Intelligence and Statistics, pages 1324-1332. PMIR, 2021.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and gener- alization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS), pages 8571-8580, 2018
11

---

Mufan Li, Mihai Nica, and Dan Roy. The future is gaggssaann Resnets and their infinite-depth-are width limit at initialization. In Advances in Neural Information Processing Systems (NeurIPS). volume 34, pages 7852-7864, 2021.
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, (1133):E6655-77777 2018.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture models. In International Conference on Learning Representations, 2016.
Radford M. Neal. Priors for Infinite Networks, pages 29-53, Springer New York, 1996.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 35:27 199-2119--721,, 2022.
Lorenzo Noci, Chuning Li, Mufan Li, Bobby He, Thomas Hofmann, Chris J Maddison, and Dan Roy. The shaped transformer: Attention models in the infinite depth-and -idth limit. Advances in Neural Information Processing Systems, 36, 2024.
Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, and Soham De. Resurrecting recurrent neural networks for long sequences. In International Conference on Machine Learning, pages 26670-26698. PMLR, 2023.
Ben Poole, Subhaneil Lahiri, Maithra Raghu, Jascha Sohk-Dickstein, and Surya Ganguli. Exponential expressivity in deep neural networks through transient chaos. Advances in neural information processing systems, 29, 2016.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models unsupervised multitask learners. 2019. h org/ CorpusID: 1600225533
Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. 1.10v2261..01232, 2016.
Pranab Sen and Julio M Singer. Large sample methods in statistics: an introduction with applica- fions, volume 25. CRC press, 1994.
Jimmy TH Smith, Andrew Warrington, and Scott Linderman. Simplified state space layers for sequence modeling. In The Eleventh International Conference on Learning Representations, 2022.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXiv:2307.09288 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Roman Vershynin. Spectral norm of products of random and deterministic matrices. Probability theory and related fields, 150(3):471-509, 2011.
Nikhil Vyas, Alexander Atanasov, Blake Bordelon, Depen Morwani, Sabarish Sainathan, and Cengiz Pehlevan. Feature-learning networks are consistent across widths at realistic scales. Advances in Neural Information Processing Systems, 36, 2024.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, et al. Huggingface's transformers: State-of-the-ar natural language processing. arXiv preprint arXiv:1910.03771, 2019.
Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and gener- alization deep neural networks. In International Conference on Machine Learning, pages 10462-10472. PMLR, 2020.
12

---

Greg Yang. Wide feedforward or recurrent neural networks of any architecture are gaussian processes. Advances in Neural Information Processing Systems. 32, 2019.
Greg Yang and Edward Hu. Tensor programs iv: Feature learning in anfinite-widthh neural networks. In International Conference on Machine Learning (ICML), 2021.
Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the fffnite-width limit. arXiv:2308.018144 0114,, 2023.
Greg Yang. Edward J Hu, Igor Babuschkin, Szymon Sidor, Xiaodong Liu, David Farhi, Nick Ryder, Jakub Pachocki, Weizhu Chen, and Jianfeng Gao. Tensor programs Tuning large neural networks via zero-shot hyperparameter transfer. arXiv:2203.03466. 2022.
Greg Yang, James B Simon, and Jeremy Bernstein. A spectral condition for feature learning. arXiv preprint arXiv:2310. 17813, 2023a.
Greg Yang, Dingli Yu, Chen Zhu, and Soufiane Hayou. Tensor programs vi: Feature learning in infinite-depth neural networks. arXiv preprint 230002231..02244, 2023b.
13

---

A Background
A.1 NE@ORT, abc-paraameterizatios and yPP
NE@ORT programs. For the precise definitions and theorems of the NE@OR T framework we refer to Yang and Littwin (2023). Here try provide an intuitive introduction.
NE@OR T program consists of a set of vectors inductively generated from set of initial ma- trices, vectors, and scalars following set of allowed instructions: vector averages (Avg), matrix multiplications (MatMul), and non-linear outer products (OuterNonlin).
All entries of an initial vector are sampled iid from N(0, 1). Every entry of an initial matrix A is sampled independently from distributions that all have mean variance n -1 and that satisfy techni- cal boundedness ussumption on all higher-orde moments. Nonlinearities vó used for OuterNonlin operations are either aseudo-Lipschin or polynomially smooth. Initial scalars c converge to 0 a.s.
New scalars c can be introduced into the program by the Avg operation over a vector v, defined as c= Ea Va- Now the NE@ORT Master Theorem implies é EZ" almost surely. Matrix multiplications Au defined usual. In the limit, non-trivial interaction term arises if A had appeared in a NESORT operation that has generated the vector v. To define the OuterNonlin operation, aggregate fixed amount |x| of NE@OR T vectors in and fix r E N. Denote the collection of all defined NE@ORT T scalars by c. Then a OuterNonlin operation with nonlinearity ah: : [ + R, is given by
1
Ya
(xxa. ..,.... X3,i)), c), $1....,,
Σ
and, due the NE@ORT T Master Theorem, behaves in the limit as
2 = f(Z*), where f R|x + R, f(Z²) = LL (,,,0)]
where 2 is an r-dimensional random vector with the limiting distribution of x, and Z7.... Z,5 are id copies of z*
The Maximal Update Parameterization (μPP). For any architecture/compnnent whose forward pass is representable as a NEROR T program, j&P can be derived following the instructions below, Let each parameter tensor W be parameterized
W n--ne n-ow where is trainable parameter with initialization Wo ~ N NV00, n (2-2bw). ). Let the learning rate be parameterized as ryn for some width-independent 1j 0. Then μP is prescribed follows: Set 0 and bw 1/2 for all parameter tensors W.
aw 0 if the input and output to W are width-independerr (e.g., scalars),
aw = 2 both input and output to W scale with width (e.g., attention matrices).
aw if input is width-indepent
and output scales with width (e.g., input embeddings). aw = 1/2 input scales with width and output is width-iindependet (e.g., output matrices).
A.2 Classical Limit Theorems
Here we recapitulate classical results about large sums of random variables that we use in our proofs. We present the versions of the Kolmogorov strong law of large numbers and the Lindenberg-Feie Feller Central Limit Theorem provided in Sen and Singer (1994, Theorems 3. and 3.3.3).
Theorem A.1 (Kolmogorov Strong Law of Large Numbers). Let X', > 1, be independent random variables such that EX;; μ4 and V aar (Xx,) = 03 exist for every 2 1. Then
Σ k 2f < 00
Σ p4 as O.
Theorem A.2 (Cindenberg-Fellle-cceoe Central Limit Theorem). Let X,, i2 be independent random variables such that EX, Mi and Var (X.) of exist for every i 2 1. Also lei oy - of and
Z, s, [2 == X; Σm=
= pi). Then the Lindenberg-Fede Feller condition
n
Ve > 0,
E [(x: - -10)=000xxxxx222222 + 0, as n + 00,
14

---

holds if and only if both
(i) maxisicn a as + and
(ii) Zn converges in distribution Gaussian with mean 0 and standard deviation
B Definitions
Definition B.1 (Softplus). The softplus fimction is defined as Softplus(z) = log(1 + exp(3x))
with default smoothing value 3 = 1.
B.1 Feature Learning in General Sequence Models
Let 9 be directed acyclic graph (DAG) with vertices V(G) and edges E(g). For each node V()). its parent nodes are denoted by P.((). In a sequence model, each node v corresponds to sequence VI-L with length L, computed from its parent node sequences with operation VI.L fo( (u1.z ju PA(())). These sequence layers can be shared instance-wise multi-layer perceptrons (MLPs), normalizations, residual summation, or transformerssrecurrent layers, etc.
Definition B.2 (Feature Learning in Sequence Models). A sequence model is in the feature learning regime for any E VG)), the features Vi.t. where v E RR and the updates of features Av after gradient update have the following scaling:
3 E [4], livil2 = ((VNNN)))) N.) (Stability at initialization)
(9)
3 I [7], ||Av/|22 O No) (Non-trivial feature updates)
(10)
C Proofs
C.1 Structured SSMs are not covered by previous approaches
Proposition C.1 (Structured SSMs are not generally representable as NESORT programs). The forward and backward signal propagation of structured SSMs including S4 and MAMBA trained via standard algorithms such as SGD ADAM are not representable as NEgoRT program. Indeed this holds for all existing hiipo-parameterization of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPÖ-LegS-D, S4D-Iny S4D-Lin, and S4D-Real.
Proof. Recall the definition of NESORT T program (Yang and Littwin, 2023). Let A be any of the HiPPO or S4D structured matrices. As the entries of A initialization are deterministic and differ, A can neither be initial NEØOR T matrix, nor be generated by an OuterNonlin of random NERORT vectors (acting on all coordinate dimensions in the same way). This is because no allowed NETORT operation
generate NESORT T vector with differing variance in its coordinates, whereas the multiplication with A clearly produces differing variances each entry. All that is left to do is to show this claim via induction.
Claim: All NE@OR vectors have the same variance in each coordinate.
To start the induction, all entries of an initial vector v are sampled iid from N(0, 1).
Now assume all vectors v currently defined in the NEgoR program have the same variance. The allowed operations to generate NEgOR T vector are matrix multiplication with an initial matrix and OuterNonlin. As initial matrices Ao have independent entries with mean 0 and variance -1, the multiplication Agu will again generate a vector with the same variance in each coordinate.
An OuterNonlin operation takes in previously defined NE@OR T vectors and treats all coordinates in the same way. Consequently will again generate a NE@ORT T vector with the same variance in each coordinate.
There is other way generate new NERORT T vector, which concludes the induction.
To see that the structured matrix A produces nonisotropic coordinate distributions, we detail the argument for S4 recurrent layer. The other choices of A follow via analogous arguments. Recall
15

---

that the S4 recurrent layer y1-z = fsi(u1::: w (B, can be viewed as a discretization of continuous-timn SSM given by
x = Ax, + Bu, and
Y: = Re[Cx,]
(11)
dt
for 1, .... L with Xg = 0 and where Re(.) gives the real part of a complex vector, and where we let A (a) with a E CN=, B E CN..Ma, C E CN, We have the discretized sequence to sequence mapping fsa as follows:
X A'x + B'u and y1 = Re(C'xi]
(12)
where we follow the Zero-Order-Har Hold (ZOH) discretization method, B' = (A' - I) -B, and C²
A' = exp(r A);
(13)
To model long-range dependency, S4 advocates for HiPPO theory (Gu et al., 2020). While there are many possible Hippo-based initializations possible for the structured transition matrices (e.g.. Hippo-Leg-S. S4D-Inv, and S4D-Real) and our result holds for all the parameterizations, let us consider a simplified initialization for clarity. Following S4D-Lin, we can set a, - + imn. B and C are initialized such that the entries are i.i.d and follow a Gaussian distribution with mean 0 and variance B and of respectively (for both real and imaginary parts).
Assume for now that the inputs to the SSM layer are i.i.d which holds asymptotically in general say if they are outputs of a previous layer such as an MLP. Let consider the distribution of the hidden states for token index O at initialization: Xg = B ugg where B' (A' I) BB, and + i sssn( (Tn)), then the eigenvalues of A' are clearly ((). To further simplify, if we set 7 = 2, then ||| e- and for all complex vectors an i.i.d Gaussian matrix, for sufficiently large width, Bug is distributed i.dd by simple Central Limit Theorem argument. However, B' also depends on 1+2mné 1/2+2=21,2- Due to the linear decay the diagonal entries of A -!.
A' = diag(a),.... d'y,) with d', = e-' 1 (cos
v C.N it holds that A'v -e 'v. Since
A A-1 diag(o, (a, with a,
the hidden states Zg already at initialization are not (even asymptotically) identically distributed. Therefore, there cannot exist coordinate distribution that represents the entries of vectors arising in modern SSMs such as MAMBA neither at initialization nor over the course of training.
Proposition C.2 (Spectral scaling does not generally imply feature learning in SSMs). Structured SSMs including S4 and MAMBA trained via standard algorithums such as SGD or ADAM that satisfy spectral scaling conditions (+) do not satisfy conditions for feature learning given in Equation (D). Indeed this holds for all well-known tiopo-parameterization of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegSS S4D-Inv, S4D-Lin, and S4D-Real.
Proof. Let A be any of the HiPPO S4D structured matrices.
Consider the scale of the hidden states at initialization for token index O. By definition, |x1 llz B'uillz = |ABuill where A = (A' I)A Then, observe that [ 12 =
N
A Bi. is sum of N₂ independent (but not identically distributed) random variables.
However, it is easy to verify that they satisfy the Kolmogorov condition and therefore the sum behaves according to the strong law of large numbers. Intuitively, this is allowed by the polynomial decay of eigenvalues of A and the correct scaling of the inputs to the S4 layer. Applying the Kolmogorov SLLN (Thcorem A. 1), we obtain
Σ
A2
B Bijuns) a., oB 10 | |1111(((22)
where ((2) denotes the Riemann zeta function at 2 and equates to 7² 72//6, and e >> 0 some width- independent constant. The zeta function appears due to cc(2) < Σnn A2 < Ein Tet - (()) for N₂ large enough (see (37) for more details). Here and throughout the paper, for the formal application of the L.IN or the CLT, we first adequately normalize all quantities to arrive at well- defined indhhinndependent limit statements, but do not write out such technicalities for conciseness.
16

---

Observe that, for spectral scaling conditions to yield the right scaling of the initialization variance, it crucial that the following condition holds:
| ll2 (II'II (|ulll))
(14)
Since the spectrum of (A' I)A- is less than 1, upper bound the spectral norm of B' (A' I) -B 'B can be found in Vershynin (2011) and is given by c(VN, + vN.) for some indeedinddeendett constant c. A matching lower bound can be easily found by not- ing that the spectral norm is lower bounded by the maximal row and column norm: |[B']] Z
|B'u|l2 ([ E O (TT) which clearly
max mmxx [Bull2 max
|By (llla
2 VNudB. Therefore
violates (14). This spectral decay is induced by of the considered choices of
C.2 S6 Mamba recurrent layer
Recall that S6 Mamba layer YI:L fMamba (Mamba(u1..)) can be written as follows:
B = Liny, iu; WB. WB.b3) E R.V,
(15)
C = Linx, ( W bc) E RN=
(16)
7 Softplus(To + Broadcasty, (Lin (u; b,),, To. RNv
(17)
For i = 1, 1,..., Nu. we have N₂. ID SSMs:
A (0) Diag( )xp(a(o))
(1) )) algg E ERR;
(+)
x ( = A
(00 (1)
x u (
B(O un E x ER",
A (0) B (i) 7OH(1
(1) A (0) Bt), 7 ER
(
(1) T yi E.R
y/
where X
(8) 0, log" TO E RRY, and all linear layer weights WB, b Wc, bc, W br are trainable parameters. The Gero-Order-Holl (ZOH) discretization procedure A B(())
ZOH(T, (0) B((i) can be written as:
A (() = exp(7) (2) A A()
B - 1)A (3) BB(1)
(18)
In nutshell, Mamba can be seen as N, SSMs, one for each input channel. Weights for these SSMs are shared and depend on the input at that recurrent step. Because of the dependency on inputs, Mamba can model non-stationary sequences. Weight matrices are initialized as follows:
WB ~ (0,0 E N.N.,, Wc NVO, I) E RNN W ~ WN II E RIXN.,
and zero initialization for all biases. For each SSM, ( is still initialized according the HiPPO theory. When ( is real-valued, we let log [j] = log(j +1) and To can be seen a bias term initialized to To Softplus -1 (0.0001 0.1)).
Below we restate Propositions 3.2 and 3.3 followed by their proofs.
Proposition 3.2 (Scale of hidden states x1 in Mamba at initialization). Under the ZOH dis- cretization procedure, as Nu then N,, approach infinity, for any E [Na], the squared 12-norm of the hidden states ||x{²) (8) [|2 is a.s. scaled as (xx{( 112 e (c(2) of |ull? (u) (u()))). where ((2) denotes the Riemann zeta function at 2.
Proposition 3.3 (Scale of outputs y of Mamba layer at initialization). Under the ZO11 discretization procedure as Nu then N₂ approach infinity, for any i E [Nu]. the output vi converges in distribution Gaussian with mean 0 and standard deviation Cabac || 2 for some width- independent constant C >
Proofs for Proposition 3.2 and Proposition 3.3. For the following ID linear state space model:
(i)
A (1)) (6)
x B
(3) (0)
u (8)
R, x (6) e RR..
(19)
11
xTx) V/ E.R.
(1) (6)
17

---

Following condition (INIT), the ID SSM admits stability at initialization, when the following condition satisfied:
If 1 1 E (1) then |xx] ( E ((v and of E @(1).
For simplicity of exposition, let's begin by considering the scale of X1 = (A)
(A ( "i
(6) before generalizing to arbitrary I [Z].
Since B1 = W yu + by with by = at initialization,
N₂
(2)
(A (10) I)(A (9) W nu wuru (1)
Σ (A:) ( (u)))((
Σ
W (u)")) ))²
where A, (A ()
Note that the inner summation can be expressed
Nu
m. (u) 1² ( (WB), 2 (2)
(1)
(n') me (ne) un
n
Since WB B has i.idd Gaussian entries and due to Assumption 3.1, as N, - oc, the first term behaves according to law of large numbers and converges almost surely to o /r/lu The second term behaves according central limit theorem and converges in distribution
Gaussian with mean 0
and variance of
N.
(nl) (n"). Therefore, N. co, ui uu
N. (Σ
W (u)"))
converges to
Gaussian distribution with mean o |uu and variance B
N
ui
(x)) (x")
(w))
For applying Kolmogorov SLLN (Theorem A. note that ((2) S
E (A:) S
Em (+-1e S ((2) and c. ((4) S Em= (Ar) S > (++1f S ((4) for some constant E (0, and for N₂ large enough (see (37) for more details). For random variables Um No C²) for some C > O, it holds that Var((A,) (A,)), Pm) I (AA)), un I (A:)m,, + C²), so that the Kolmogorov condition is fulfilled, and we have that N₂ then N₂ approach infinity, for some width- outth-independent constant > O,
N
Na
(A (m) (1))(((
Σ
W B) (u))))) a.s, (u
|| (
Therefore, for stability at initialization, the scale of initialization OB E E0(V N
Scale of t/ ( When -
-Cix(")
where x("
B((0) u0 and we have
No
(6)
(4) (A.m.
Σ
(W Wo)m. ( u ( ( (K)
Vi
j,k=1
Applying the Lindenberg-Feller CLT (Theorem A.2), we have that, as N,, then N₂ approach infinity, u1
N
(4)
(A4),,
(Wo)ms
(elmm(WBBnmt WBBm,,uu (1) (k)
ui
converges a Gaussian distribution with mean
3.h== (3) 0 and variance e((2) OB0C 1ul
(6) Accordingly, imposing the conditions for stability of yi requires that ac E @ N.
(i) Stability of xi and d for arbitrary E [L]. For all E [4], we have x(")
Σ (A (i) ym B (1)
(8) fm um First, observe that since A' diag(a',...,, d'y.) with d'
(cos(T,' Tn) + i ( Tn)), we have that, for all complex vectors CNo it holds
18

---

that for any m [7], II(A')"vvl2 e-mmr
Therefore, the operator (A)) does /21v///
/21/
change the width-scaling. Since, for any ui setting oB E @ N.. yields B (0) 4 2
((vN2)-
Therefore, each term in the summation is of order @ and unless, for every I, the term BB()) iy perfectly cancels out with the terms before to affect the width scaling, we have that ||x] ( The same argument can be used to show the stability yt (()
|| O ((vN)).
Proposition 3.4 (Scale of the updates of hidden states Ax after step of SGD). Under the ZOH discretization procedure, as Nv then No approach infinity, for every ID SSM, the squared 12-norm of the updates ||Ax] (6) ll2 of the hidden states after one step of SGD is a.s. scaled as ||xx|) (8) 112 E ④ (IB 1 aC ||n 112 <(4)t). where ((4) denotes the Riemann zeta function at 4.
Proof of Proposition 3.4. Following condition (A), the features of the ID SSM evolve non-trivially with width, when the following condition satisfied:
If Au (
(1) then | (O(VN,) and Av) 0(1).
Scale of the updates Any after 1 SGD step. First, note that, for the discretization step ZOH, we have
x = D1
(0)
C
A (6)
((() A (
(0) + (7()) (0) A 6- (A (A (B, B 1
Nu
A (1)
B
A (
() (1) 'B (0)
1=1
B B(0
(1) (1)
4 X
Diag(
A (1) A R
log
N
WB B
C
01
X
(i)T
For any quantity letting denote the updated quantity after one step of SGD, we can write
A (i) = Diag(e exp (alog - Ao log = A ( exp (-n. (3) 1).
B, Wn nu = (Wn - 1) Buu (u/ + Au:).
For clarity, let us again begin by considering the scale of x We can follow the same argumentation as before to generalize the result arbitrary I [L]. The updates Ax(" can therefore be computed
as B (i) B "1 where B( ((A (0)) () )0) B, and A' exp(rÁ())) Since
(4) and in are assumed to be @(1) due Assumption 3.1, the scale of Ax is determined by that
"
of BB0 and B ((0) B(O)
(
Scaling B(()) Recall that B(O0) = (i)
= (A ()(A()))((W,
AWa B) (u) + Aun).
Let's begin by understanding the scale of (A I Since AWB
-Mn BuuB letting Å, (A (A)--) we can express the 22-norm of (Ā (1)
(A ()) JAW zul follows:
| (0) - ()(A nu || = a
= |nal [aml||A, ||
1l | 12
Note that since Á, W. B is a rank-one matrix, its spectral norm can be decomposed as MAWWB,ll |l. = Inal || 11 |1ll|
19

---

Therefore, since ||ul|l2 E (vvN,), if the spectral norm of A, AWB W is scaled as 0 e(/x.)
(A I)(A()) gu
Claim. If JAWalll. E Φ(y N: then (A (A gu
(
Essentially, this follows if (À I) (0 does not affect the width-scaling of the spectral norm of
AWB. B. that is, ||,AW all. = | This is equivalent to showing that | | 2 = - 2'
Note the following Lemma which states that the scaling of [|A, 122 is identical to that of | B 2
Lemma C.3. Let Qi (exp A A (.). For any complex vector E CNN, the width-scalingg
v remains invariant under the operator (A (() -I) that is, [| (() I)Q7 @(1lv]12).
Proof of Lemma C.3. Recall that A; = (A (1) - I),," where A, = A (8) Q and
Q = (exp (-n. n/ (8) 11 I) = cxp Ta ((7 A (A 61) (A - [B) B (TT
o o))) OI I.
Since [A [1] -j + 1), it's easy to see that, for any positive value of 7(),, the eigenvalues of 7(1) A (i) (A) (1 (A - I) are negative, in (l),, and converge to 0 from below. The expression
evaluates to
exp(-70) (1) -1 (i+1)
700 A (1)
7 exp(->(0) 1))
Since B1()) = Ui and x
(1)
D1 CT.
(1)
No
No.
[ B (0T lis
W wlijm (m))( VIa
Σ
[
u
(ml')
yui Di
(1)
Due to Assumption 3.1, Vi (0) is distributed i.idd as N₂, Do and Vi (1) E 0 ( (Yang and Hu, 2021). A CLT argument reveals that, under the scaling ds E 0 and oe @(V N, the eigenvalues of (B"()) (T) @ I) are in O N.. Lemma C.4 follows as consequence of this result.
Lemma C.4. As No then N₂ approach infinity, (|A()) A ( ||. (11) and || ||. 2(11) if and only na E ((N.).
Accordingly, all the eigenvalues of Q are in (1)) and therefore, for any complex vector v, ||Q,, E @(Ivvl2). Together with the fact that the spectrum of Ar(i) is @(1), we have A, E (()). Further note that the entries of (0) and consequently those A, are independent of each other since the entries [ (+)
B (6)T)
Im. are independent for different m [Nz].
Σ (A (A ( () it follows from Lemma C.3 that | ||, |B
Since B,
Therefore, for the correct scaling of the updates, Me should be scaled as O Ax) ). Similar computa- tions reveal that nic should be scaled O (,
Proposition 3.5 (Scale of the updates of outputs Ay ( after 1 step of SGD). Under the ZOH
discretization procedure, as N,, then N₂, approach infinity, for every ID SSM, the squared bz-norm of
the updates of the hidden states after one step of SGD scales as م (6) (nooby /,, || (1)
20

---

Before we present the proof of Proposition 3.5, we first provide the following Lemma.
N
(t) U1 converges to normal distribution with mean 0 and variance
Lemma C.5. As Nu
Na Var(y) (1) ( Note that since we assume that the rest of the network is scaled correctly (under μP), V1 (1) distributed iid usymptotically and Vi E ((())) (Yang and Hu, 2021). Therefore
(0) u1 =
Proof. For i E [Nu]. the output of the 1D SSM is given by
= Cixi where C₁ = (Wou)) and xi B un (A - I)A ())
W WBu1. (20)
Using the notation to denote the updated quantity after step of SGD,
vi x = (Woūn) (A, wiū
(21)
The updated quantities W and We are derived as
WB WB WB+AWB, Wc= Wo+AWo with AWB = - WB,
=
= - We,
where W B and We denote the gradient of the loss with respect the quantities WB and We respectively and can be computed according to
N Σ WB
N
91
U1
A. W Wunu T and Wc =
Σ
9i 21
A
uu
(22)
t=1
-=1
Also, A; (A7() ()
(1)
where A Diag(exp (a
log
1/e A
A with
Q exp (- A (8) A (1) @ and
( = (7 (4) (0) -A(0) (I) ( @I).
Therefore V1 (1) can be represented as follows
= [WC - AWc)(u, + Au,) )] [A,(W + +
+ {u + Au("), (23)
and the updates Vi vi (6)
can be expressed as
V1 vi (1) (Wau)) A W.Wuu - (Weu)) A, W Bunu + M,
(0)
where M = [W c - A AW)(uu + Au))] T [A + AW (u, + Au:) + Au("))
(Wou)) A, W W ( (Weui)? A, W Wauni represents the remainder of the terms in (23).
First, let us consider the scaling of the update term without remainder term,
N
Nu Wcm. W Bmmm (mn' (om'mn (m²"
(W A,) W.)Wuurll
Km
m=1
mI mm
where K. =(A,, -4)).
Letting V... _m
W WCm. W Bm, mm Un (ml) (m²),, (m²") N,, 00. {vm N are independent
and are distributed normally with mean 0 and variance o 0 || || by application of Lindenberg- Feller Central Limit Theorem.
Independence of entries of K. Next, we show that the entries of the diagonal matrix I. are independent of each other.
Observe that
K CAAAi -.) (ATO) 1)W (A ())A
21

---

where A Diag exp (a; log noA (og)) = A Q.. with Q; ( eexp -/ A (A (1)
1) and
A (1)
(7 A (6) (A - (A (A (B) B (iT
(0)
+ I).
Since the entries of A are deterministic, and assume that 7 is some fixed deterministic scalar, the entries of K are independent of each other iff the entries of (B( B (i) I) are independent. Now consider
Nu
B
B1 (17)
W Wmmml W Cmmm" (m)) (mn") uu
Since WB and We are Gaussian matrices with i.dd entries and the entries of u, are assumed to be ii.d, the entries [B]1) Bj
(8)TT (BOB,
Im.m are independent for different m E [Na]]
Furthermore, following the same line of argumentation as in Lemma C.3, the matrix K. can be represented as = K'' A1, where the entries of K' are (1) and accordingly for any vector the scaling of remains invariant under the operator K".
N₂ Km.m Vm. we can therefore instead consider the scaling of m=1
No Σ
To compute the scale of
Σ m=1
A Vm
where {vm
Mm are independent and are distributed normally with mean 0 and variance o (u
00, c((2)o un _ where ((2) denotes the Riemann zeta function at 2 and e<(2) evaluates to width-independent constant.
As N₂
A=1 m.m Vm converges to normal distribution with mean 0 and variance
When on E @(v No and oe E O ((vR.N.) we get (Wou)) -, A,) ( e @(1)
Scale of M..
Note that since uj and Auj" are scalars and in order 1, the scale of M is determined by that of [(Wc AW + ( AW (u; + Au1)].
Recall that both uı and Au; assumed have i.i.d coordinates and have the scaling ||u,|| |l e @(vN.) and Auj|l 0( (2)).. Therefore, we only need to consider the scaling of
[(Wc A
AWB))
It is easy to verify that the scaling of ( - A,) W.)W Wau is the same as that of (Weu,) (A,) uu following the same arguments as above.
Scale of (Awcu)) ) W u
Recall that the updates AW B and AWo are defined
AW B = 7B W and AWc = -1c
- We,
where WB and We denote the gradient of the loss with respect to the quantities Wa and Wo respectively and computed as
N₂
N
W B =
D
(4) ui
A. T and Wc =
Σ
D u
( A.C.T
(24)
No
No (A, is iden-
So, AWAu, = -nc
Σ
Di U1
(1)
A, W
u = - 2A
W u
Σ
(8)
1=1
tical for all i E [N,))
No
Therefore, letting E -
un
2
Σ
V1 ui
(6) (1)
we get
i=1
22

---

(AWou:) (A, W Bui) = E(A, W W ) (A, WBu))
N
A Asmmm immm
W mm W Bm.mn (mm) uj (m")
m=1
No W Bm.m' W Bm.mn (m²) uj (mn"), it is easy verify that as No +
Letting each is distributed normally with mean of /| and variance of Furthermore {Wm] N. are independent. Therefore, due Kolmogorov's SLLN, as N₂ oc, it holds that
N. m m/="m==
uj (m') un (mm")
N A Ammmm m WB 3 3 1
V7A
Combining this with the result from Lemma C.5, we conclude that as N₂ - oo then N₂ + oc. (AWcu,)) (A, WBu) Ofnco VN.))
Therefore, for the updates Ayj (1) to be order 1, ne must be scaled as O N, N.
Scale of (A (AWour cuu) (A.))
Recall that the updates AWp and A AWo are defined as
AW = 7 WB and AWC --1c -/c We,
where WB and We denote the gradient of the loss with respect to the quantities W a and We respectively and computed as
We
1 ( 9 1
A. and WB
9
i1) (
T.T
(25)
N
Hence, we get AWou: = -1c
Σ
yi A W Tu - W W Σ
91 mi
iml
(A; is identical for all i [NV..) and AW Bun A, A.u,
vi
uj
which yields
(A mu)) = ② (0) Mo ||
aA, (A
Following the same line of argumentation as before, we obtain that as N,, 0o then N₂ -> co, E 6(0eecnssecllls] N.).
yields the same scaling: no Substituting оB E Ø(v N. oc e O O ((xc)))) n O N.
VN.
O N. V.
Scale of (T Bun Following the same steps as for the previous terms, we obtain
(Woun) I w (8)
O(πs A o || [1).))). This term vanishes under the scaling di
O N and oc E O VN
1.).
23

---

C.3 Time-invariant S4 recurrent layer
The S4 recurrent layer VI.L = fualu (B, C}) is sequence to sequence mapping:
X A A'x1-1 B'u;
(26)
y1 R[C'xi]
(27)
for I = 1...., L with Xy = 0, which is a discretization of the continuous-time SSM:
d dt x = Ax; + Bu,
(28)
y: = R[Cx.]
(29)
where R(-) gives the real part of complex vector, A diag(a) and E C.2,, B CN.XN..
Ce C.N, x
The Hero-Order-Hold (ZOH) discretization method gives that:
A' exp(T) (9) -A) B' = (A' - I)A BB
(30)
(31)
C'=C
(32)
To model long-range dependency, S4 advocates for HiPPO theory. While our results hold for every
Hippo initialization matrix, here, we show the result for S4D-Lin and set a, = = + itn. Note that A is not parameterized and will not be trained. It has been observed that complex values are not essential and S4D-real with a, = -(n + 1) may perform similarly well.
The backward pass for fsa is as follows:
X = CTTy + A
(33)
L
C²
Σ
y(X]
(34)
1=1
L
B' Σ
Xuu
(35)
u. B
(36)
Since we want all learnable weight matrices to learn features, we split each $4 layer into two sublayers and demand them both admit feature learning. The first one being the state space equation (26) and the second one being the decoder (27).
Claim C.I (Seale of hidden states xı in S4 at initialization). Under the ZOH discretization procedure, as N₂ and N,, approach infinity with N,/N, E ( for any i E [Na]. the squared
12-norm of the hidden states (0) 2
is almost surely scaled as x E θ (c(2) Ilull"). (lul)) where 2
(1)
2
((2) is the Riemann zeta function at 2.
Proof. Stability at initialization.
First, we show that the size of xi = B'u, at initialization is @(vNa). By definition,
x ll = B u 1|2 = ||
where A = 3(A' -A'--) (IA-1) Then, writing un for readability, observe that
B'u
22
a sum of N₂ independent not identically distributed random variables. However, we show that they satisfy the Kolmogorov condition and therefore the sum behaves according to the strong law

---

of large numbers. Intuitively, this is allowed by the polynomial decay of eigenvalues of A and the correct scaling of the imputs to the S4 layer.
Verifying that the sequence satisfies Kolmogorov condition. For any i E [[v]]
No B
N₂ E[A}
Var(A, Σ j=1
Σ
> j=1
B uyy
j=1
N, B1,2,1
- E A
Σ
j=1
N₂
= E
Σ
Bey Bu uks up!
kJk
N2
N,.
A A1((ET E
B
E
B
j=1
j*1'=1
N2
N,.
A (30
2
+ oB
Σ
u5u3.
5=1
1/3'=1
S 300 ||u||
N₂
Nu.
N₂
Σ Var(A?(())
3 ||u/|/ Besu,)
1
j=1
N₂ (exp(- 1T) ((0)) 1)
30 ||ul/2
(i+1)
N₂
S $30'3 ||u||
( 1)
30 || C(4),
(37)
where ((4) denotes the Riemann zeta function at for which closed form is given by (2x))
2.4!
Note that, due to the rapid decay of exp(-(i 1) (8) + 0 with - co, the zeta function also lower bounds the spectral sums
A3 2 ((2) and A ((4) for some width vidth-independent
constant E (0, 1) for N₂ large enough.
By assumption, the inputs to the S4 layer are scaled such that ||||| E @(vN,) and therefore, as long as oB E O(1) the sequence of random variables satisfies the Kolmogorov condition.
Limiting behavior of [B'ul
Applying SLLN, we have
No
No
N
No F.
Bu
A?
B
==1
N₂
N,.
Σ
A
E
BL n
1=1
2
Σ A3 Σ 2
o
1=1
=2( |ul2
25

---

where ((2) denotes the Riemann zeta function at 2 and equates to n² n2/6, and E (0, is some indeh-independent constant.
Spectral scaling does not yield the correct scale of initialization for d-
For spectral scaling conditions to yield the right scaling of the initialization variance, it is crucial that the following condition holds:
[Buul|2 E O l1))))
Since the spectrum of (A' I)A- is less than 1, an upper bound on the spectral norm of B' (A' A 'B can be found in (Vershynin, 2011) and is given by C(vN, + vN,) for some width-independent constant C.
A matching lower bound can be easily found by noting that the spectral norm is lower bounded by the maximal row and column norm.
BBIl. 2 max max || max
|| |B:y)lb)
2 max |B/,1l
AA, max
B
No. AA -max
B
j=1
2 max A V NNOB
2 / /n
|B'u| Therefore (I.
E 0
2
VN_
Correct scale of og for stability at initialization.
We know that BB'ull2 E O(aB ||ulll)) and since ||u|| ((vN.) (since we assumed stablity of activations in the previous layer), imposing stability of [B'u Blull yields the following scaling:
OB 0( N N.
Stability of x; for arbitrary / E [Z].
A' (A'm"B'u ()'u m) + B'u, First, ob- alx.) with al, = 1 (cos( (T Tn) + issin ( (1) Tn)), we have that, for all complex vectors v E C. it holds that for any m [L]. || (A'''vvl = m / || Therefore, the operator (A')m does not change the width-scaling. We showed earlier that for any ur, setting aB @(v Ns N. yields ||Buujll2 ((VN.). Therefore, each term in the summation is of order O (vVN)) and unless, for every I, the term B' B'u, perfectly cancels out with the terms before to affect the width scaling, we have that E @(VN,) for all [Z].
For all / E [Z], we have x = (A')" ul (A') serve that since A' = diag(a/
Claim C.2 (Scale of output ( of S4 at initialization). Under the ZOH discretization procedure as N2 and N,, approach infinity with N,/N, E @(1), for any [N.), the output y converges in distribution to a Gaussian with mean O and standard deviation Casac || for some width- independent constant C.
Proof. Correct scaling of updates (and learning rate) can be achieved by correctly scaling the spectral norm of the updates of weight matrices (an and ac).
To give concrete update/learning rate scaling rules, we need to choose a concrete update rule. We first consider SGD as an example. First, we show that while spectral scaling of weight matrices does
26

---

not imply the correct scaling of (pre-)activations spectral scaling conditions on the updates of weight matrices imply the correct scaling of the activation updates.
To simplify cumbersome notation, lets first consider the scale of updates x1 U after a single step of SGD on ((u1 91)).
LABuiil; ||AABuill ||Anyiuy un |n) ||11|// || ll2
Since the update matrix AB' is rank matrix, we have
AAB|l I1, ||AB|||
| [Ax::
")
lol ||| 1ll ||Axi
The updates Ax₃ are given by Ax, (B)) + Auj) - B'u; B'Au, + AB'u, AB' Auj.
If the spectral norm of AB' is scaled as @
N= N.
then AB'u, E ((vN)). We also have that
AB'Au, I ||A'II.. Auull E O(vN.). Since | (VVN,,, when dg is set Ø B'Au, E ((vN,,) Therefore, unless the scale of AB'u; u1 perfectly cancels with the remaining two terms, | E @(vNs).
Generalizing to more gradient steps (and / 1) follows immediately if we assume that updates do not perfectly cancel out initial quantities, combined with the fact that A' does not change the width scaling.
D Additional Experiments
We conducted additional experiments to validate our theoretical derivations and to compare with baseline parameterizations for SSMs. Focusing exclusively on the SSM components, as parame- terization for other parts of the model are already addressed in Yang and Hu (2021), Yang al. (2023a), we decoupled the learning rates for SSM and non-SSM layers. First, we tuned the learning rate for the non-SSM layers and then evaluated test performance across various SSM learning rates, using the optimal non-SSM learning rate. The results, presented in Figure 4, show that the jP-SSM parameterization exhibits better monotonicity and stability.
SP
pP (Heuristic)
JPP-SSM (Ours)
3.0
20
512
2010
277
4335
Test
28
25-
23 -10
-30
log, (ITS5M)
log (IM)
log, (Ir ssar)
Figure 4: Decoupled learning rates. Test loss against SSM learning rate in Mamba with varying
widths (N, and N₂) on the 4M tokens sampled from the WikiText- 103 dataset. The learning rate for non-SSM components is fixed and chosen via byper-parameter tuning.
In Figure 5, we present results on a randomly sampled subset of the Fineweb dataset. While computational constraints prevented from training the entire dataset using larger model widths, at small scales, our observations on Fineweb align with the wikitext-103 results reported in the main paper.
27

---

SP
JPP (Heuristic)
pPP-SMM (Ours)
12.
Test loss
80
log,
log (Ir)
log (1r)
Figure 5: Results on the Feb dataset. Test loss against SSM learning rate in Mamba with varying widths (N₂ and N₂) the 20M tokens from the Web dataset.
28

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We state all our contributions while acknowledging related work. All main claims are theoretically proven or empirically verified.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
Limitations are discussed in Section 6.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]

---

We state all assumptions in the main paper and provide all formal proofs in the appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer:[Yes
Experimental details relevant for reproducing the results to the extent that affects the main claims are disclosed in the main paper.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution primarily new algorithm, the paper should make clear how reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution a new model (c.g.. a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
30

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We only propose a correction of the width-dependent scaling of hyper- parameters of existing architectures, and precisely specify the corrected scalings in the main paper.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
cublic/guuides/CodeSuubbat sotiibmisssooooccccc for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: / //iiss... /tipps.cc/ppiii wtaliiigguies///eeee Sodeeee//oodSSumiiiccc
onPolicy) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permirted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [No]
Justification:
Some experimental details are disclosed in Section 5 but not all details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: In Figure 1, we report a-confidence intervals over 10 runs. Due to limited
computational resources, we were unable repeat the large-scale experiment for Figure multiple times.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emoo bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
31

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: All experiments run within 24 hours on 24 NVIDIA A10G GPUs.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https ptttps//neerre http:///neui.pp.eee cc Cwwbblcc/Ehisscuieeee
Answer: [Yes]
Justification: We provide a theoretical analysis widely used architectures and propose a correction. We do not foresee any ethical concerns.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper provides fundamental research toward understanding and improving existing neural architectures. We do not release any model data.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
32

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), faimess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: We only use standard datasets in our experiments and do not release any data or models.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging. and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
We cite the code and data assets according to standard practice.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include a URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
33

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. www.coo///aoo//o datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
34

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
35


Reviews:
[{"review_id": "sCWhFeq0xz", "reviewer": "Reviewer_QtDa", "summary": "Following the tensor program and maximal update parameterization, this paper studies the parameterization of initialization and learning hyperparameters in structured state space models (SSMs), e.g. S6 Mamba. The authors consider the input dimension and latent dimension of each vector in the sequences to go to infinity and analyze the proper initialization and parameterization for hyperparameters such that the initial output at each layer is stable when passing through multiple layers and the feature updates are compatible with initialization when taking one step gradient. The paper provides a detailed analysis of signal propagation in SSMs, both in the forward and backward passes, as the width of the network increases. And it also reveals that established scaling rules, such as the maximal update parameterization and spectral scaling conditions, fail to maintain feature learning in SSMs.", "questions": "1. Typo in Line 47: $\\mathbb{R}_u^N$\n\n2. Can you provide more motivation for defining feature learning in layer-wise sequential models in Definition 2.1? In (1) and (2), why do you only require the existence of one output satisfying stability and feature-update assumptions? Do we need to ensure these bounds for all outputs?\n\n3. In Line 70, you mentioned for feature learning we need the inputs to the layer to be asymptotically independent and identically distributed. Does this mean we need to assume $u_1,\\ldots,u_L$ are asymptotically i.i.d?\n\n4. Line 76: typo ''Parmeterization''\n\n5. In (6), $W$ should be $W_l$\n\n6. Explain the functions in (8-10) for completeness.\n\n7. Below Line 172: typo $\\sigma_B^2$ should be $\\sigma_C^2$", "limitations": "N/A", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper is the first paper studying hyperparameter scaling for infinite-width SSMs, a topic that has not been extensively explored compared to other neural network architectures like MLPs and CNNs. And this work has practical implications for improving the training and performance of large-scale SSMs and sets the stage for future research in this area. By tackling practical issues such as vanishing or exploding gradients, the paper provides solutions that may enhance the stability and efficiency of training state-of-art state-space models, making it a valuable resource for practitioners.", "weaknesses": "1. There should be more experiments and empirical comparisons among standard parameterization, maximal update parameterization, spectral scaling, and the $\\mu $P SSM parameterization, in terms of test loss and parameterization transferability on different types of datasets and learning tasks. The only empirical experiments in Fig 2 in the paper seem to indicate that spectral scaling performs similarly to $\\mu$P SSM parameterization. It would be more convincing if we could compare these parameterizations in various situations.\n\n2. There is a lack of explanation of why we need to consider $N_x$ and $N_u$ go to infinity. In terms of SSMs, we usually consider the length of the sequence to be large and SSMs can preserve the long-range dependency. Besides, in the analysis, the authors only consider the inputs of the layer i.i.d. which is quite different from the practical sequential dataset. It requires more explanations for these assumptions and the definition of feature learning in this paper.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive feedback and positive evaluation of our work. We appreciate your insights and we have carefully addressed your major concerns in the following response. We will rectify minor issues like typos in our revision. We hope that our revisions warrant your consideration for an increased score.\n\n**Additional Experiments.** We have conducted additional experiments to further validate our theory. We have summarized them in the global response and the results are provided in the attached PDF. Addressing your concerns, we conduct experiments on a randomly sampled subset of the Fineweb dataset which suggest similar conclusions as those on wikitext-103 in Figure 2 (main paper). muP-SSM outperforms SP in terms of test loss, training stability, monotonic improvement of generalization with scale, and transferability of optimal learning rate across scales. Observations in comparison to spectral scaling are more nuanced. muP-SSM markedly improves training stability at large learning rates over spectral scaling, and slightly improves generalization performance. First tuning the learning rate in non-SSM layers and subsequently tuning a separate learning rate in SSM layers separates muP-SSM and spectral scaling more clearly: While muP-SSM continues to monotonically improve with scale, spectral scaling no longer improves beyond a certain width, suggesting that the reasonable performance of spectral scaling stems from the non-SSM layers. In SSM layers, the spectral scaling approach lacks rigorous theoretical justification.\n\n**Clarification on the asymptotically i.i.d assumption.** In Section 2.2 & A.1 of our paper, we discuss that activations and their updates in neural networks representable as Tensor Programs (TPs) become asymptotically independent and identically distributed (i.i.d.) with increasing width. This result underlies our i.i.d assumption. Note that, we consider the practical setting where SSM layers are embedded within a network of non-SSM layers (such as MLPs or Layernorms) representable via TP. Accordingly, it's natural to assume that inputs to the SSM layer (e.g., activations from the previous non-SSM layer) are asymptotically i.i.d at each recurrence step. In fact, this i.i.d. assumption isn't strictly necessary. It suffices to assume that inputs are correctly scaled, i.e., $\\vert \\vert u\\vert \\vert_2 \\in \\Theta(\\sqrt{N_u})$. This scaling has been demonstrated for correctly scaled network modules representable via TP. Crucially, we do not assume that the sequence of inputs $u_1, u_2, \\cdots, u_L$ are i.i.d. Rather, we assume that the coordinates of a single input $u_i$ are asymptotically i.i.d. (or at least correctly scaled) in the sense described above.\n\n**Why should N_u and N_x go to infinity?** We would like to clarify that our analysis **does not require** both $N_u$ (input dimension for SSM component) and $N_x$ (latent states dimension for SSM component) to go to infinity. It merely allows both quantities to scale up. To derive the results for when only $N_u$ goes to infinity, one can simply set $N_x \\in \\theta(1)$. In this simplified scenario where the dimension of the latent states $N_x$ is fixed, heuristic muP spectral scaling yields $\\eta_B,\\eta_C=\\Theta(\\frac{1}{N_u})$ while muP-SSM gives $\\eta_B,\\eta_C=\\Theta(\\frac{1}{\\sqrt{N_u}})$ (see Table 1 in the main paper). Experiments carried out independently by researchers (to be acknowledged in the public version), indeed verify that the correct width-independent scaling aligns with muP-SSM (results as shown in Figure 4 in the attachment).\nNote, however, that as models are scaled up both $N_u$ as well as $N_x$ may be increased in practice, albeit at very different scales. For example, see Table 10 in the original Mamba paper [1].\n\n**In Definition 2.1, why do we require that there exists one input that is correctly scaled?** This follows the same logic as the paper that proposed muP [3, Definition H.9], which defines a parameterization to be feature learning iff there exists a training routine and input that result in the correct update scaling. This is a technicality we have to adopt as there exist degenerate combinations of inputs and learning rates that result in smaller scalings. Even under this definition, there exists only one choice of layerwise initialization and learning rate scalings that achieves feature learning.\n\n[1] Gu, Albert, and Tri Dao. \"Mamba: Linear-time sequence modeling with selective state spaces.\" *arXiv:2312.00752* (2023).\n\n[2] Wortsman, M., Liu, P. J., Xiao, L., Everett, K., Alemi, A., Adlam, B., ... & Kornblith, S. . Small-scale proxies for large-scale transformer training instabilities. *arXiv:2309.14322 (2023)*.\n\n[3] Yang, Greg, and Edward J. Hu. \"Feature learning in infinite-width neural networks.\" *arXiv:2011.14522* (2020).", "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": "LhwgZWOEhR", "reviewer": "Reviewer_JuBx", "summary": "This paper investigates the scaling behavior of state-space models (SSMs) and structured variants like Mamba as their width approaches infinity. The authors demonstrate that established scaling rules such as maximal update parameterization (μP) and spectral scaling conditions fail to yield feature learning in SSMs at infinite width. They derive a new scaling rule, μP SSM, enabling feature learning in SSMs as width approaches infinity. Empirical results show that μP SSM leads to improved stability, generalization, and hyperparameter transferability compared to standard parameterization and spectral scaling.", "questions": "1. Can you provide a more intuitive explanation of why SSMs violate the tensor program ansatz?\n2. How do you expect these results to generalize to other SSM variants and tasks?\n3. What are the computational trade-offs of applying μP SSM scaling in practice?\n4. Can you elaborate on the implications of spectral scaling leading to parts of the SSM being in the lazy regime?\n5. Have you explored whether μP SSM enables training of wider SSMs than previously possible?\n6. Can you outline a path for extending your analysis to the proportional limit case?", "limitations": "The authors discuss the key limitation of restricting analysis to the N_u then N_x approaching infinity setting. However, they should more explicitly address limitations of their empirical evaluation and potential challenges in applying results to practical scenarios.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Addresses an important theoretical question about SSM scaling behavior\n- Provides rigorous analysis of forward and backward signal propagation in SSMs\n- Identifies limitations of existing scaling approaches and proposes a principled correction\n- Empirically validates results on real SSM architectures like Mamba\n- Has potential implications for training larger, more efficient SSMs", "weaknesses": "- Theoretical analysis limited to N_u then N_x approaching infinity setting\n- Narrow scope of empirical validation (text generation with Mamba on a single dataset)\n- Lacks comparison to other recent SSM variants beyond Mamba\n- Insufficient discussion of potential negative implications of enabling feature learning in larger SSMs\n- No clear roadmap for extending results to more practical settings", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to express our gratitude for your thorough and insightful review of our manuscript, and useful feedback for improving our work.\n\n**Order of limits and practical applicability.** In SSM models, $N_u$ (input dimension for SSM component) typically increases much faster than $N_x$ (latent states dimension for SSM component) during scaling (see Table 10 in the original Mamba paper [1]). This makes the limit where $N_u$ approaches infinity before $N_x$ more practically relevant. Furthermore, fortunately, these limits commute which allows us to extend our results to the proportional limit setting.\n\n**Additional Experiments.** We have conducted additional experiments to further validate our theory. We have summarized them in the global response and the results are provided in the PDF attached to the global response. Addressing your concerns, we now also present results on a randomly sampled subset of the Fineweb dataset, suggesting similar conclusions to the wikitext-103 results in Figure 2 (main paper). Separately tuning the learning rates in SSM layers and non-SSM layers shows the benefit of muP-SSM over spectral scaling more clearly: While muP-SSM continues to improve with scale, spectral scaling no longer monotonically improves beyond a certain width threshold. Even without separate learning rates, muP-SSM markedly improves training stability at larger learning rates over spectral scaling. SP consistently performs worse than muP-SSM in terms of stability and generalization. This indicates that muP-SSM can indeed improve performance at larger scale.\n\n**Implications for other SSM variants.** Mamba is among the most complex SSM architectures, containing all components present in other recent SSM variants. By providing a scaling analysis for Mamba analogous analyses for S4 or LRUs follow as corollaries. We will clarify implications for other SSM variants in the revision (see our answer to Reviewer iUa6).\n\n**Implications of enabling feature learning and computational considerations.** We see no negative implications in enabling feature learning in every layer. If feature learning is undesired in a specific layer, training can be explicitly disabled, saving computation instead of letting the updates to implicitly vanish with scale. muP-SSM therefore provides the flexibility to achieve feature learning in SSM layers when desired. muP-SSM introduces no additional computational complexity, as it's merely a different parameterization of the weights.\n\n**Have you explored whether μP SSM enables training of wider SSMs than previously possible?** While we lack the computational resources to experiment at such large scales, there's evidence suggesting muP-SSM could enable training of wider SSMs. Even at relatively small scales, muP-SSM shows markedly improved training stability at larger learning rates compared to spectral scaling. According to [3], instabilities in small-scale models at large learning rates can predict general training instabilities at larger scales. This evidence provides some support for the possibility of training wider SSMs with muP-SSM, though direct large-scale experiments are needed for confirmation. We hope that the SSM community would conduct experiments at larger scales to further investigate this potential.\n\n**Intuitive explanation of why SSMs violate the TP ansatz.** There are two primary reasons for why SSMs such as Mamba cannot be represented via TP:\n\n1. **Structured transition matrix (A).** Typically, the state transition matrix A is highly structured and chosen according to (or loosely based on) Hippo theory [2]. One example is the diagonal matrix with the $i$th diagonal entry being $i+1$. The TP framework crucially relies on matrices with i.i.d entries (such as i.i.d Gaussians) and cannot represent such structured matrices.  \n2. **Selection mechanism.** Activations and weights play different roles under TP. The selection mechanism first computes activations (linear transformations of input parameterized via some weights) and uses them as weights in the linear recurrence (see the pdf attached to the global response for a visualization). This underlies a second source of incompatibility of SSMs with the TP framework. \n\n[1] Gu, Albert, and Tri Dao. \"Mamba: Linear-time sequence modeling with selective state spaces.\" *arXiv preprint arXiv:2312.00752* (2023).\n\n[2] Gu, A., Dao, T., Ermon, S., Rudra, A., & Ré, C. \"Hippo: Recurrent memory with optimal polynomial projections.\" *NeurIPS 2020*.\n\n[3] Wortsman, M., Liu, P. J., Xiao, L., Everett, K., Alemi, A., Adlam, B., ... & Kornblith, S. \"Small-scale proxies for large-scale transformer training instabilities.\" *arXiv:2309.14322 (2023)*.", "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": "xCnYwhagcX", "reviewer": "Reviewer_iUa6", "summary": "This paper studies the large width scaling behavior of a recently popular class of models known as structured state space models (SSMs). The authors demonstrate that the previous work on large width neural scaling, which prescribes a parametrization known as the Maximal Update Parametrization (muP) as the optimal scaling for neural networks, does not cover the case of SSMs. Furthermore, muP turns out to be suboptimal for SSMs and does not achieve the desired consequence of stable feature learning and hyperparameter transfer. The authors then propose the proper scaling for SSMs and demonstrate numerically that hyperparameter transfer is achievied.", "questions": "None.", "limitations": "Yes.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Understanding proper scaling of hyperparameters with model size is a practically important problem especially in the age of scaling. State space models are an important class of models with competitive performance and desirable attributes such as fast inference speeds. Identifying the proper scaling for SSMs is important for ensuring stable performance as SSM model sizes increases and for unlocking hyperparameter transfer which can greatly reduce the cost of hyperparameter tuning. The authors provide a fairly simple and thorough explanation for the correct scaling for SSMs which turns out to be different from the previous muP scaling. The numerical experiments give further confidence to the correctness of the results.", "weaknesses": "The presentation is at times a bit hard to follow. Some more introduction to state space models would be helpful to help readers who are not already familiar. Many terms are introduced with very little explanation to orient the reader. Some diagrams could be very helpful, so that the reader can better visualize the SSM forward pass etc.\n\nOn the mathematical side it would be great if the things could be cleaner and better organized. I think the original spectral scaling paper [1] is a great example of this. The notation and results used should be clearly established (e.g. definition of spectral norm, Kolmogorov's strong law of large numbers, etc.). It would be probably the most helpful if the proofs could be \"modularized\" into basic operations. Right now all of the analysis goes through very specific instantiations for specific models. The original paper [1] focuses on what happens for elementary operations which can be transparently composed. The results for SSMs should then follow as corollaries.\n\n[1] Yang, G., Simon, J. B., & Bernstein, J. (2023). A spectral condition for feature learning. arXiv preprint arXiv:2310.17813.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely appreciate your thoughtful review and positive assessment of our work. Your feedback has been invaluable in helping us improve the paper. We have addressed your concerns regarding improved presentation and organization of the proofs. We hope that our revisions will merit your consideration for an increased score.\n\n**Improved presentation.** To enhance the visualization of the Mamba forward pass, we have created a new figure (please see Figure 2 in the PDF attached to the global response). We will integrate this figure into the main paper and expand Section 3.1 to provide additional background on SSMs for enhanced readability. We will also introduce all notation and results used such as the definition of the spectral norm, Kolmogorov’s SLLN and Lindenberg-Feller CLT, in the appendix.\n\n**Modularization of the proofs.** Thank you for your suggestion. We concur that modularizing the proofs will certainly enhance the paper's accessibility. However, analyzing signal propagation in the Mamba architecture is inherently more complex than in MLPs, necessitating some analysis of specific modules. Nevertheless, we came up with a strategy to modularize the analysis by noting that the Mamba layer can be mechanistically viewed in the forward pass as a sequence of 3 components/sub-layers:\n\n1. Selection\n2. Discretization\n3. Linear recurrence\n\nOur visualization of the Mamba layer is based on this perspective (in the PDF attached to the global response). To our knowledge, all SSM layers comprise some or all of these sub-layers, allowing us to modularize both forward and backward signal propagation analysis through each sub-layer. It's worth noting that different discretization rules (e.g., ZOH or Euler) would require separate analyses. Beyond this, while SSMs primarily differ in their initialization of the transition matrix A, our analysis remains applicable to any initialization chosen according to HiPPO theory.\n\nWe believe these changes will significantly enhance the clarity and accessibility of our work, and we look forward to incorporating them in our revised manuscript.", "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:
Dear Reviewers,

We sincerely appreciate the time and effort you have invested in evaluating our work. Your insightful comments and constructive feedback have been invaluable in helping us improve the clarity and quality of our research. We have included a pdf attachment to this response which contains additional figures that address some of the questions raised by multiple reviewers. We will address individual points raised by each reviewer in detail in the reviewer-specific responses.

**Improved Accessibility.** Following the suggestion from Reviewer iUa6, we have included a new illustration (Figure 2 in the attached PDF) that demonstrates the forward pass of the Mamba SSM layer. This addition aims to enhance the accessibility of our work for readers less familiar with State Space Models.

**Additional Experiments.** To address concerns raised by Reviewers QtDa and JuBx and to strengthen our empirical evidence, we conducted additional experiments. Specifically, 

1. **Results on the** **Fineweb Dataset.** We present results on a randomly sampled subset of the Fineweb dataset (Figure 3 in the attachment). While computational and time constraints prevented us from training on the entire dataset or using larger model widths, our observations on Fineweb are consistent with the wikitext-103 results presented in the main paper (Figure 2).

2. **Decoupled Learning Rates.** To isolate the effects of SSM scaling, we decoupled learning rates for SSM and non-SSM layers. We first tuned the learning rate of non-SSM layers, then compared test performance across different scales for various SSM learning rates, using the optimal non-SSM learning rate (Figure 1 in attachment).

3. **Verify Correct Scaling**. In a simplified scenario where the dimension of the latent states $N_x$ is fixed, heuristic muP spectral scaling yields $\eta_B,\eta_C=\Theta(\frac{1}{N_u})$ while muP-SSM gives $\eta_B,\eta_C=\Theta(\frac{1}{\sqrt{N_u}})$ (see Table 1 in the main paper). In experiments carried out independently by researchers (to be acknowledged in the public version), their results as shown in Figure 4 in the attachment verify that the correct width-independent scaling indeed aligns with muP-SSM.

These new experiments further validate our theoretical results and address the specific concerns raised in the reviews. We plan to incorporate these results into the revised manuscript. Below, we summarize our key findings based on both the main paper experiments and the additional experiments included in the pdf.

**Summary of empirical findings.** Our work evaluates the scaling behaviour of different parameterizations for State Space Models (SSMs) using four criteria: generalization (test loss), training stability, monotonic improvement of generalization with scale, and transferability of optimal learning rate across scales. Our experiments reveal that while **Standard Parameterization (SP) consistently performs worse than muP-SSM across all metrics**, the comparison between muP-SSM and spectral scaling yields nuanced results at the scales we test:

1. **Generalization.** Spectral scaling has only slightly worse test loss compared to muP-SSM at the largest scale we were able to test across all the experiments.
2. **Training stability.** However, muP-SSM demonstrates markedly improved training stability at larger learning rates compared to the spectral scaling approach already at the relatively small scales we test. Note that, instabilities in small scale models at large learning rates can be predictive of general training instabilities at larger scales [2].
3. **Monotonic improvement of generalization with scale.**
    - Using a single global learning rate for the entire network (SSM + non-SSM layers), similar to muP-SSM, spectral scaling appears to improve performance monotonically with scale up to the largest scale tested (Figure 2 in main paper, Figure 3 in rebuttal attachment).
    - However, we hypothesized this might be an artifact due to the correct scaling of non-SSM modules under spectral scaling. To test this, we decoupled learning rates for SSM and non-SSM layers. That is, we first tune the learning rate of the non-SSM layers and then under the optimal choice of LR for the non-SSM layers, we compare performance of the models across different scales for different SSM learning rates. Figure 1 in the rebuttal attachment shows that beyond a certain width threshold, performance under spectral scaling no longer improves monotonically with scale, contrasting sharply with muP-SSM.
4. **HP transferability.** At tested scales, both muP-SSM and spectral scaling demonstrate similar transferability. This is not entirely unexpected when using a global learning rate, as both methods identically parameterize non-SSM layers.

It's important to note that unlike muP-SSM, the spectral scaling approach used in our experiments is heuristically derived and lacks rigorous theoretical justification.

Meta-review: This paper studies feature learning in SSMs in the infinite-width regime. There are several contributions: (1) showing muP parameterization leads to unbounded outputs (2) identifying this problem is due to tensor programs not being able to represent SSMs (3) developing a correction that leads to more stable training and bounded outputs (4) proposing a scaling rule for SSMs and empirically tests it on several architectures.  

The reviewers appreciate the interesting problem and the theoretical analysis. Even though the empirical validation could be more comprehensive, given the initial promising results and the potential impact on the SSM space, I recommend acceptance.
Decision: Accept (poster)