Title: On Feature Learning in Structured State Space Models
Venue: NeurIPS 2024
Accepted: yes
Keywords: State space models

feature learning

hyperparameter transfer

scaling theory

deep learning
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

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.

1 Introduction

State-space models (SSMs), such as Mamba (Gu and Dao, 2023), have become popular in deep learning as alternatives to transformers like GPT and BERT series (Radford et al., 2019, Brown et al., 2020, Achiam et al., 2023, Devlin et al., 2018, Touvron et al., 2023, Chowdhery et al., 2023, Gemini Team et al., 2023). SSMs integrate elements from RNNs, CNNs, and control models, excelling in 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 scale-both in terms of data and model size. However, this increased scale often introduces challenges, such as precision issues due to instability or the vanishing/exploding gradient problems. Additionally, the sequential nature of state-space models (SSMs) makes them notoriously difficult to train. Therefore, developing a rigorous understanding of how SSMs scale as their dimensions increase and identifying optimal scaling rules is crucial.

In this vein, infinite-width asymptotics, such as Neural tangent kernel (NTK) analyses have been a central tool providing key insights in DL theory. However, a 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 hyperparameters, which includes layer-wise scaling of the learning rates, one can find a unique parameterization called Maximal Update pPrameterization ( µ P) that induces non-trivial feature evolution

SP

(Heuristic)

(Ours)

SSM Latent States, ZOH

SSM Output States, ZOH

SSM Latent States, Euler

SSM Output States, Euler

-SSM

log2(

√N

log2(

√N

log2(

√Nu

log2(

√N

log2

√N

log2(

√N

log2(

√Nu

log2(

√N

15

17

19

16

18

20

12

14

11

13

15

12

9

8

Layer Index

三

2#

1

2

√N

23

219

211212

16

22

74

14

14

18

20

13

16

17

10-

11

12

13

15

-14

2#

√N

2

2#

211212

20.0

1 10.0

1 15.0

17.5

1 22.5

12.5

7.5

10

10

11

12

14

-11

12

13

-14

13

9

9

sZ

Layer Index

23

1

3

2

√Nu

2

2

10

15

20

25

10

10

11

12

13

14

11

12

13

14

2

2

√Nu

2#

2

I

I

I

I

I

I

I

I

1

1

10

10

14

16

11

13

11

12

13

10

12

12

8

6

9

6-

2”

Layer Index

三

2#

2

1

E

√N

2

221212

10.0

12.5

15.0

17.5

5.0

7.5

n n

11

10

11

10

12

13

14

12

13

14

2

2#

√N

221211212

10

15

20

14

11

12

-10

10

13

14

16

15

ZT-

-5

-9

2

Layer Index

三

2²

1

2

√Nu

2#

2²2²2

20

10

15

25

10

12

10

11

14

12

13

15

16

=8

5

2

4

6

9

2

2²

√Nu

2

2”

even in the infinite width limit for a 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 to 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 µ P 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 as Tensor Programs ( cf. , Figure 1).

We derive a 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 to large-scale SSMs, similar to 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. A sequence layer can be represented as y 1: L = f seq ( u 1: L ) where u 1: L is a compact notation for sequence u 1 , . . . , u L , with u l ∈ R N u , y l ∈ R N y . Note that f seq can be easily generalized to accept and output multiple sequences as in residual connections and multi-branch architectures (see Appendix B.1). Here N u and N y are treated as widths of the sequence model, which can be scaled up to create larger sequence models. To denote the backward pass of a model concisely, let L be the overall training loss, and we write ¯ · = ∂ L ∂ · . Instead of studying the scaling of each element in a vector u ∈ R N u , we study the norm ∥ u ∥ 2 . When elements in u are in Θ(1) , we have ∥ u ∥ 2 ∈ Θ( √ N u ) . Throughout this work, we will use Θ to represent the asymptotic order in the limit ( N u →∞ in this case), but omit the precise notation for readability.

In layerwise sequence models, where the k -th layer is denoted as u ( k ) 1: n = f k ( u ( k -1) 1: n ) with u ( k ) l ∈ R N k , 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 ∈ [ K ] , the features u ( k ) 1: L and the updates of features ∆ u ( k ) 1: L after one gradient update have the following scaling:

Note that the first condition is a stability condition that demands the activations to remain coordinatewise Θ(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.1.

For a 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 ( ∆ ), we say this sequence layer admits feature learning. The corrected scaling µP -SSM we propose in Section 3 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 is 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 MLPs, but requires different initialization and learning rate scaling rules for the trainable weights in SSMs, as 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 NE ⊗ OR ⊤ program (Yang and Littwin, 2023). A NE ⊗ OR ⊤ program constitutes a sequence of vectors in R n and a sequence of scalars in R inductively generated from an initial set of random scalars, vectors, and matrices following a 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 as a NE ⊗ OR ⊤ program (which includes many modern architectures used in practice including transformers and convolutional networks), there exists a unique scaling rule called Maximal Update Parametrization ( µ P) under which features in every single layer evolve in a width-independent fashion with scale. Interestingly, it has been shown that unlike standard parameterizations, µ P allows transferability of optimal hyper-parameters from small to large scale models.

The key idea of the 'Tensor Program Ansatz' (Yang, 2019) is that pre-activations and their gradients arising in training most neural network architectures via standard update rules (such as SGD and ADAM) can be represented as vectors in a NE ⊗ OR ⊤ program. The Ansatz suggests that vectors in a TP become asymptotically independent as well as identically distributed both at initialization as well as during training. Leveraging this, the framework allows for a mechanistic tracking of the behaviour of such vectors by assigning a random variable Z v to represent the (asymptotically identical) distribution of the coordinates of the vector v . This principle underlies the development of the key theoretical 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 a 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 1: L = f S4 ( u 1: L ; w = { B , C } ) can be viewed as a discretization of a continuous-time SSM given by

for l = 1 , . . . , L with x 0 = 0 and where Re ( · ) gives the real part of a complex vector, and where we let A = Diag( a ) with a ∈ C N x , B ∈ C N x × N u , C ∈ C N y × N x . We have the discretized sequence to sequence mapping f S4 as follows:

where we follow the Zero-Order-Hold (ZOH) discretization method,

In practice, simple Euler discretization can also be used for B , i.e., B ′ = τ B . 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 NE ⊗ OR ⊤ 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 NE ⊗ OR ⊤ 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 a NE ⊗ OR ⊤ program. Indeed this holds for all existing Hippo-parameterizations of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegS-D, S4D-Inv, S4D-Lin, and S4D-Real.

Proof sketch. The formal proofs are provided in Appendix C. The key architectural component that is not representable in a NE ⊗ OR ⊤ program is the Hippo-based structured transition matrix A . Due to the rapid decay of the diagonal entries of A -1 , the hidden states x 0 are not even asymptotically identically distributed, already at initialization. Therefore coordinates of x 0 cannot be generated by a sequence of NE ⊗ OR ⊤ computations, neither at 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 generalization 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 a 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 l and ∆ W l ∈ R N l × N l -1 denote the l th layer weight matrices and their gradient updates. The condition states that for feature learning to hold, the spectral norms of the matrices should satisfy:

Under µ P, architectures that are representable as a NE ⊗ OR ⊤ 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 ( ∆ ) . Indeed this holds for all well-known Hippo-parameterizations of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegS-D, S4D-Inv, S4D-Lin, and S4D-Real.

Proof sketch. The formal proofs are provided in Appendix C. To gather some intuition for this result in a simplified setting, consider the scale of the hidden states for token index 0 , which is already wrong at initialization. By definition, ∥ x 1 ∥ 2 = ∥ B ′ u 1 ∥ 2 = ∥ ΛBu 1 ∥ 2 , where Λ = ( A ′ -I ) A -1 . Now, observe that ∥ B ′ u 1 ∥ 2 2 is a sum of independent random variables that satisfy the Kolmogorov condition, so that the sum behaves according to 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:

However, using standard tools from random matrix theory one can show that ∥ B ′ u ∥ 2 ( ∥ B ′ ∥ ∗ ∥ u 1 ∥ 2 ) ∈ Θ( 1 √ N u ) which clearly violates (4).

3 Identifying the unique scaling for effective feature learning in SSMs

In this section, we 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 as S4, S5, H3, or DSS. In the appendix, we also provide a 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 in the Mamba architecture y 1: L = f Mamba ( u 1: L ) can be written as follows:

For i = 1 , . . . , N u , we have N u 1D SSMs:

where x ( i ) 0 = 0 , a ( i ) log , τ 0 ∈ R N u , and all linear layer weights W B , b B , W C , b C , W τ , b τ are trainable parameters. The Zero-Order-Hold (ZOH) discretization procedure A ′ ( i ) l , B ′ ( i ) l = ZOH( τ, A ( i ) , B ′ ( i ) l ) can be written as:

In a nutshell, Mamba can be seen as N u 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:

and zero initialization for all biases. For each SSM, A ( i ) is still initialized according to the HiPPO theory. When A ( i ) is real-valued, we let a ( i ) log [ j ] = log( j + 1) and τ 0 can be seen as a bias term initialized to τ 0 ∼ Softplus -1 ( U (0 . 001 , 0 . 1)) . We will assume that τ 0 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 τ as a trainable parameter.

Per Channel Linear

Recurrence

......

......

Selection

Discretization

Per Channel Parameter Generation

3.2 Forward signal propagation through a S6 Mamba layer

To derive the correct choice of initialization scalings σ B and σ C , 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 u →∞ then N x →∞ . We believe that our results would also hold in the proportional limit where N u →∞ and N x →∞ with N x N u ∈ Θ(1) . However deriving the results in this setting would incur a significant technical overhead and we defer this analysis to a future work. We briefly discuss this in Section 6. 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, or 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 µ P.

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 B 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 are summarized in Table 1.

For simplicity of exposition, first consider the scale of x ( i ) 1 = ( A ′ ( i ) l -I ) A ( i ) -1 B 1 u ( i ) 1 before generalizing to arbitrary l ∈ [ L ] .

Proposition 3.2 ( Scale of hidden states x ( i ) 1 in Mamba at initialization ) . Under the ZOH discretization procedure, as N u then N x approach infinity, for any i ∈ [ N u ] , the squared l 2 -norm of the hidden states ∥ x ( i ) 1 ∥ 2 2 is a.s. scaled as ∥ x ( i ) 1 ∥ 2 2 ∈ Θ ( ζ (2) σ 2 B ∥ u 1 ∥ 2 ( u ( i ) 1 ) 2 ) , where ζ (2) denotes the Riemann zeta function at 2.

Following condition (INIT), the 1D SSM admits stability at initialization under the following conditions:

Therefore, for stability at initialization, the initialization should scale as σ B ∈ Θ( √ N x N u ) . Note that, under both standard parameterization (e.g., Kaiming or LeCun initialization) and spectral scaling conditions, σ B is initialized as Θ( √ 1 N u ) , which leads to vanishing hidden states according to Proposition 3.2. We empirically verify this fact in Figure 1.

Proposition 3.3 provides the scale of the output of a Mamba layer.

Proposition 3.3 ( Scale of outputs y ( i ) 1 of a Mamba layer at initialization ) . Under the ZOH discretization procedure as N u then N x approach infinity, for any i ∈ [ N u ] , the output y ( i ) 1 converges in distribution to a Gaussian with mean 0 and standard deviation Cσ B σ C ∥ u 1 ∥ 2 2 for some widthindependent constant C > 0 .

Accordingly, imposing the conditions for stability of y ( i ) 1 according to condition (INIT) implies the initialization scaling condition σ C ∈ Θ( √ 1 N x N u ) . Note again that standard parameterization suggests initializing σ C ∈ Θ( √ 1 N u ) under which the outputs would diverge with scale. Under spectral scaling, σ C 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 to arbitrary l ∈ [ L ] . The corrected scalings of σ B and σ C derived above generalize to the entire sequence, as a sum over the sequence usually does not cancel out the scaling. More formally, for all l ∈ [ L ] , we have

First, observe that the operator ( A ′ ( i ) l ) m does not change the width-scaling. To see this note that since A ′ ( i ) l = Diag( a ′ 1 , . . . , a ′ N x ) with a ′ n = e -1 2 τ ( i ) l (cos( τ ( i ) l πn ) + i sin( τ ( i ) l πn )) , we have that, for all complex vectors v ∈ C N x it holds that ∥ ( A ′ ( i ) l ) m v ∥ 2 = e -mτ ( i ) l / 2 ∥ v ∥ 2 for any m ∈ [ L ] . Now since, for any u l , setting σ B ∈ Θ( √ N x N u ) yields ∥ ∥ ∥ B ′ ( i ) l u l ∥ ∥ ∥ 2 ∈ Θ( √ N x ) , each term in the summation is of order Θ( √ N x ) . Unless, for every l , the term B ′ ( i ) l u l perfectly cancels out with the terms before to affect the width scaling , we have that ∥ ∥ ∥ x ( i ) l ∥ ∥ ∥ 2 ∈ Θ( √ N x ) . The same argument can be used to show the stability of y ( i ) l . Concluding this argument, we have derived the correct scaling of the initialization variances σ B and σ C 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 y 1: L = f mamba ( u 1: L ; w ) are stable at initialization in the infinite-width limit under the following scaling conditions:

.

3.3 Backward signal propagation in a S6 Mamba Layer

In this section, we provide the correct scaling of the learning rates η a , η B , and η C by a detailed analysis of the backward signal propagation in the limit of N u →∞ then N x →∞ . Specifically, we analyze the scale of the activation updates for both hidden states and outputs in the first backward pass through a 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 1.

Proposition 3.4 ( Scale of the updates of hidden states ∆ x ( i ) 1 after 1 step of SGD ) . Under the ZOH discretization procedure, as N x then N u approach infinity, for every 1D SSM, the squared l 2 -norm of the updates ∥ ∆ x ( i ) 1 ∥ 2 of the hidden states after one step of SGD is a.s. scaled as ∥ ∆ x ( i ) 1 ∥ 2 ∈ Θ ( η B 1 √ N u σ C ∥ u 1 ∥ 3 2 ζ (4) 1 2 ) , where ζ (4) denotes the Riemann zeta function at 4.

The Riemann zeta function at 4 evaluates to a width-independent constant and σ C ∈ Θ( √ 1 N x N u ) due to Proposition 3.3 for stability of outputs. Therefore, for the scale of the hidden state updates to be Θ( √ N x ) , the correct scaling of the learning rate η B is given by Θ( N x √ N u ) . On the other hand, spectral scaling suggests that η B must scale as Θ( N x N u ) . Under this scaling, however, updates of the 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 1.

Next, to derive the correct scaling of the learning rate η C , we consider the scale of output updates.

Proposition 3.5 ( Scale of the updates of outputs ∆ y ( i ) 1 after 1 step of SGD ) . Under the ZOH discretization procedure, as N u then N x approach infinity, for every 1D SSM, the squared l 2 -norm of the updates of the hidden states after one step of SGD scales as | ∆ y ( i ) 1 | 2 ∈ Θ ( η C σ 2 B √ 1 N u ∥ u 1 ∥ 4 2 ) .

The result suggests that for the correct scaling of the updates, the learning rate η C must scale as Θ( 1 N x √ N u ) . Under spectral scaling, η C scales much larger as Θ( N x N u ) . However, since the updates of the hidden states vanish under spectral scaling, this larger incorrect scaling of η C downward corrects the scale of the updates in the outputs as shown in Table 1 and empirically verified in Figure 1.

On the correct scaling of η a . It turns out that the scaling of the learning rate η a does not play a role in either stability at initialization or for non-trivial updates with scale. However, if η a is not scaled correctly, then the transition matrix A is not updated. In particular, as shown in Appendix C.2, η a needs to scale as Θ( N u ) . Below we summarize the correct scaling conditions to achieve non-trivial feature updates in the infinite-width limit of a Mamba layer.

In the same vein as the discussion in Section 3.2, it is straightforward to verify that our results hold for arbitrary l ∈ [ L ] 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 in a S6 Mamba recurrent layer y 1: L = f mamba ( u 1: L ; w ) evolve non-trivially in the infinite-width limit under the following conditions:

.

4 µ P-SSM implies stability and feature learning in Mamba

Empirical verification of different scalings in Table 1. First, we verify that our derived µ P-SSM scaling (see Table 1) for Mamba indeed leads to feature learning, i.e., it ensures stability at initialization as defined in condition (INIT) and non-trivial feature updates during training as defined in condition ( ∆ ). We scale up the SSM latent state size N x and the SSM output dimension N u simultaneously and track the scaling of both features and feature updates. Due to the linear decay in the eigenvalues of the transition matrix A -1 , we typically observe a strong finite sample effect at small N x . Constrained by computational resources, we opt for a much smaller N u ( N u = N x / 8 ) than is usually employed in practice. This adjustment enables us to scale up N x 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 in this section, we train Mamba with 3 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 µ P 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 l , 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 Euler discretization is used for B l , the latent states will have the correct scaling, but the output signals will explode. The results clearly highlight the importance of

correcting the scaling of the µ P parameterization. When using the corrected scaling which we call µ P-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 in scaling when the width is small is due to finite sample effects. It stabilizes once the width is sufficiently large.

Stability, generalization, and hyper-parameter transfer. In Figure 3, we employ Mamba as a generative model on the wikitext-103 dataset and conduct single-epoch training for 20 K iterations. We plot the test loss against the learning rate on a logarithmic scale and compare the results across different model widths (both N u and N x ). In this experiment, we use the standard setting where N u ≫ N x ( N u = 16 N x in this case). Using µ P-SSM scaling or µ P (heuristic) significantly improves test performance compared to standard parameterization of Mamba for this task. For larger learning rate, µ P-SSM shows better stability compared to µ P (heuristic), highlighting the importance of deriving the correct scaling for SSMs rather than heuristically adopting µ P (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 hyper-parameters 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 depth-dependent scaling of the residual branches (Bordelon et 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 a more thorough understanding of both necessary and sufficient conditions of the transferrability of hyper-parameters is warranted.

SP

(Heuristic)

Test loss

-SSM (Ours)

4.5

256

512

4.0

1024

2048

4096

3.5

3.0

2.5

6

-4

2

0

-6

-5

-4

-3

2

-1

-6

-5

-4

E-

-2

-1

log2(lr)

log2(lr)

log2(lr)

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 mean-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 architectures, optimization algorithms and arbitrary abc -parameterizations. Each abc -parameterization is essentially defined by a layerwise scaling of initialization variance and learning rate as a function of

network width. Seminal work by Yang and Hu (2021) shows that there exists a unique maximal update parameterization ( µ P) that attains a stable feature learning infinite-width 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, µP , 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 hyper-parameters 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 to the N u then N x tend to infinity setting. A completely rigorous analysis of SSMs in the proportional limit where N u and N x approach infinity with N x N u held roughly constant requires us to carefully track how the different 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 is that the different vectors in a TP (such as activations or updates) are asymptotically identically distributed. As discussed earlier, this assumption crucially underlies the key theoretical results of TP called Master Theorems. However, a 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, it 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.

References

Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 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 arXiv:2405.04517 , 2024.

Blake Bordelon, Lorenzo Noci, Mufan Bill Li, Boris Hanin, and Cengiz Pehlevan. Depthwise hyperparameter transfer in residual networks: Dynamics and scaling limit. arXiv:2309.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 L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Albert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mixing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv: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: a family of highly capable multimodal models. arXiv preprint arXiv:2312.11805 , 2023.

Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv: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:3597135983, 2022.

Ankit Gupta, Albert Gu, and Jonathan Berant. Diagonal state spaces are as 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. PMLR, 2021.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural Tangent Kernel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems (NeurIPS) , pages 8571-8580, 2018.

Mufan Li, Mihai Nica, and Dan Roy. The future is log-gaussian: Resnets and their infinite-depth-andwidth 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 twolayer neural networks. Proceedings of the National Academy of Sciences , 115(33):E7665-E7671, 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 in transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems , 35:27198-27211, 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-width 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 Sohl-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 are unsupervised multitask learners. 2019. URL https://api.semanticscholar.org/ CorpusID:160025533 .

Samuel S Schoenholz, Justin Gilmer, Surya Ganguli, and Jascha Sohl-Dickstein. Deep information propagation. arXiv:1611.01232 , 2016.

Pranab K Sen and Julio M Singer. Large sample methods in statistics: an introduction with applications , 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, et 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, Łukasz Kaiser, and Illia Polosukhin. Attention is 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 Funtowicz, et al. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019.

Lechao Xiao, Jeffrey Pennington, and Samuel Schoenholz. Disentangling trainability and generalization in deep neural networks. In International Conference on Machine Learning , pages 10462-10472. PMLR, 2020.

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 J. Hu. Tensor programs iv: Feature learning in infinite-width neural networks. In International Conference on Machine Learning (ICML) , 2021.

Greg Yang and Etai Littwin. Tensor programs ivb: Adaptive optimization in the infinite-width limit. arXiv:2308.01814 , 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 v: 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 arXiv:2310.02244 , 2023b.

A Background

A.1 NE ⊗ OR ⊤ , abc-parameterizations and µ P

NE ⊗ OR ⊤ programs. For the precise definitions and theorems of the NE ⊗ OR ⊤ framework we refer to Yang and Littwin (2023). Here we try to provide an intuitive introduction.

A NE ⊗ OR ⊤ program consists of a set of vectors inductively generated from a set of initial matrices, vectors, and scalars following a set of allowed instructions: vector averages (Avg), matrix multiplications (MatMul), and non-linear outer products (OuterNonlin).

All entries of an initial vector v are sampled iid from N (0 , 1) . Every entry of an initial matrix A is sampled independently from distributions that all have mean 0 , variance n -1 and that satisfy a technical boundedness assumption on all higher-order moments. Nonlinearities ψ used for OuterNonlin operations are either pseudo-Lipschitz 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 = 1 n ∑ n α v α . Now the NE ⊗ OR ⊤ Master Theorem implies c → ˚ c := E Z v almost surely. Matrix multiplications Av are defined as usual. In the limit, a non-trivial interaction term arises if A ⊤ had appeared in a NE ⊗ OR ⊤ operation that has generated the vector v . To define the OuterNonlin operation, aggregate a fixed amount | x | of NE ⊗ OR ⊤ vectors in x and fix r ∈ N . Denote the collection of all l defined NE ⊗ OR ⊤ scalars by c . Then a OuterNonlin operation with nonlinearity ψ : R | x | ( r +1)+ l → R , is given by

and, due to the NE ⊗ OR ⊤ Master Theorem, behaves in the limit as

where Z x is an r -dimensional random vector with the limiting distribution of x , and Z x 1 , . . . , Z x r are iid copies of Z x .

The Maximal Update Parameterization ( µ P). For any architecture/component whose forward pass is representable as a NE ⊗ OR ⊤ program, µ P can be derived following the instructions below. Let each parameter tensor W be parameterized as W = n -a W w where w is a trainable parameter with initialization w αβ ∼ N (0 , n -2 b W ) . Let the learning rate be parameterized as ηn -c for some width-independent η > 0 . Then µ P is prescribed as follows: Set c = 0 and b W = 1 / 2 for all parameter tensors W .

a W = 0 if the input and output to W are width-independent (e.g., scalars),

a W = 2 if both input and output to W scale with width (e.g., attention matrices).

a W = 1 if input is width-independent and output scales with width (e.g., input embeddings).

a W = 1 / 2 if input scales with width and output is width-independent (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-Feller Central Limit Theorem provided in Sen and Singer (1994, Theorems 3.2.10 and 3.3.3).

Theorem A.1 (Kolmogorov Strong Law of Large Numbers) . Let X i , i ≥ 1 , be independent random variables such that E X i = µ i and V ar ( X i ) = σ 2 i exist for every i ≥ 1 . Then

Theorem A.2 (Lindenberg-Feller-Central Limit Theorem) . Let X i , i ≥ 1 , be independent random variables such that E X i = µ i and V ar ( X i ) = σ 2 i exist for every i ≥ 1 . Also let s 2 n = ∑ n i =1 σ 2 i and Z n = s -1 n ( ∑ n i =1 X i -∑ n i =1 µ i ) . Then the Lindenberg-Feller condition

holds if and only if both

(ii) Z n converges in distribution to a Gaussian with mean 0 and standard deviation 1 .

B Definitions

Definition B.1 (Softplus) . The softplus function is defined as Softplus ( x ) = 1 β log(1 + exp( βx )) with default smoothing value β = 1 .

B.1 Feature Learning in General Sequence Models

Let G be a directed acyclic graph (DAG) with vertices V ( G ) and edges E ( G ) . For each node v ∈ V ( G ) , its parent nodes are denoted by PA ( v ) . In a sequence model, each node v corresponds to a sequence v 1: L with length L , computed from its parent node sequences with operation v 1: L = f v ( { u 1: L | u ∈ PA ( v ) } ) . These sequence layers can be shared instance-wise multi-layer perceptrons (MLPs), normalizations, residual summation, or transformers/recurrent layers, etc.

Definition B.2 (Feature Learning in Sequence Models) . A sequence model is in the feature learning regime if for any v ∈ V ( G ) , the features v 1: L where v ∈ R N v and the updates of features ∆ v 1: L after one gradient update have the following scaling:

C Proofs

C.1 Structured SSMs are not covered by previous approaches

Proposition C.1 ( Structured SSMs are not generally representable as NE ⊗ OR ⊤ programs ) . The forward and backward signal propagation of structured SSMs including S4 and MAMBA trained via standard algorithms such as SGD or ADAM are not representable as a NE ⊗ OR ⊤ program. Indeed this holds for all existing Hippo-parameterizations of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegS-D, S4D-Inv, S4D-Lin, and S4D-Real.

Proof. Recall the definition of a NE ⊗ OR ⊤ program (Yang and Littwin, 2023). Let A be any of the HiPPO or S4D structured matrices. As the entries of A at initialization are deterministic and differ, A can neither be an initial NE ⊗ OR ⊤ matrix, nor be generated by an OuterNonlin of random NE ⊗ OR ⊤ vectors (acting on all coordinate dimensions in the same way). This is because no allowed NE ⊗ OR ⊤ operation can generate a NE ⊗ OR ⊤ vector with differing variance in its coordinates, whereas the multiplication with A clearly produces differing variances in 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 NE ⊗ OR ⊤ program have the same variance. The allowed operations to generate a NE ⊗ OR ⊤ vector are matrix multiplication with an initial matrix and OuterNonlin. As initial matrices A 0 have independent entries with mean 0 and variance n -1 , the multiplication A 0 v will again generate a vector with the same variance in each coordinate.

An OuterNonlin operation takes in previously defined NE ⊗ OR ⊤ vectors and treats all coordinates in the same way. Consequently it will again generate a NE ⊗ OR ⊤ vector with the same variance in each coordinate.

There is no other way to generate a new NE ⊗ OR ⊤ vector, which concludes the induction.

To see that the structured matrix A produces nonisotropic coordinate distributions, we detail the argument for an S4 recurrent layer. The other choices of A follow via analogous arguments. Recall

that the S4 recurrent layer y 1: L = f S4 ( u 1: L ; w = { B , C } ) can be viewed as a discretization of a continuous-time SSM given by

for l = 1 , . . . , L with x 0 = 0 and where Re ( · ) gives the real part of a complex vector, and where we let A = diag ( a ) with a ∈ C N x , B ∈ C N x × N u , C ∈ C N y × N x . We have the discretized sequence to sequence mapping f S4 as follows:

where we follow the Zero-Order-Hold (ZOH) discretization method,

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 n = -1 2 + iπn . B and C are initialized such that the entries are i.i.d and follow a Gaussian distribution with mean 0 and variance σ 2 B and σ 2 C 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 us consider the distribution of the hidden states for token index 0 at initialization: x 0 = B ′ u 0 , where B ′ = ( A ′ -I ) A -1 B , and A ′ = diag ( a ′ 1 , . . . , a ′ N x ) with a ′ n = e -1 2 τ (cos( τπn ) + i sin( τπn )) , then the eigenvalues of A ′ are clearly Θ(1) . To further simplify, if we set τ = 2 , then ∥ A ′ ∥ ∗ = e -1 and for all complex vectors v ∈ C N x it holds that A ′ v = -e -1 v . Since B is an i.i.d Gaussian matrix, for sufficiently large width, Bu 0 is distributed i.i.d by a simple Central Limit Theorem argument. However, B ′ also depends on A -1 = diag ( a -1 n ) with a -1 n = -1+2 πni 1 / 2+2 π 2 n 2 . Due to the linear decay of the diagonal entries of A -1 , the hidden states x 0 already at initialization are not (even asymptotically) identically distributed. Therefore, there cannot exist a 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 algorithms such as SGD or ADAM that satisfy spectral scaling conditions ( ∗ ) do not satisfy conditions for feature learning given in Equation ( ∆ ) . Indeed this holds for all well-known Hippo-parameterizations of the structured matrix A including HiPPO-LegS, HiPPO-LegS-N, HiPPO-LegS-D, S4D-Inv, S4D-Lin, and S4D-Real.

Proof. Let A be any of the HiPPO or S4D structured matrices.

Consider the scale of the hidden states at initialization for token index 0 . By definition, ∥ x 1 ∥ 2 = ∥ B ′ u 1 ∥ 2 = ∥ ΛBu 1 ∥ 2 , where Λ = ( A ′ -I ) A -1 . Then, observe that ∥ B ′ u 1 ∥ 2 2 = N x ∑ i =1 Λ 2 i ( N u ∑ j =1 B i,j u 1 ,j ) 2 is a sum of N x 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 Λ and the correct scaling of the inputs to the S4 layer. Applying the Kolmogorov SLLN (Theorem A.1), we obtain

where ζ (2) denotes the Riemann zeta function at 2 and equates to π 2 / 6 , and c > 0 is some widthindependent constant. The zeta function appears due to c ζ (2) ≤ ∑ N x i =1 Λ 2 i ≤ ∑ N x i =1 1 ( i +1) 2 ≤ ζ (2) for N x large enough (see (37) for more details). Here and throughout the paper, for the formal application of the LLN or the CLT, we first adequately normalize all quantities to arrive at a welldefined width-independent limit statements, but do not write out such technicalities for conciseness.

Observe that, for spectral scaling conditions to yield the right scaling of the initialization variance, it is crucial that the following condition holds:

Since the spectrum of ( A ′ -I ) A -1 is less than 1 , an upper bound on the spectral norm of B ′ = ( A ′ -I ) A -1 B can be found in Vershynin (2011) and is given by c ( √ N x + √ N u ) 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: ∥ B ′ ∥ ∗ ≥

max { max i ∥ B ′ i : ∥ 2 , max j ∥ ∥ B ′ : j ∥ ∥ 2 } ≥ √ N u σ B . Therefore ∥ B ′ u ∥ 2 ( ∥ B ′ ∥ ∗ ∥ u 1 ∥ 2 ) ∈ Θ( 1 √ N u ) which clearly violates (14). This spectral decay is induced by all of the considered choices of A .

C.2 S6 Mamba recurrent layer

Recall that a S6 Mamba layer y 1: L = f Mamba ( u 1: L ) can be written as follows:

For i = 1 , . . . , N u , we have N u 1D SSMs:

where x ( i ) 0 = 0 , a ( i ) log , τ 0 ∈ R N u , and all linear layer weights W B , b B , W C , b C , W τ , b τ are trainable parameters. The Zero-Order-Hold (ZOH) discretization procedure A ′ ( i ) l , B ′ ( i ) l = ZOH( τ, A ( i ) , B ′ ( i ) l ) can be written as:

In a nutshell, Mamba can be seen as N u 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:

and zero initialization for all biases. For each SSM, A ( i ) is still initialized according to the HiPPO theory. When A ( i ) is real-valued, we let a ( i ) log [ j ] = log( j +1) and τ 0 can be seen a bias term initialized to τ 0 ∼ Softplus -1 ( U (0 . 001 , 0 . 1)) .

Below we restate Propositions 3.2 and 3.3 followed by their proofs.

Proposition 3.2 ( Scale of hidden states x ( i ) 1 in Mamba at initialization ) . Under the ZOH discretization procedure, as N u then N x approach infinity, for any i ∈ [ N u ] , the squared l 2 -norm of the hidden states ∥ x ( i ) 1 ∥ 2 2 is a.s. scaled as ∥ x ( i ) 1 ∥ 2 2 ∈ Θ ( ζ (2) σ 2 B ∥ u 1 ∥ 2 ( u ( i ) 1 ) 2 ) , where ζ (2) denotes the Riemann zeta function at 2.

Proposition 3.3 ( Scale of outputs y ( i ) 1 of a Mamba layer at initialization ) . Under the ZOH discretization procedure as N u then N x approach infinity, for any i ∈ [ N u ] , the output y ( i ) 1 converges in distribution to a Gaussian with mean 0 and standard deviation Cσ B σ C ∥ u 1 ∥ 2 2 for some widthindependent constant C > 0 .

Proofs for Proposition 3.2 and Proposition 3.3 . For the following 1D linear state space model:

Following condition (INIT), the 1D SSM admits stability at initialization, when the following condition is satisfied:

For simplicity of exposition, let's begin by considering the scale of x ( i ) 1 = ( A ′ ( i ) l -I )( A ( i ) ) -1 B ( i ) 1 u ( i ) l before generalizing to arbitrary l ∈ [ L ] .

Since B 1 = W B u 1 + b B with b B = 0 at initialization,

Note that the inner summation can be expressed as

̸

Since W B has i.i.d Gaussian entries and due to Assumption 3.1, as N u →∞ , the first term behaves according to law of large numbers and converges almost surely to σ 2 B ∥ u 1 ∥ 2 . The second term behaves according to central limit theorem and converges in distribution to a Gaussian with mean 0

̸

̸

a Gaussian distribution with mean σ 2 B ∥ u 1 ∥ 2 and variance σ 4 B N u ∑ n ′ = n ′′ =1 u ( n ′ ) 1 u ( n ′′ ) 1 .

For applying Kolmogorov SLLN (Theorem A.1), note that c · ζ (2) ≤ ∑ N x m =1 (Λ i ) 2 m,m ≤ ∑ N x m =1 1 ( i +1) 2 ≤ ζ (2) and c · ζ (4) ≤ ∑ N x m =1 (Λ i ) 4 m,m ≤ ∑ N x m =1 1 ( i +1) 4 ≤ ζ (4) for some constant c ∈ (0 , 1) and for N x large enough (see (37) for more details). For random variables v m ∼ N ( σ 2 B , C 2 ) for some C 2 > 0 , it holds that V ar ((Λ i ) 2 m,m v m ) ≤ E [(Λ i ) 4 m,m v 2 m ] ≤ (Λ i ) 4 m,m ( σ 4 B + C 2 ) , so that the Kolmogorov condition is fulfilled, and we have that as N u then N x approach infinity, for some width-independent constant c > 0 ,

Therefore, for stability at initialization, the scale of initialization σ B ∈ Θ( √ N x N u ) .

Scale of y ( i ) l . When l = 1 , y ( i ) 1 = C 1 x ( i ) 1 , where x ( i ) 1 = B ′ ( i ) 1 u ( i ) 1 and we have

Applying the Lindenberg-Feller CLT (Theorem A.2), we have that, as N then N approach infinity, u x u ( i ) 1 N x ∑ m =1 (Λ i ) m,m N u ∑ j,k =1 ( W C ) m,j ( W B ) m,k u ( j ) 1 u ( k ) 1 converges to a Gaussian distribution with mean

0 and variance c ζ (2) u ( i ) 1 σ B σ C ∥ u 1 ∥ 2 2 .

Accordingly, imposing the conditions for stability of y ( i ) 1 requires that σ C ∈ Θ( √ 1 N x N u ) .

Stability of x ( i ) l and y ( i ) l for arbitrary l ∈ [ L ] . For all l ∈ [ L ] , we have x ( i ) l = ∑ l -1 m =0 ( A ′ ( i ) l ) m B ′ ( i ) l -m u ( i ) l -m . First, observe that since A ′ = diag ( a ′ 1 , . . . , a ′ N x ) with a ′ n = e -1 2 τ ( i ) l (cos( τ ( i ) l πn ) + i sin( τ ( i ) l πn )) , we have that, for all complex vectors v ∈ C N x it holds

that for any m ∈ [ L ] , ∥ ( A ′ ) m v ∥ 2 = e -mτ ( i ) l / 2 ∥ v ∥ 2 . Therefore, the operator ( A ′ i ) m does not change the width-scaling. Since, for any u l , setting σ B ∈ Θ( √ N x N u ) yields ∥ ∥ ∥ B ′ ( i ) l u l ∥ ∥ ∥ 2 ∈ Θ( √ N x ) . Therefore, each term in the summation is of order Θ( √ N x ) and unless, for every l , the term B ′ ( i ) l u l perfectly cancels out with the terms before to affect the width scaling, we have that ∥ x ( i ) l ∥ ∈ Θ( √ N x ) . The same argument can be used to show the stability of y ( i ) l .

Proposition 3.4 ( Scale of the updates of hidden states ∆ x ( i ) 1 after 1 step of SGD ) . Under the ZOH discretization procedure, as N x then N u approach infinity, for every 1D SSM, the squared l 2 -norm of the updates ∥ ∆ x ( i ) 1 ∥ 2 of the hidden states after one step of SGD is a.s. scaled as ∥ ∆ x ( i ) 1 ∥ 2 ∈ Θ ( η B 1 √ N u σ C ∥ u 1 ∥ 3 2 ζ (4) 1 2 ) , where ζ (4) denotes the Riemann zeta function at 4.

Proof of Proposition 3.4 . Following condition ( ∆ ), the features of the 1D SSM evolve non-trivially with width, when the following condition is satisfied:

Scale of the updates ∆ x ( i ) l after 1 SGD step.

First, note that, for the discretization step ZOH , we have

For any quantity · , letting ˜ · denote the updated quantity after one step of SGD, we can write

For clarity, let us again begin by considering the scale of x ( i ) 1 . We can follow the same argumentation as before to generalize the result to arbitrary l ∈ [ L ] . The updates ∆ x ( i ) 1 can therefore be computed as ˜ B ′ ( i ) 1 ˜ u ( i ) 1 -B ′ ( i ) 1 u ( i ) 1 , where ˜ B ′ ( i ) 1 = (( ˜ A ′ ) ( i ) -I )( ˜ A ( i ) ) -1 ˜ B 1 and ˜ A ′ ( i ) = exp( τ ˜ A ( i ) ) . Since u ( i ) l and ˜ u ( i ) l are assumed to be Θ(1) due to Assumption 3.1, the scale of ∆ x ( i ) 1 is determined by that of ˜ B ′ ( i ) 1 and ˜ B ′ ( i ) 1 -B ′ ( i ) 1 .

Let's begin by understanding the scale of ( ˜ A ′ ( i ) -I )( ˜ A ( i ) ) -1 (∆ W B u 1 ) . Since ∆ W B = -η B ¯ B 1 u T 1 , letting ˜ Λ i = ( ˜ A ′ ( i ) -I )( ˜ A ( i ) ) -1 , we can express the l 2 -norm of ( ˜ A ′ ( i ) -I )( ˜ A ( i ) ) -1 ∆ W B u 1 as follows:

Note that since ˜ Λ i ∆ W B is a rank-one matrix, its spectral norm can be decomposed as

Therefore, since ∥ u 1 ∥ 2 ∈ Θ( √ N u ) , if the spectral norm of Λ i ∆ W B is scaled as Θ( √ N x N u ) ,

Claim. If ∥ ∆ W B ∥ ∗ ∈ Θ( √ N x N u ) , then ∥ ∥ ∥ ( ˜ A ′ ( i ) -I )( ˜ A ( i ) ) -1 ∆ W B u 1 ∥ ∥ ∥ 2 ∈ Θ( √ N x ) .

Essentially, this follows if ( ˜ A ′ ( i ) -I )( ˜ A ( i ) ) -1 does not affect the width-scaling of the spectral norm of ∆ W B , that is, ∥ Λ i ∆ W B ∥ ∗ = ∥ ∆ W B ∥ ∗ . This is equivalent to showing that ∥ ∥ ∥ ˜ Λ i ¯ B 1 ∥ ∥ ∥ 2 = ∥ ∥ ¯ B 1 ∥ ∥ 2 .

Note the following Lemma which states that the scaling of ∥ ˜ Λ i ¯ B 1 ∥ 2 is identical to that of ∥ A -1 ¯ B 1 ∥ 2 .

Lemma C.3. Let Q i, 1 = (exp ( -η a ¯ A ( i ) A ( i ) ) . For any complex vector v ∈ C N x , the width-scaling of v remains invariant under the operator ( ˜ A ′ ( i ) -I ) Q -1 i, 1 , that is, ∥ ( ˜ A ′ ( i ) -I ) Q -1 i, 1 v ∥ 2 ∈ Θ( ∥ v ∥ 2 ) .

Proof of Lemma C.3. Recall that ˜ Λ i = ( ˜ A ′ ( i ) -I ) ˜ A -1 i , where ˜ A i = A ( i ) Q i, 1 and

Since [ A ( i ) ] j = -( j + 1) , it's easy to see that, for any positive value of τ ( i ) , the eigenvalues of τ ( i ) A ′ ( i ) l -( A ( i ) ) -1 ( A ′ ( i ) l -I ) are negative, in Θ(1) , and converge to 0 from below. The expression evaluates to

Since ¯ B ′ ( i ) 1 = u ( i ) 1 ¯ x ( i ) 1 and ¯ x ( i ) 1 = ¯ y ( i ) 1 C T 1 ,

Due to Assumption 3.1, ¯ y ( i ) 1 is distributed i.i.d as N u → ∞ and ¯ y ( i ) 1 ∈ Θ( 1 N u ) (Yang and Hu, 2021). A CLT argument reveals that, under the scaling σ B ∈ Θ( √ N x N u ) and σ C ∈ Θ( √ 1 N x N u ) , the eigenvalues of ( ¯ B ′ ( i ) 1 B ( i ) T 1 ⊙ I ) are in Θ( 1 N u ) . Lemma C.4 follows as a consequence of this result.

Lemma C.4. As N u then N x approach infinity, ∥ ˜ A ( i ) -A ( i ) ∥ ∗ ∈ Θ(1) and ∥ ˜ A ( i ) ∥ ∗ ∈ Ω(1) if and only if η a ∈ Θ( N u ) .

Accordingly, all the eigenvalues of Q i, 1 are in Θ(1) and therefore, for any complex vector v , ∥ Q i, 1 v ∥ 2 ∈ Θ( ∥ v ∥ 2 ) . Together with the fact that the spectrum of A ′ ( i ) is Θ(1) , we have ˜ Λ i ′ ∈ Θ(1) . Further note that the entries of ˜ A ( i ) and consequently those of ˜ Λ i are independent of each other since the entries [ ¯ B ′ ( i ) 1 B 1 ( i ) T ] m,m are independent for different m ∈ [ N x ] .

Therefore, for the correct scaling of the updates, η B should be scaled as Θ( N x √ N u ) . Similar computations reveal that η C should be scaled as Θ( 1 √ N x N u ) .

Proposition 3.5 ( Scale of the updates of outputs ∆ y ( i ) 1 after 1 step of SGD ) . Under the ZOH discretization procedure, as N u then N x approach infinity, for every 1D SSM, the squared l 2 -norm of the updates of the hidden states after one step of SGD scales as | ∆ y ( i ) 1 | 2 ∈ Θ ( η C σ 2 B √ 1 N u ∥ u 1 ∥ 4 2 ) .

Before we present the proof of Proposition 3.5, we first provide the following Lemma.

Lemma C.5. As N u →∞ , N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 converges to a normal distribution with mean 0 and variance N u Var (¯ y ( i ) 1 u ( i ) 1 ) . Note that since we assume that the rest of the network is scaled correctly (under µ P), ¯ y ( i ) 1 is distributed i.i.d asymptotically and ¯ y ( i ) 1 ∈ Θ( 1 N u ) (Yang and Hu, 2021). Therefore N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 ∈ Θ( 1 √ N u ) .

Proof. For i ∈ [ N u ] , the output of the 1 D SSM is given by y ( i ) 1 = C 1 x ( i ) 1 where C 1 = ( W C u 1 ) T and x ( i ) 1 = B ′ ( i ) 1 u ( i ) 1 = ( A ′ ( i ) -I ) A ( i ) -1 W B u 1 . (20) Using the notation ˜ · to denote the updated quantity · after 1 step of SGD,

The updated quantities ˜ W B and ˜ W C are derived as

˜ W B = W B +∆ W B , ˜ W C = W C +∆ W C with ∆ W B = -η B ¯ W B , ∆ W C = -η C ¯ W C , where ¯ W B and ¯ W C denote the gradient of the loss with respect to the quantities W B and W C respectively and can be computed according to

Also, ˜ Λ i = ( ˜ A ′ ( i ) -I ) ˜ A ( i ) -1 , where ˜ A ( i ) = Diag(exp( a log i -η a ¯ A ( i ) log )) = A ( i ) Q i , with Q i = ( exp( -η a ¯ A ( i ) A ( i ) ) ⊙ I ) and

Therefore ˜ y ( i ) 1 can be represented as follows

˜ y ( i ) 1 = [ ( W C +∆ W C )( u 1 +∆ u 1 ) ] T [ ˜ Λ i ( W B +∆ W B )( u 1 +∆ u 1 ) ] ( u ( i ) 1 +∆ u ( i ) 1 ) , (23) and the updates ˜ y ( i ) -y ( i ) can be expressed as

1 1

where M = [ ( W C + ∆ W C )( u 1 + ∆ u 1 ) ] T [ ˜ Λ i ( W B + ∆ W B )( u 1 + ∆ u 1 ) ] ( u ( i ) 1 + ∆ u ( i ) 1 ) -( W C u 1 ) T ˜ Λ i W B u 1 u ( i ) 1 -( W C u 1 ) T Λ i W B u 1 u ( i ) 1 represents the remainder of the terms in (23). First, let us consider the scaling of the update term without remainder term,

where K = ( ˜ Λ i -Λ i ) .

Letting v m = N u ∑ m ′ ,m ′′ =1 W Cm,m ′ W Bm,m ′′ u 1 ( m ′ ) u 1 ( m ′′ ) , as N u →∞ , { v m } N x m =1 are independent and are distributed normally with mean 0 and variance σ 2 B σ 2 C ∥ u 1 ∥ 4 by an application of LindenbergFeller Central Limit Theorem.

Independence of entries of K . Next, we show that the entries of the diagonal matrix K are independent of each other.

Observe that

Since the entries of A are deterministic, and we assume that τ ( i ) is some fixed deterministic scalar, the entries of K are independent of each other iff the entries of ( ¯ B ′ ( i ) 1 B ( i ) ⊺ 1 ⊙ I ) are independent. Now consider

Since W B and W C are Gaussian matrices with i.i.d entries and the entries of u 1 are assumed to be i.i.d, the entries [ ¯ B ′ ( i ) 1 B 1 ( i ) T ] m,m are independent for different m ∈ [ N x ] .

Furthermore, following the same line of argumentation as in Lemma C.3, the matrix K can be represented as K = K ′ A -1 , where the entries of K ′ are Θ(1) and accordingly for any vector the scaling of v remains invariant under the operator K ′ .

To compute the scale of N x ∑ m =1 K m,m v m , we can therefore instead consider the scaling of N x ∑ m =1 A -1 m,m v m where { v m } N x m =1 are independent and are distributed normally with mean 0 and variance σ 2 B σ 2 C ∥ u 1 ∥ 4 .

As N x → ∞ , N x ∑ m =1 A -1 m,m v m converges to a normal distribution with mean 0 and variance c ζ (2) σ 2 B σ 2 C ∥ u 1 ∥ 4 , where ζ (2) denotes the Riemann zeta function at 2 and c ζ (2) > 0 evaluates to a width-independent constant.

Scale of M .

Note that since u ( i ) 1 and ∆ u ( i ) 1 are scalars and in order 1 , the scale of M is determined by that of [( W C +∆ W C )( u 1 +∆ u 1 )] T [ ˜ Λ i ( W B +∆ W B )( u 1 +∆ u 1 )] .

Recall that both u 1 and ∆ u 1 are assumed have i.i.d coordinates and have the scaling ∥ u 1 ∥ ∈ Θ( √ N x ) and ∥ ∆ u 1 ∥ ∈ Θ( √ N x ) . Therefore, we only need to consider the scaling of

It is easy to verify that the scaling of ( W C u 1 ) T ( ˜ Λ i -Λ i ) W B u 1 is the same as that of ( W C u 1 ) T ( ˜ Λ i ) W B u 1 u ( i ) 1 following the same arguments as above.

Recall that the updates ∆ W B and ∆ W C are defined as

where ¯ W B and ¯ W C denote the gradient of the loss with respect to the quantities W B and W C respectively and computed as

So, ∆ W C u 1 = -η C N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 Λ i W B u 1 u 1 T u 1 = -η C ∥ u 1 ∥ 2 Λ i W B u 1 N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 ( Λ i is identical for all i ∈ [ N u ] ).

Therefore, letting Ξ = -η C ∥ u 1 ∥ 2 N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 , we get

Letting w m = N u ∑ m ′ ,m ′′ =1 W Bm,m ′ W Bm,m ′′ u 1 ( m ′ ) u 1 ( m ′′ ) , it is easy to verify that as N x → ∞

̸

each v m is distributed normally with mean σ 2 B ∥ u 1 ∥ 2 and variance σ 4 B N u ∑ m ′ = m ′′ =1 u 1 ( m ′ ) 2 u 1 ( m ′′ ) 2 . Furthermore { w m } N x m =1 are independent.

Therefore, due to Kolmogorov's SLLN, as N x →∞ , it holds that

Combining this with the result from Lemma C.5, we conclude that as N u → ∞ then N x → ∞ , (∆ W C u 1 ) T ( ˜ Λ i W B u 1 ) ∈ Θ( η C σ 2 B ∥ u 1 ∥ 4 1 √ N u ) .

Therefore, for the updates ∆ y ( i ) 1 to be of order 1 , η C must be scaled as Θ( 1 N x √ N u ) .

Scale of (∆ W C u 1 ) T ( ˜ Λ i )∆ W B u 1 u ( i ) 1 .

Recall that the updates ∆ W B and ∆ W C are defined as

where ¯ W B and ¯ W C denote the gradient of the loss with respect to the quantities W B and W C respectively and computed as

Hence, we get ∆ W C u 1 = -η C N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 Λ i W B u 1 u 1 T u 1 = -η C ∥ u 1 ∥ 2 Λ i W B u 1 N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 ( Λ i is identical for all i ∈ [ N u ] ) and ∆ W B u 1 = Λ i W C u 1 u 1 T N u ∑ i =1 ¯ y ( i ) 1 u ( i ) 1 , which yields

Following the same line of argumentation as before, we obtain that as N u → ∞ then N x → ∞ , (∆ W C u 1 ) T ( ˜ Λ i ∆ W B u 1 ) ∈ Θ( σ B σ C η B η C ∥ u 1 ∥ 6 1 N u ) .

Substituting σ B ∈ Θ( √ N x N u ) , σ C ∈ Θ( 1 √ N x N u ) , η B ∈ Θ( N x √ N u ) yields the same scaling: η C ∈ Θ( 1 N x √ N u ) .

Scale of ( W C u 1 ) T ( ˜ Λ i )∆ W B u 1 u ( i ) 1 . Following the same steps as for the previous terms, we obtain ( W C u 1 ) T ( ˜ Λ i )∆ W B u 1 u ( i ) 1 ∈ Θ( η B 1 √ N u σ 2 C ∥ u 1 ∥ 2 ) . This term vanishes under the scaling η B ∈ Θ( N x √ N u ) and σ C ∈ Θ( 1 √ N x N u ) .

C.3 Time-invariant S4 recurrent layer

The S4 recurrent layer y 1: L = f S4 ( u 1: L ; w = { B , C } ) is a sequence to sequence mapping:

for l = 1 , . . . , L with x 0 = 0 , which is a discretization of the continuous-time SSM:

where R ( · ) gives the real part of a complex vector, A = diag ( a ) and a ∈ C N x , B ∈ C N x × N u , C ∈ C N y × N x .

The Zero-Order-Hold (ZOH) discretization method gives that:

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 n = -1 2 + iπn . 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 = -( n +1) may perform similarly well.

The backward pass for f S 4 is as follows:

Since we want all learnable weight matrices to learn features, we split each S4 layer into two sublayers and demand them both to admit feature learning. The first one being the state space equation (26) and the second one being the decoder (27).

Claim C.1 ( Scale of hidden states x 1 in S4 at initialization ) . Under the ZOH discretization procedure, as N x and N u approach infinity with N u /N x ∈ Θ(1) , for any i ∈ [ N u ] , the squared l 2 -norm of the hidden states ∥ ∥ ∥ x ( i ) 1 ∥ ∥ ∥ 2 2 is almost surely scaled as ∥ ∥ ∥ x ( i ) 1 ∥ ∥ ∥ 2 2 ∈ Θ ( ζ (2) σ 2 B ∥ u ∥ 2 ) , where ζ (2) is the Riemann zeta function at 2.

Proof. Stability at initialization.

First, we show that the size of x 1 = B ′ u 1 at initialization is Θ( √ N x ) . By definition,

where Λ = ( A ′ -I ) A -1 . Then, writing u = u 1 for readability, observe that

is a sum of N x independent but 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 Λ and the correct scaling of the inputs to the S4 layer.

Verifying that the sequence satisfies Kolmogorov condition.

For any i ∈ [ N x ] ,

̸

where ζ (4) denotes the Riemann zeta function at 4 for which a closed form is given by |B 4 | (2 π ) 4 2 · 4! . Note that, due to the rapid decay of exp( -( i +1) · τ ( i ) l ) → 0 with i →∞ , the zeta function also lower bounds the spectral sums N x ∑ i =1 Λ 2 i ≥ c · ζ (2) and N x ∑ i =1 Λ 4 i ≥ c · ζ (4) for some width-independent constant c ∈ (0 , 1) for N x large enough.

By assumption, the inputs to the S4 layer are scaled such that ∥ u ∥ ∈ Θ( √ N x ) and therefore, as long as σ 2 B ∈ O (1) the sequence of random variables satisfies the Kolmogorov condition.

Limiting behavior of ∥ B ′ u ∥ 2 2 .

Applying SLLN, we have

̸

where ζ (2) denotes the Riemann zeta function at 2 and equates to π 2 / 6 , and c ∈ (0 , 1) is some width-independent constant.

Spectral scaling does not yield the correct scale of initialization for σ B .

For spectral scaling conditions to yield the right scaling of the initialization variance, it is crucial that the following condition holds:

Since the spectrum of ( A ′ -I ) A -1 is less than 1 , an upper bound on the spectral norm of B ′ = ( A ′ -I ) A -1 B can be found in (Vershynin, 2011) and is given by C ( √ N x + √ N u ) 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.

Correct scale of σ B for stability at initialization.

We know that ∥ B ′ u ∥ 2 ∈ Θ( σ B ∥ u ∥ 2 ) and since ∥ u ∥ 2 ∈ Θ( √ N u ) (since we assumed stablity of activations in the previous layer), imposing stability of ∥ B ′ u ∥ 2 yields the following scaling: σ B ∈ Θ( √ N x N u ) .

Stability of x l for arbitrary l ∈ [ L ] .

For all l ∈ [ L ] , we have x l = ∑ l -1 m =1 ( A ′ ) m B ′ u l -m = A ′ ( ∑ l -2 m =1 ( A ′ ) m B ′ u l -m )+ B ′ u l . First, observe that since A ′ = diag ( a ′ 1 , . . . , a ′ N x ) with a ′ n = e -1 2 τ ( i ) l (cos( τ ( i ) l πn ) + i sin( τ ( i ) l πn )) , we have that, for all complex vectors v ∈ C N x it holds that for any m ∈ [ L ] , ∥ ( A ′ ) m v ∥ 2 = e -mτ ( i ) l / 2 ∥ v ∥ 2 . Therefore, the operator ( A ′ ) m does not change the width-scaling. We showed earlier that for any u l , setting σ B ∈ Θ( √ N x N u ) yields ∥ B ′ u l ∥ 2 ∈ Θ( √ N x ) . Therefore, each term in the summation is of order Θ( √ N x ) and unless, for every l , the term B ′ u l perfectly cancels out with the terms before to affect the width scaling, we have that x l ∈ Θ( √ N x ) for all l ∈ [ L ] .

Claim C.2 ( Scale of output y ( i ) of S4 at initialization ) . Under the ZOH discretization procedure as N x and N u approach infinity with N u /N x ∈ Θ(1) , for any i ∈ [ N u ] , the output y ( i ) 1 converges in distribution to a Gaussian with mean 0 and standard deviation Cσ B σ C ∥ u 1 ∥ 2 2 for some widthindependent 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 ( σ B and σ C ).

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

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 x 1 = B ′ u 1 after a single step of SGD on ( ( u 1 , y 1 ) ).

Since the update matrix ∆ B ′ is a rank one matrix, we have

The updates ∆ x 1 are given by ∆ x 1 = ( B ′ + ∆ B ′ )( u 1 + ∆ u 1 ) -B ′ u 1 = B ′ ∆ u 1 + ∆ B ′ u 1 + ∆ B ′ ∆ u 1 .

If the spectral norm of ∆ B ′ is scaled as Θ( √ N x N u ) , then ∆ B ′ u 1 ∈ Θ( √ N x ) . We also have that

∆ B ′ ∆ u 1 ≤ ∥ ∆ B ′ ∥ ∗ ∥ ∆ u 1 ∥ 2 ∈ O ( N x ) . Since ∥ ∆ u 1 ∥ 2 ∈ Θ( N u ) , when σ B is set to Θ( √ N x N u ) , B ′ ∆ u 1 ∈ Θ( √ N x ) . Therefore, unless the scale of ∆ B ′ u 1 perfectly cancels out with the remaining two terms, ∥ ∆ x 1 ∥ ∈ Θ( √ N x ) .

Generalizing to more gradient steps (and l > 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 parameterization for other parts of the model are already addressed in Yang and Hu (2021), Yang et 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 µ P-SSM parameterization exhibits better monotonicity and stability.

Test loss

-SSM (Ours)

SP

(Heuristic)

3.0

2.9

256

512

2.8

1024

2048

4096

2.7

2.6

2.5

2.4

2.3

-10

8-

-6

-4

-10

-8

-6

-4

-10

8-

9-

-4

log2 (lr ssm)

In Figure 5, we present results on a randomly sampled subset of the Fineweb dataset. While computational constraints prevented us from training on the entire dataset or using larger model widths, at small scales, our observations on Fineweb align with the wikitext-103 results reported in the main paper.

SP

Test loss

-SSM (Ours)

(Heuristic)

12.0

10.0

8.0

6.0

4.0

-10

8-

-6

4

-2

-10

-8

-6

-4

-2

-10

-8

-6

-4

-2

log2(lr)

log2(lr)

log2(lr)

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 in 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.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of 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 in 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 on 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 a few datasets or with a 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, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a 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 be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We state all assumptions in the main paper and provide all formal proofs in the appendix.

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 crossreferenced.

All assumptions should be clearly stated or referenced in 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 a short proof sketch to 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

Question: Does the paper fully disclose all the information needed to reproduce the main experimental 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)?

Answer:[Yes]

Justification: Experimental details relevant for reproducing the results to the extent that it affects the main claims are disclosed in the main paper.

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 a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to 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 to 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 a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe 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 be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to 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 hyperparameters 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 NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a 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: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are 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) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, 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 a 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, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: In Figure 1, we report σ -confidence intervals over 10 runs. Due to limited computational resources, we were unable to repeat the large-scale experiment for Figure 3 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 error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

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.

It 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 a 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 of range (e.g. negative error 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 in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

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 the type of compute workers CPU or GPU, internal cluster, or 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 in 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 in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: We provide a theoretical analysis of 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 NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration 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 or data.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

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), fairness 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 not tied to particular applications, let alone deployments. However, if there is a 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 as 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 or 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 a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

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 a 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 not use 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 is used and, if 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.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a 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 is not available online, the authors are encouraged to reach out to the asset'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/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is 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. Crowdsourcing 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 as 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 in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as 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 an equivalent approval/review based on the requirements of your country or institution) were obtained?

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.

Depending on 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.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nUnderstanding 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\nNone.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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\n\nWEAKNESSES:\n- 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\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nFollowing 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.\n\nSTRENGTHS:\nThis 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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$\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-29"}