Quantifying Memory Utilization with Effective State-Size
Authors: Rom Parnichkun
Neehal Tumma
Armin W Thomas
Alessandro Moro
Qi An
Taiji Suzuki
Atsushi Yamashita
Michael Poli
Stefano Massaroli
Venue: ICML 2025 poster (2025)
Keywords: model analysis
interpretability
linear systems
attention
state-space models
sequence models
memory utilization
context utilization
Primary area: deep_learning->sequential_models_time_series
TL;DR: We propose the Effective State-Size (ESS) metric to quantitatively analyze how sequence models utilize memory and context, offering insights into model in-context recall, initialization strategies, efficiency, and architecture design.

Abstract: As the space of causal sequence modeling architectures continues to grow, the need to develop a general framework for their analysis becomes increasingly important. With this aim, we draw insights from classical signal processing and control theory, to develop a quantitative measure of *memory utilization*: the internal mechanisms through which a model stores past information to produce future outputs. This metric, which we call ***effective state-size*** (ESS), is tailored to the fundamental class of systems with *input-invariant* and *input-varying linear operators*, encompassing a variety of computational units such as variants of attention, convolutions, and recurrences. Unlike prior work on memory utilization, which either relies on raw operator visualizations (e.g. attention maps), or simply the total *memory capacity* (i.e. cache size) of a model, our metrics provide highly interpretable and actionable measurements. In particular, we show how ESS can be leveraged to improve initialization strategies, inform novel regularizers and advance the performance-efficiency frontier through model distillation. Furthermore, we demonstrate that the effect of context delimiters (such as end-of-speech tokens) on ESS highlights cross-architectural differences in how large language models utilize their available memory to recall information. Overall, we find that ESS provides valuable insights into the dynamics that dictate memory utilization, enabling the design of more efficient and effective sequence models.

Quantifying Memory Utilization with Effective State-Size
Rom N. Parnichkun 12 Neehal Tumma "1 Armin W. Thomas'
Alessandro Moro 2 Qi An Taiji Suzuki Atsushi Yamashita
Michael Poli 14 Stefano Massaroli 133
Abstract
1. Introduction
In recent years, the success of autoregressive sequence modeling in the context of deep learning has largely been driven by advancements in highly parallelizable causal ar- chitectures, such the transformer (Vaswani al., 2017). However, despite their strong performance and hardware efficiency, understanding the inner workings of these neu- ral networks remains a challenging task due to their non- linearity and the diversity of fundamental building blocks used. To this end, we leverage a new class of model ab- stractions, allowing for the development of a unified frame- work for the analysis of these computational units.
As the space of causal sequence modeling ar- chitectures continues to grow, the need to de- velop a general framework for their analysis be- comes increasingly important. With this aim, we draw insights from classical signal processing and control theory, to develop a quantitative mea- sure of memory utilization: the internal mecha- nisms through which model stores past infor- mation to produce future outputs. This metric, which we call effective state-size (ESS), is tai- lored to the fundamental class of systems with imput-invarian and input-varying linear opera- tors, encompassing a variety of computational units such variants of attention, convolutions, and recurrences. Unlike prior work on mem- ory utilization, which either relies on raw oper- ator visualizations (e.g. attention maps), or sim- ply the total memory capacity (i.e. cache size) of a model, our metrics provide highly inter- pretable and actionable measurements. In partic- ular, we show how ESS can be leveraged to im- prove initialization strategies, inform novel regu- larizers and advance the performance--fficcenee frontier through model distillation. Furthermore, we demonstrate that the effect of context delim- iters (such as end-of -f--ooeeeeech tokens) on ESS high- lights cross-architectural differences in how large language models utilize their available memory to recall information. Overall, find that ESS provides valuable insights into the dynamics that dictate memory utilization, enabling the design of more efficient and effective sequence models.
In particular, note that the majority of sequence mod- els of practical interest can formally be expressed as either linear systems (y = Tu) or systems with input-varying linear operators (y = fr(u)),, the latter of which we abbreviate as LIV'. Like the decoupling of the "attention matrix" from the full nonlinear attention operation, the LIV framework decouples the input-varying construction of matrix T (u + T fr())) from the linear mapping = Tu). This decoupling enables a wide array of deep learning primitives beyond attention variants [67; 34; 64] to be uniformly formulated and analyzed, including sequence models like convolutions [40; 46; 35; 55) and linear state- transition recurrences [15: 18; 27; 60; 29; 26; 74; 48; 16).
Current approaches analyzing the inner workings of LIVs often rely upon simple visualizations of T (or the ag- gregation of 7 across multiple layers and residuals) [44; 68; 1; 4: 73; 62]. However, these visualizations alone often fail to highlight critical properties that explain how different models construct internal representations of the input data. Moreover, prior attempts in obtaining quantitative metrics, such as through spectral analysis of the operator T (Min & Li, 2024; Bhojanapalli et al., 2020), are either limited to specific model class or do not appropriately take into ac- count important conflating factors like the causal masking
'Equal contribution. The order of authorship was determined by coin flip. 'Liquid AI 'The University of Tokyo 'RIKEN 4 "Stanford University. Correspondence to: Rom N. Parnichkun <rom coom.pamichkune gmail.com>.
We abbreviate them as LIVs instead of IVLs to conform to the classical convention, where linear time-varying and time- invariant systems are abbreviated as LTTSs and LTIs respectively, as well as recent works such as (Zancato et al., 2024), which call them "linear input-varying systems".
Proceedings of the 1nnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Quantifying Memory Utilization with Effective State-Size
Causal Sequence
Realized Linear Structure (y = Tw)
Effective State-Sizz (ESS)
ESS-En Enabled
Models
rank(Te,,-1) i [1,4]
Antalysis and Techniques
Input-Vary Vere-arriing//vaiinnttt
Linear Operators
Model Distitaatin
Attention
Convolttion
Initializer and Featurizer Analysis
Recurencc
Sequeemce Pider
Other Structures
->
Memory Utilization Analysis
ANN LSTM GRU, elc.
Figure I. An overview of the effective siate-size (ESS) metric and its various downstream applications. A linear matrix T is extracted from causal sequence models, and the rank of its submatrix used as a proxy for memory utilization model analysis and other heuristic-based techniques.
of T which significantly distorts the metric as discussed in (Wu et al.).
crucial for in-context recall.
In this work, we focus our analysis on the working mem- ory² of model architectures and examine two aspects of model memory in particular: memory capacity (i.e. cache/state size) and memory utilization. Notably, mem- ory capacity alone can be misleading, models with sim- ilar capacities may learn to utilize their available memory to varying degrees. Therefore, we introduce the notion of memory utilization measure that provides deeper in- sight into the differences between architectures with com- parable computational efficiency.
2. Related Work
In this section, we briefly describe previous works that are closely related to the core findings of the paper. For an extended discussion of the related works, refer Section B.
Classifying sequence models. sequence model can be defined as mapping from input sequences to output sequences that leverages state to maintain information across time. Within the scope of deep learning, the early attempts at constructing sequence models fall under the class of non-linear RNNs (Rumelhart & McClelland, 1987; Hochreiter & Schmidhuber, 1997) which exhibit non-linear state space dynamics. However, in recent years, mod- els with linear state transitions have grown in popularity due to inherent trainability (Pascanu et al., 2013) and ef- ficiency (Martin & Cundy, 2018) limitations in the non- linear regime. These models fall under two categories: linear systems and systems with input-varying linear op- crators (LIVs). We note that linear systems can further be broken down into time-invariant (LTI) and time-varying (LTV) systems, depending on whether or not the param- eters of the realized recurrence change as function of time (i.e. sequence index). This taxonomy has also been used in prior work (Zancato et al., 2024) to motivate the construction of novel architectures. In modern deep Iearn- ing, LTI systems (Lindquist & Picci, 1979; Ljung, 1999; Ho & Kalman, 1966) can be found in frameworks like $4 (Gu et al., 2022a). However, given the lack of expressiv- ity afforded to LTIs, the current state-of-the-ar models fall under the more expressive LIV class (Gu & Dao, 2024; Poli et al., 2023; Vaswani et al., 2017) which parameterizes the state-space dynamics as function of the input. Since LIVs constitute the minimal superclass containing most se- quence models of practical interest, the goal of this work is to develop a quantitative means of analyzing models under
Our main technical contributions can be summarized as fol- lows:
We draw from classical signal processing and control theory, and propose effective state-size (ESS) met- ric computed by taking the rank of a: proxy for memory utilization in LIVs (Section 3).
We validate ESS beyond its theoretical interpretation by demonstrating its correlation with performance across a wide range of models and memory-intensin synthetic tasks, including associative recall and selec- tive copying (Section 4).
We explore the use of the ESS metric as a means of enhancing the performance-efficiency trade-off by demonstrating its ability to inform model distillation (Section 5.2), initialization schemes (Section 5.1), and regularization strategies (Section A).
We extend the utility of ESS to language, demonstrat- ing how it captures previously uncharacterized prop- erty of I.LMs: state modulation (Section 5.3). This phenomenon provides concrete intuition as to why in- troducing input dependencies into state transitions is
"Here we refer to "memory" in the sense commonly asso- ciated with the "state" of dynamical systems, as described by (Willems, 1989), as opposed to the notion of language models memorizing some fact encountered during training (Allen-Zhu & Li, 2025).
2

---

Quantifying Memory Utilization with Effective State-Size
the LIV abstraction. To do so, we draw from minimal re- alization theory (DeWilde & van der Veen, 1998; Akaike, 1974), which has previously been applied within the scope of classical linear systems but has yet to be extended to modern LIV sequence models.
A unified representation of sequence models. While typically nonlinear, most sequence models of interest can effectively materialize a linear operator T,, where the equa- tion y = Tu faithfully expresses the computation per- formed by the model (see Section D.2 for further elabo- ration):
Tú = QiKj = C,B;
linear attention, recurrence, convolution,
Quantifying memory utilization in sequence models. As discussed in Section I, the current landscape of ex- amining memory in sequence models primarily reduces to qualitative measures or limited quantitative ones. Here, we expand on the latter, noting that commonly used measures such model size (Hoffmann et al., 2022) and cache/state size (which we define formally as theoretically realizable state-size in Section 3) are imperfect means of measuring memory in sequence models. Namely, while these metrics serve as reasonable proxies for the capacity of a model to leam, they fail to capture how much of that capacity is real- ized. As such, they ignore important aspects of the model pipeline such as data, initialization, and optimization. To capture these aspects of network training, we apply the no- tions of numerical and effective rank (Roy & Vetterli, 2007) to the minimal realization problem, which gives rise to the ESS metric. We elaborate on this in Section 3.
T = C.AA-I A5t1
Tú = -li-j
Tu -CHHi-jBB
gated convolution, aftention.
Tú = Ko)==
==
We make distinction between linear systems (such as convolutions) and LIVs (such as attention and gated con- volutions). The operator T of linear systems are input- invariant, whereas for LIVs, they are materialized via causal featurizers that map past inputs into features, i.e. fn Uu + B,, which are then used construct the ele- ments Tf I as outlined above. We begin our formulation by restricting the scope to linear systems.
Operator-recurrence duality and the minimal recurrent realization of linear systems. Consider a general linear recurrence formulated as follows:
We highlight the concept of semiseparable rank (Vandebril et al., 2005; Xia al., 2010), which has been recently ex- plored in (Dao & Gu, 2024). Similar to ESS, semiseparable rank examines the rank of specific submatrices of T. How- ever, as semiseparable rank is primarily motivated by the design of efficient algorithms, we find it most closely aligns with the concept of theoretically realizable state size (TSS) in the context of our work. In this study, we provide com- plete derivations of ESS, clearly distinguish it from TSS (and by extension semiseparable rank), and demonstrate its effectiveness as a proxy for memory utilization in practical settings for modern LIV sequence models.
8i+1 = A;s; + Bug Vi =CC, + Dju
(1)
where (A, E R B, E R xd C, E Rdd D, E Rdxd) 8; and n, are the state and state-size at quence index i respectively. Classic results (DeWilde & van der Veen, 1998, ch. 3) establish that for any given input-invariant operator T (i.c. linear system), there exist infinite recurrent realizations in the form of Equation (1), motivating the search for the minimal one.
Given any causal iapur-invarian operator T, there exist infinite variations of linear recurrences in the form of Equation (1) that realize an equivalent input- output operator.
Theorem 3.1.
3. Theory
This section begins by showing that most modern sequence models can be formulated to materialize linear operator
A simple extension of the proof of this theorem (in Section C.2.4) demonstrates the following:
T. We then formally define ESS metric derived from T. providing theoretical insights grounded in minimal re- alization theory. Finally, we detail the practical compu- tation of ESS, proposing two variants: tolerance-ESS and entropy-ESS.
Theorem 3.2. The rank of the operator submatrix (H₃ =
Te,s-1) determines the minimal state size required to represent the causal operation (y Tu) as a recurrence.
Refer to Section C.3 for the proof. Figure 2 provides a high-level visualization of how the minimal state-size is computed from a causal operator T.
Using a flattened notation, we let T E ddtxdt u,y E gddt denote the operator, inputs, and outputs respectively, l de- note the sequence length and d denote the channel dimen- sion. We index sequence indices with subscripts, i.c. Ty E Rdxd U, E Rd and channels (and other nom-temporal di- mensions) with superscripts, i.c. T E R u² e Rd. For additional details on notation, refer to Section C.I.
Note that since we use the flattened notation, where T E dd rank(H;) determines the minimal state-size re- quired for processing all channels in a layer via a recur- rence at the ith sequence index. We formally refer to this metric as per-sequence index effective state-size (ESS) and
3

---

Quantifying Memory Utilization with Effective State-Size
T
linear attention variants, TSS is equal to the state-size de- fined by their recurrent formulation. We refer readers to Section D.2 for detailed operator-specific derivations of TSS. Recall that ESS depends not only on the model's 'ss func- tional form, but also on the input data, optimization, and more generally anything that impacts the realization of T'; in contrast, TSS is limited in that it is determined solely by the model's structure.
Hoz
3.1. Computing Effective State-Size
752
In practice, computing ESS requires a few additional con- siderations due to the numerical errors and approximations involved in computing the rank of matrices. We propose two approaches both of which rely on singular values (E,) from taking the singular value decomposition (SVD) of H; (which equals fr( for LIVs) - that provide complementary perspectives on the same metric.
(aak
Figure 2. This figure provides a high-level example of how ESS is computed from the operator T: (1) submatrix H. (highlighted in pink) is extracted from T, and (2) its rank is computed to obtain the effective state-size (ESS) for each index i. TSS is the maxi- mum possible ESS of an operator, determined by the operator's hyperparameters.
Here, a tolerance value T is manually se- lected to threshold the singular values of Hi, determining the tolerance- metric as follows:
Tolerance-ESS,
discuss its interpretation for both linear systems and LIVs below.
Doleranc--ES(HH,17) 7) = Ha,," : o >)),
(2)
where of"" Σ,, U.diag(E.)V SVD(H.). According to the Sckart-Young Saannngg-irkyy theorem, the tolerance-ESt metric can be interpreted the minimum state size nec- essary for an input- input-invariant recurrence to approximate the original operator, such that the spectral norm of the approx- imation error remains below the specified tolerance level (I] - Tylle S).
Interpreting effective state-size. As shown in Theorem 3.2, the ESS of an input-invariant linear system is given by its minimal state-size which is directly interpretable as a measure of model memory utilization. For LIVs, however, ESS is also function of the input (rank(Jr(w) (.))e.)-))) which means that the minimal realization process outlined in the proof of Theorem 3.1 is no longer guaranteed to obtain recurrences that preserve causality (i.e. the mini- mally realized features A B C depend on future inputs Uk: k > i) Nevertheless, ESS lower bounds the state-size n, (refer to Section C.2.2), meaning that for any LIV, an equivalent recurrence must necessarily materialize state- size at least large as its ESS. Therefore, we claim that ESS serves proxy for memory utilization in not only linear systems, but in LIVs as well. We empirically vali- date the usefulness of this interpretation of ESS for LIVs in Sections 4 and 5.
One drawback of tolerance-ESs. is its re- liance on the somewhat arbitrary selection of a tolerance value. One can instead compute the effective rank (Roy & Vetterli, 2007), which involves exponentiating the normal- ized spectral entropy (perplexity) of H₃:
Entropy-ESS.
cnSropy-ESS(HA)) == =exp( > P." log()")), m
(3)
where p. = To Entropy-ESS is particularly useful for summarizing metrics across the entire tolerance space, whereas tolerance-Ei offers a more precise and readily interpretable depiction of rank concerning approximation error. Unless tolerance is specified, we use entropy-ESS throughout all our experiments. An additional discussion comparing tolerance and entropy ESS, along with our code for computing them, can be found in Section D.I.
Memory capacity in LIVs. The memory capacity of LIV' is given by the state-size n,, We formally refer to it as theoretically realizable state-size (TSS), as it serves as a tight upper bound for ESS. For models without realization- agnostic minimal recurrent formulations, such as softmax attention, we resort to the trivial realization as shown in Equation C.2.5, in which TSS (for single channel) is equal to the sequence index i. For models like SSMs and
Computational complexity of ESS. Since SVD scales cubically with the size of a square matrix, the time com- plexity of computing ESS for a single layer with chan- nels, processing an input of sequence lenghh e, is O ((d))²).
An example of causality preserving realization of LIVs is the trivial realization shown Equation (C.2.5).
4

---

Quantifying Memory Utilization with Effective State-Size
GLA
LA
WLA
SA
1.0
1.0
acturayy
Securayy
aocutacy
acmmney
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
2-3
2-1
2-
2-1
2-
2-
ESS/Kv
ESS/kv
ESS/
(total ESS) /lev
Figure 3. Scatter plots of accuracy vs ESS/kv across featurizers, Within each featurizer plot, all task-model configurations from the sweep corresponding to each featurizer are shown.
Fortunately, most modern sequence models (outside of S5 (Smith et al.)) process the d channels independently (i.e. they are SISO). This means that each layer's 's dl x dl opera- tor T can be decomposed into d independent fx operators, which reduces the ESS computation to an SVD of d inde- pendent operators (T := T a (d). Consequently, per-channel ESS can be computed and summed ob- tain the equivalent per-sequence index ESS. This reduces the computational cost to (dd*). Note that in models like attention, where channels share the same recurrence within a single head, the time complexity is further reduced by a factor proportional to the head dimension. Additionally, in recurrent models with bounded TSS (n < E), truncated SVD can be employed to decrease the time complexity to nddn)).
4. Empirical Validation of Effective State-Size
To demonstrate the practical utility of ESS beyond its the- oretical interpretation discussed in Section 3, we next turn to an empirical analysis. In this section, examine ESS across wide range of tasks and models in order to under- stand how it varies across different regimes, with particular focus placed on its relationship with model performance on memory-intensie tasks.
Task space. To explore ESS in an extensive, yet con- trolled, manner, we iterate on a set of synthetic tasks pro- posed by Poli et al. which have been shown to effectively approximate model performance on large-scale language tasks. Specifically, we train models on the multi-query associative recall (MQAR), selective copying, and com- pression tasks, each of which probes the ability of models to effectively utilize their working memory. We note that here, we restrict the presentation of our results to MQAR and refer the reader to Section E. for the results on selec- tive copying and compression, which showcase analogous trends.
Different forms of ESS. Initially, we formulated ESS as a per-sequence index metric, which is general in that it ap- plies to all models that realize an operator T. Recall in the last section, we extended ESS along the channel di- mension for SISO models. For LIV SISO models par- ticular, we can further extend ESS along the batch size di- mension, since ESS is function of the input. This means that for multi-layer LIV SISO model with m layers, d channels, batch size b, and sequence length e (which encap- sulates the entire set of models we analyze in this work), ESS E mmxdxxxxx Given this tensor, we define average ESS and total ESS as follows:
Model space. We explore four models within the scope of this analysis: gated linear attention (GLA), weighted linear attention (WLA), linear attention (LA) and softmax attention (SA). We choose this set of frameworks since, to- gether, they span various classes of modern sequence mod- els while maintaining similarities in aspects unrelated to their classification, making them suitable for comparison. The key distinctions between these models are as follows (specific details and model formulations can be found in Section D.2):
average ESS mdbl Σ
Σ
ESS, e
aE(d] [E(m) TE]M [E[(]
total ESS- = average ESS d
GLA layer: This layer implements the gated linear at- tention formulation described in (Yang al., a), where the recurrent feature A (gating term) is input-varying, placing it in the same class as models like Liquid-S4 (Hasani et al.. 2023) and Mamba (Gu & Dao, 2024; Dao & Gu, 2024).
We analyze both forms of ESS because for non-recurrent models like softmax attention, the average TSS (which is computed analogously
average ESS) does not depend on model width. In contrast, total TSS does, enabling to capture a measure of ESS and/or TSS that is model- dependent. For more details, refer to Sections D.1 and D.2.
Finally, since average ESS is the metric used throughout most of our experiments, we will refer to it as ESS unless otherwise specified.
WLA layer: This layer is nearly identical to GLA, but with an input-invariant A matrix. This lies in the same class as Hyena-S4D S4D (Poli et al., 2023), RetNets (Sun
5

---

Quantifying Memory Utilization with Effective State-Size
GLA
LA
WLA
S.A
(toda) 1:850-000 OBF
349
27
3.0
T58
2
00.0
0.9
1
ESS. ESS.aco,
02000
0206
2
0..5
0.6
D3
ESS 000
247
102
0.3
07
-9.0
Figure 4. Correlation between ESS and accuracy over the course of model training bucketed by TSS and kv.
et al., 2023), and gated-convolutions in general.
structing an adjusted form of ESS given by ESS/Kv.
Measured over the entire task-model space, ESS/Kv ex- hibits a significantly higher correlation with accuracy than TSS/kv (Figures 3, 5, 1a, 1 1bb).
LA layer: This layer is based on (Katharopoulos et al., 2020); A is not trainable and is instead fixed as the identity matrix.
Note that the strong correlation between ESS/kv and accu- racy highlights the efficacy of ESS as a proxy for memory utilization. Furthermore, this finding underscores a signif- icant gap in the explanatory power between ESS and TSS, emphasizing the importance of analyzing models beyond just their memory capacity.
SA layer: This is the canonical attention layer which is similar to linear attention, but with the addition of softmax non- on-liinarrty applied to the attention matrix (Vaswani et al., 2017), enabling unbounded TSS.
Experimental setup. In our analysis, we exhaustively sweep across the tasks and models (which are comprised of two sequence-mixing and two channel-mixing layers) detailed above. Within each task, we also sweep across varying task difficulties. In the case of MQAR, we do so by modulating the number of key-value (kv) pairs the mod- els are tasked to match, well as the total sequence length of the prompt. Within each model, we sweep across vary- ing TSS. For each task-model configuration, we compute the ESS and accuracy on a validation set every 10 epochs. We will refer to the entire space of tasks and models across which we sweep as the task-model space. Finally, we split our profiling of ESS into two sections: cross tadk-model analysis (Section 4. and within task-model analysis (Sec- tion 4.2). For more details on the setup, refer to Section D.3.
4.2. Within Task-Model Analysis
Next, to further establish ESS proxy for memory uti- lization, we study how ESS evolves as a function of MQAR performance in a regime where TSS is kept fixed and, therefore, does not correlate with accuracy. We do this by analyzing ESS-accuracy correlation on a per-model, per- task basis over the course of training, uncovering several insights that serve as the basis for our subsequent analysis.
For less memory-intensive tasks trained using models with high TSS, we observe a lower correlation be- tween ESS and performance compared to more memory- intensive tasks trained using a lower TSS (Figure 4).
This in line with the interpretation of ESS as a measure of memory utilization. For easier tasks that are learned by model with high memory capacity, the model is not in- centivized to increase its memory utilization beyond where it resides at initialization. In contrast, for difficult tasks that operate in a conttry-constrained regime, the model is forced to increase its memory utilization in order to learn, resulting in strong positive correlations between accuracy and ESS over training. 4 Digging bit deeper, that this form of ESS analysis reveals two failure modes of model learning in recurrent frameworks (which recall have bounded TSS): state saturation and state collapse.
4.1. Cross Task-Model Analysis
sOC.
find
ESS 1599/kr
netall BSS (tesel Tis)
Figure 5. ESS/Kv TSS/kv as proxy for model performance as measured by correlation.
'In Figure 4, the empty spot in the WLA grid corresponds to NaN from the entropy-ESS computation. The empty spots in the SA grid correspond to MQAR task constraints discussed Section D.3.
Our first goal is to understand how ESS empirically cap- tures memory utilization by studying its correlation with post-training MQAR performance across the entire task- model space. To appropriately analyze ESS across tasks, we normalize it by the memory demands of MQAR, con-
6

---

tion refers to the scenario in which a model has insufficient TSS to fully learn task, resulting in its ESS con- verging near its TSS. This is reflected in its ESS/TSS (which we refer to state Quantifying Memory Utilization with Effective State-Size
TSS (student) 16
TSS (student)
TSS (student)
TSS (student)
128
101a
Distillation Loss
10-2
engression
2
2
2
2
2
2
2²
ESS
(teacceer)
ESS
[teacher)
ESS
(teacher)
ESS (teacber)
Figure 6. ESS-distillation loss (activation) correlation.
State satura-
isting performance-effiiieee frontier in sequence models. We partition our results based on the stage of model train- ing at which we apply ESS analysis: initialization-phar (Section 5.1), post-training (Section 5.2), and mid mid-training (Section A). Finally, we move beyond synthetics and ex- tend ESS to language in Section 5.3.
1.0-
accurasy
FSS 7T33
0.5
0..5
0.0
0.0
5.1. Initialization Phase Analysis
Figure 7. Accuracy and state utilization as function of kv for low and high TSS mod- cls.
ESS TSS Soalaz
GLA
GLAA-8 (TS8 206)
utilization) residing near 1. We observe this in Figure 7 where we note that models with TSS of 8 perform worse as the task difficulty scales due to a saturated state. State collapse, on the other hand, refers to the scenario in which a model has sufficient TSS to learn (or partially learn) a task, but its ESS fails to increase during training. resulting in a heavily underutilized state. With respect to state collapse, we observe the following:
(a)
(b)
(c)
Figure 8. (a) ESS-TSS scaling in the S6, GLA and GLA-S6 fea-
turizers. (b) ESS and accuracy on MQAR as a function of TSS
GLA. (c) ESS and accuracy on MQAR as a function of normal- ization factor for initialization in GLA-S6.
Initialization in weight space plays crucial role in ma- chine learning, significantly impacting model convergence and training stability (Glorot & Bengio, 2010). We extend this concept to the initialization of recurrent models in state space, leaning on the intuition from Figure 3 that suggests higher ESS can enhance performance. Namely, we illus- trate how ESS at initialization can be used inform featur- izer selection - the selection of the function that maps the imput to the operator T = fu) or equivalently the recurrent features (A,( (,,),, (u.)) D2 (w,)))EE and ini- tialization schemes. In doing so, we uncover design flaws of a prominent model, S6 (Mamba) (Gu & Dao, 2024).
For GLA and WLA. state collapse occurs in the high kv bucket of task-mode) space (i.e. kv 27) whereas for LA it does not (Figure 7). More generally, we find that LA has higher state utilization than GLA and WLA.
While state saturation can only be solved by increasing TSS, state collapse can in principle be solved by increas- ing ESS. Unlike TSS, which is a fixed hyperparameter of the model, one can modulate ESS by changing various as- pects of the model pipeline. Furthermore, even outside of the state collapse regime, given the positive correlation be- tween ESS and performance across the task-model space, increasing ESS is generally viable approach to improv- ing model performance without sacrificing efficiency. We explore this idea in the results to follow.
ESS-informed featurizer selection. To study the rela- tionship between memory capacity and memory utiliza- tion in S6, we remove the short convolutional layer in the Mamba block and stack two of these modified blocks between SwiGLUs (Shazeer, 2020). Under the default MQAR task settings outlined in (Poli et al.) (see Tables 2, 3, and 4 for details), we observe that S6 is entirely unable to learn MQAR (accuracy R 0) across multiple scales of TSS (16 256) as shown in Figure 27. This aligns with results from (Yang et al., b), which independently demonstrate the
5. Applications of Effective State-Size
In Section 4, we showed that changes in ESS are correlated with changes in performance, both across models and dur- ing model training, indicating its importance beyond just interpretability. In this section, we aim to push this insight further by understanding how we can leverage ESS to both improve upon and improve our understanding of the ex-
7

---

increase the Quantifying Memory Utilization with Effective State-Size
poor performance of S6 without the additional short convo- lutional layer on different in-context recall task. To inves- tigate the cause, we look into how S6 is preconditioned to utilize its memory by computing its ESS when processing a Gaussian noise input, prior to training.
analysis helps identify potentially weak featurization and initialization schemes, enabling us to pinpoint shortcom- ings in the S6 featurizer and implement a straightforward fix.
Figure 8a demonstrates that the ESS of S6 layers at ini- tialization scales poorly with respect to TSS, notably fail- ing to increase monotonically. In contrast, GLA lay- ers (Yang et al., a), configured with hyperparameters to match the TSS, model width, number of layers, and hidden-state normalization of the S6 model (see Sec- tion D.2 and Table 2). exhibit greater and monotoni- cally increasing ESS-TSS scaling at initialization (Figure 8a). Despite the architectural similarities between the S6 and GLA layers, Figure 8b demonstrates that, unlike S6, GLA achieves accuracy improvements that correlate with increases both TSS and ESS.
5.2. ESS-Informed Model-Order Reduction
Model-order reduction refers to the process of improving model efficiency by reducing state-size while retaining per- formance. Previous works, such as (Massaroli et al., 2023), have explored the distillation of linear time-invariant (LTI) operators (Ti; Ti+k.i+k) into linear recurrences with small state-sizes using backpropagation. Other techniques for model-order reduction such as modal truncation and balanced truncation (Beliczynski et al., 1992; Gawronski & Juang, 1990) are also applicable to LTIs. In this study, however, we are concerned with improving the efficiency of general LIVs. Since ESS serves as a lower bound for the minimally realizable TSS (Section 3), we postulate that ESS can be used as heuristic for conducting model-order reduction.
Based on these findings, we conjecture that the poor ESS- TSS scaling of S6 prevents the model from effectively uti- lizing all of its states, irrespective of increases in memory capacity.
To test this, we distill GLA layers (with TSS = 256) of various models trained on different MQAR task regimes, to see how the ESS of the original model layer (i.e. the teacher) influences its ability to be distilled into a smaller layer (i.e. the student). We apply the technique outlined in (Bick et al., 2024), where the process can be divided into two steps. 1) matching the operators (min(||T(s)
ESS-informed initialization scheme. To further investi- gate the differences between the aforementioned S6 model and GLA model, we construct a composite model termed GLA-S6. This model adopts the fcature-sharing structure of GLA (dividing dimensions into heads and sharing com- putations within a head), but applies the S6 featurization to the matrix as follows:
To /l ( Il?))) and 2) matching the output activations (min (Ilv(s) - Mo (11/ll//)))))) More details can be found in Section D.6.
GLA:
A diag( ding(sigmoiddWuu)
(4)
GLA-S6:
A = diag exp ( - |1/a 2/a
n/a]" T
Figure 6 plots ESS of each of the teacher model along with its activation distillation loss across student models of dif- ferent sizes (TSS).
softplus(Wui))
(5)
A clear trend emerges, higher teacher ESS correlates with greater distillation loss throughout all student model sizes tested.
Like S6. GLA-S6 fails to learn MQAR across the same range of TSS (see Figure 27) and exhibits poor initialization-ESS scaling as shown in Figure 8a. Upon fur- ther inspection, we identify the cause of poor ESS seal- ing: with each new state introduced, the arange term (1 2
This finding also aligns with downstream MQAR evalua- tion scores shown in Figure 31, where distilling the layer with higher teacher ESS results in greater performance degradation. Moreover, as shown in Figure 32, student ESS can be compared against teacher ESS to provide addi- tional feedback on the effectiveness of the distillation pro- cess itself. These findings position ESS as useful heuris- tic for predicting model compressibility, enabling efficient estimation of the potential for state-size reduction without extensive experimentation.
n)) exponentially pushes new entries of A towards zero, negating the effects of additional states de- spite the increase in TSS. Therefore, to ameliorate the poor ESS scaling, we propose simple solution: normalization factor.
By scaling the normalization factor (a), Figure 8c shows that GLA-S6 achieves improvements in MQAR accu- racy post-training, reflecting the impact of increasing its initialization-ESS. despite the models having identical memory capacities.
5.3. State Modulation of Large Language Models
In contrast to synthetic tasks like MQAR, selective copy- ing, and compression, we find that strong recall perfor- mance on language depends not only on a model having
These experiments demonstrate the efficacy of analyzing ESS at initialization, as they reveal how different models are preconditioned to utilize their working memory. This
8

---

Quantifying Memory Utilization with Effective State-Size
The Eileet of Separslor Tokena on: the ESS of Falson Moniia 7B
Perplexity Cap
petiod
Standard
Rigraa Recal
(Meal) 238%
0-
LA GLA
W.A
SA
1000
2000
2000
4000
Bigran Recall Distance
(b)
Figure 9. (a) The effect of separator tokens over Falcon Mamba 7B. See Section E.5 for plots of other open-weight models. (b) Compar-
ison standard perplexity and bigram recall perplexity (Arora et al., 2024a).
sufficient ESS, but also on its ability to dynamically mod- ulate its ESS in response to imputs. We demonstrate that this explains why linear attention, though effective on syn- thetic experiments (Section 4), is widely known to perform poorly on more complex language tasks (Katharopoulos et al., 2020; Arora et al.. 2024b). We begin by evaluating the (total ESS); of open-weight pre-trained models. Bor- rowing notation from Section 3.1, (total ESS); is defined as follows:
though standard perplexities (computed over a subset of the FineWeb dataset (Penedo et al., 2024)) are similar across SA. WLA, and GLA, significant differences emerge when considering the bigram recall perplexity metric introduced by Arora et al. (2024a).
The ability of a model to recall information, measured by bigram recall perplexity across pre-training dataset (rather than within narrow task space), reveals a perfor- mance hierarchy that closely mirrors the observed state modulation capabilities.
(total ESS);
mb EE E aE(d, [E(m) YE(M
Σ
ES
This finding suggests that state modulation serves as key mechanism enabling models to effectively manage com- plex context dependencies, directly impacting their perfor- mance on recall heavy tasks.
Since we do not marginalize along the sequence dimension, will examine this measure qualitatively by observing function of the sequence index. Our
how it changes analysis shown in Figure 9a (and more broadly in Section E.5) reveals an intriguing phenomenon: ESS undergoes a noticeable dip whenever an end-of-speech (EOS) token is encountered (refer to Section D.8 for experimental details). This behavior aligns with our intuition regarding the role of EOS tokens and provides quantitative measure of how ef- fectively a model can 'reset' or 'forget' past contexts when transitioning between distinct segments of text. To inves- tigate these effects in more controlled environment, we trained four 1B parameter models (LA, WLA, GLA, and SA as described in Section 4) under identical conditions (see Table 9).
6. Conclusion
In this work, we propose effective state-size (ESS), mea- sure of memory utilization in sequence models derived using dynamical systems theory. We motivate this met- ric as a valuable tool for analyzing memory utilization in LIVs by demonstrating its strong correlation with perfor- mance across wide range of synthetic tasks. In doing so, we find that ESS offers versatile framework for un- derstanding both the performance and efficiency of causal sequence models. Leveraging these insights, we are able to construct novel, ESS-informed initializers, regularizers, and distillation strategies that improve beyond the existing performance-fficieneyy trade-offs found in recurrent mod- els. Finally, we extend the ESS framework to language tasks, introducing the idea of state modulation a concept that proves crucial for performance on bigram recall tasks. Overall, this work establishes ESS as a foundational tool for understanding and improving sequence model perfor- mance, opening new avenues for optimizing memory uti- lization and, more generally, model efficiency.
We observe a clear hierarchy in the degree of state mod- ulation, which can be summarized as follows:
SA GLA WLA > LA (Figure 42).
SA exhibits the most pronounced state modulation, begin- ning at a tolerance level of 1e-2, while also realizing the largest ESS. GLA follows, with modulation emerging a tolerance of 1c-1. WLA shows minimal modulation, only detectable at tolerance of 1.0, while LA displays no dis- cernible state modulation in response to separator tokens, demonstrating clear lack of ability to modulate ESS. The importance of state modulation becomes apparent when ex- amining model performance. Figure 9b illustrates that al-
9

---

Quantifying Memory Utilization with Effective State-Size
Impact Statement
Arora, S., Byuboglu, S., Zhang, M., Timalsina, A., Al- berti, S., Zou, J., Rudra, A., and Re, C. Simple linear attention language models balance the recall- throughput tradeoff. In ICLR 2024 Workshop on Math- ematical and Empirical Understanding of Foundation Models, 2024b. URL https:: / forum? wwwmmmdddddohhhheeeee
This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of our work, none which we feel must be specifically highlighted here.
opwnrevieeeee..eeeeeeeeeeee.eewe nnt// (pages 9, 18).
References
Ba, J., Hinton, G., Mnih, V., Leibo, J. Z., and lonescu, C. Using fast weights to attend to the recent past. In Proceedings of the 30th International Conference on Neural Information Processing Systems, NIPS' 16, pp. 4338-4346 Red Hook, NY, USA, 2016. Curran Asso- ciates Inc. ISBN 9781510838819, (page 18).
Abnar, S. and Zuidema, W. Quantifying attention flow in transformers. In Jurafsky, D., Chai, J., Schluter, N., and Tetreault, J. (eds.), Proceedings the 58th Anmual Meet- ing of the Association for Computational Linguistics, pp. 4190-45 4197, Online, July 2020. Association for Compu- tational Linguisties. doi: 10. 11118665/112222.ac-maaaaaa 385. URL https: http:://acaaatthloogooooo org/2020. acl-main.20 385/. (pages 1, 18).
Beliczynski, B., Kale, I., and Cain, G. D. Approximation of fir by iir digital filters: An algorithm based on balanced model reduction. IEEE Transactions on Signal Process- ing. 40(3):532-542, 1992. (page 8).
Akaike, H. Stochastic theory of minimal realization.
IEEE Transactions on Automatic Control, 19(6):667-
674, 1974. doi: 10.1 1A.11///AA.1994.1100000000
(page 3).
Bhattamishra, S., Patel, A., Blunsom, P., and Kanade, V. Understanding in-context learning in transformers and LLMs by learning to learn discrete functions. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https: / openreview. net/ forum? www.iiddeeeeyeeeeeeeee (page 18).
Akyürck, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: architectures and algorithms. In Pro- ceedings of the 41st International Conference on Ma- chine Learning, IC, JMLR. 2024. (page 18).
Ali, A.,
Zimerman, I., and Wolf, L. The hidden attention of mamba models, 2024. URL https: /tt//aaxvvv/////////////////aavm org/ abs/2403. 01590. (pages 1, 18).
Bhojanapalli, S., Yun, C., Rawat, A. S., Reddi, S., and Kumar, S. Low-rank bottleneck in multi-head attention models. In Proceedings of the 37th International Confer- ence on Machine Learning, ICML'20. JMLR.org, 2020. (pages 1, 18).
Allen-Zhu, Z. and Li, Y. Physics of language models: Part 3.3, knowledge capacity scaling laws. In The Thirteenth International Conference on Learning Representations, 2025. URL https: wttp//pwwweeeeeeee/eeee wwwp:nntttffommme id id-FxNNiUgtfa (page 2).
Bick, A., Li, K., Xing, E. P., Kolter, J. Z., and Gu, A. Transformers to SSMs: Distilling quadratic knowledge to subquadratic models. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024. URL https: / openreview net/ forum? id=FJ1rSZBMCD, (page 8).
Alzahrani, N., Alyahya, H., Alnumay, Y., AlRashed, S.,
Alsubaie, S., Almushayqih, Y., Mirza, F, Alotaibi, N., Al-Twairesh, N., Alowisheq, A., Bari, M. S., and Khan, H. When benchmarks are targets: Revealing the sensitiv- ity of large language model leaderboards. In Ku, L.-W., Martins, A., and Srikumar, V. (eds.), Proceedings of the 62nd Annual Meeting of the Association for Computa- tional Linguistics (Volume I: Long Papers). pp. 13787- 13805, Bangkok, Thailand, August 2024. Association for Computational Linguistics. doi: 10. 10.18653/v1/2044 acl-long.744 744. URL https:/ / /awllaaaannooooooooeeoa org prg/ acl - long. 744/. (page 19).
Cammarata, N., Carter, S., Goh, G., Olah, C., Petrov, M., Schubert, L., Voss, C., Egan, B., and Lim, S. K. Thread: Circuits. Disuill, 2020. doi: 10.2391 5/distill..0024 (page 18).
Chen, C.-T. Linear System Theory and Design. Oxford University Press, Inc., USA, 3rd edition, 1998. ISBN 0195117778. (page 1).
Dao, T. and Gu, A. Transformers are ssms: generalized models and efficient algorithms through structured state space duality. In Proceedings of the 41st International Conference on Machine Learning, ICML'24. JMLR. 2024. (pages 1, 5, 17, 18, 25, 26).
Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., and Re, C. Zoology: Measuring and improving recall in efficient language models. In The Twelfth International Conference on Learning Rep- resentations, 2024a. URL https:/ openreview. (pages 9, 9, 18, 19,
net/
De, S., Smith, S. L., Fernando, A., Botev, A., Cristian- Muraru, G., Gu, A., Haroun, R., Berrada, L.,, Chen, Y.,
19, 26, 26, 26).
10

---

Quantifying Memory Utilization with Effective State-Size
Srinivasan, S., Desjardins, G., Doucet, A., Budden, D., Teh, Y. W., Pascanu, R., Freitas, N. D., and Gulcehre, C. Griffin: Mixing gated linear recurrences with lo- cal attention for efficient language models, 2024. URL https: / http://wwwwvv../aaaee/ abs/ 2402. 9427. (page 52).
In Teh, Y. W. and Titterington, M. (eds.), Proceed- ings of the Thirteenth International Conference Ar- tificial Intelligence and Statistics, volume 9 of Pro- ceedings of Machine Learning Research, Pp. 249- 256, Chia Laguna Resort, Sardinia, Italy, 13-15 May 2010. PMLR. URL https: http:///poooeeieee/eeeeoe mml.. press/v9/ peess//9/gloootol//
html. (page 7).
DeWilde, and van der Veen, A. Time- Varying Sys- tems and Computations. Springer US, 1998. ISBN 9780792381891. URL postts::/bboot google. co. 21).
Gu, A. and Dao, T. Mamba: Linear-time sequence modeling with selective state spaces. In First Con- ference on Language Modeling, 2024. URL https: / //openreview. net forum? (pages 1, 2, 17, 18, 25, 26).
(pages I, 3,
http/forrnm---Etssskk1F Pid-tEYskw1V22.
Dong, Y., Cordonnier, J.-B., and Loukas, A. Attention is not all you need: pure attention loses rank doubly expo- nentially with depth. In Meila, M. and Zhang. T. (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Ma- chine Learning Research, pp. 2793-2803, PMLR, 18- 24 Jul 2021. URL https: http///ppceeee....eee// mlr. press /v//s//119/// 1131///o/9111111 html. (page 18).
Gu, A., Goel, K., and Re, C. Efficiently mod- eling long sequences with structured state spaces. International Conference on Learning Representa- tions, 2022a. URL https: forum?id=uYLFozi https://openeev/eeee http:////pprreee..ee...m.. net/ (pages 1, 2, 17, 22, 25).
Gu, A., Gupta, A., Goel, K., and Ré, C. On the parameter- ization and initialization of diagonal state space models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA, 2022b. Curran Associates Inc. ISBN 97817130 011113871066 1088. (page 17).
Dupont, E., Doucet, A., and Teh, Y. W. Augmented neural ODEs. Curran Associates Inc., Red Hook, NY, USA, 2019. (page 18).
Elhage, N., Nanda, N., Olsson, C., Henighan, T., Joseph, N., Mann, B., Askell, A., Bai, Y. Chen, A., Conerly, T., DasSarma, N., Drain, D., Ganguli, D., Hatfield- Dodds, Z., Hernandez, D., Jones, A., Kernion, J., Lovitt, L., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. A math- ematical framework for transformer circuits. Trans- former Circuits Thread, 2021. https//transffrmmr- siwiiitspuub22211ffaa
Hasani, R., Lechner, M., Wang. T.-H., Chahine, M., Amini, A., and Rus, D. Liquid struc- tural state-space models. In The Eleventh In- ternational Conference Learning Representations, 2023. URL https: /ttpp:://penreiieer openreview net/ forum?
id=gAOTKKKES7R, (pages 1,5, 18).
Hendrycks, D., Burns, C., Basart, S., Zou, A., Mazeika, M.. Song, D., and Steinhardt. J. Measuring massive multitask language understanding. In In- ternational Conference on Learning Representations, 2021. URL https: wpwp/ppeeeee..eee////e hwtpvviewwwet/fooo...oe forum? id=d7KB jmI (page 19).
html. (page 18).
Fu, D. Y., Dao, T., Saab, K. K., Thomas, A. W., Rudra, A., and Re, C. Hungry hungry hippos: Towards language modeling with state space models. In The Eleventh International Conference on Learning Representations, 2023. URL https:/ wwwprewiwwwweeeeeeeep netooruumm SesCOzDDOYYGG (page 18).
Ho, B. L. and Kalman, R. E. Effective construction of linear state-variable models from input/output functions. at Automaatsierunnsttehhiii, 14(1-12):545-548, 1966. doi: doi: doi:10..1524aatto116666 14.122544. (page 2).
Gawronski, W. and Juang, J.-N. Model reduction in lim- ited time and frequency intervals. International Jour- nal of Systems Science, 21( 1990. doi: 10. 00.10000000772299010006666 URL https::/p///........ /doi. (page 8).
Hochreiter, S. and Schmidhuber, J. Long short-term memory. Neural Comput., 98)) nov 1997. ISSN 0899-7667. doi: 10.1162/neco.199777
http:/10.100000020010000
1735. URL ht https:
Glorioso, P., Anthony, Q. Tokpanov, Y., Whittington, J., Pilault, J., Ibrahim, A., and Millidge, B. Zamba: A compact 7b ssm hybrid model, 2024. URL https: (page 52).
1 1997. 9. 8. 1735. (page 2).
Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, F.,, Cai, T., Rutherford, E., de Las Casas, D.,, Hendricks, L. A., Welbl, J., Clark, A., Hennigan, T., Noland, E., Millican, K., van den Driessche, G., Damoc, B., Guy, A., Osindero, S., Simonyan, K., Elsen, E.,, Vinyals, O., Rae,
//arxiv. //aaxiv.org/abl......aaaa http:/aab22220.1..7.7
Glorot, X. and Bengio. Y. Understanding the diffi- culty of training deep feedforward neural networks.

---

Quantifying Memory Utilization with Effective State-Size
J. W., and Sifre, L. Training compute-optimal large lan- guage models. In Proceedings of the 36th International Conference on Neural Information Processing Systems, NIPS '22, Red Hook, NY, USA. 2022. Curran Asso- ciates Inc. ISBN 9781713871088. (page 3).
Min, Z. and Li, Z. On the efficiency of transformers: The effect of attention rank, 2024. URL https: //operreviww wwweenrrveww.et/forreee (pages 1, 18).
Olsson, C., Elhage, N., Nanda, N., Joseph, N., Das- Sarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Gan- guli, D., Hatfield-Dodds, Z., Hernandez, D., John- ston, S., Jones, A., Kernion, J., Lovitt, I.., Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads. Transformer Circuits Thread, 2022a. hutps//transtoomrrrcicuuttstst pttp::/iiiiiiiiuuu////////oooooo hate:aaddd-nuuuttioo-eddddniiiii htmll. (pages I, 18).
Katharopoulos, A., Vyas, A., Pappas, N., and Fleuret, F. Transformers are mas: fast autoregressive transformers with linear attention. In Proceedings of the 37th Inter- national Conference Machine Learning, ICML'20. JMg,.org, 2020. (pages 6. 9, 17, 18, 25).
Kiranyaz, S., Avci, O., Abdeljaber, O.. Ince, T., Gab- bouj, M., and Inman, J. Id convolutional neural net- works and upplications: A survey, 2019. URL https: //arxiv.or .rrg/ abs / 9955 03554. (page 1).
Olsson, C., Elhage, N., Nanda, N., Joseph, N., DasSarma, N., Henighan, T., Mann, B., Askell, A., Bai, Y., Chen, A., Conerly, T., Drain, D., Ganguli, D., Hatfield-Dodds,
Lieber, O., Lenz, B., Bata, H., Cohen, G., Osin, J.. Dalmedigos, I,, Safahi, E., Meirom, S., Belinkov, Y., Shalev-Shwartz, S., Abend, O., Alon, R., Asida, T., Bergman, A., Glozman, R., Gokhman, M., Manevich, A., Ratner, N., Rozen, N., Shwartz, E., Zusman, M., and Shoham, Y. Jamba: A hybrid transformer-mant language model, 2024. URL https::///// www...ivrrrr abs abs/2403. 19887. (pages 30, 52, 53).
Z., Hermandez, D., Johnston, S., Jones, A., Kernion, J., Lovitt, L.,, Ndousse, K., Amodei, D., Brown, T., Clark, J., Kaplan, J., McCandlish, S., and Olah, C. In-context learning and induction heads. Transformer Circuits Thread, 2022b. https:/ttannsfomercrrruuittt pit.rri...ii..u//22/ime htt..d/tm/dddmmmmmml.m.mmm (page
Lindquist, A. and Picci, G. On the stochastic realization problem. SIAM J. Control Optim., 17(3))365-389 May 1979. ISSN 0363-0129. doi: 10.1137/031702. URL
18).
Oppenheim, A. V., Willsky, A. S., and Nawab, S. H. Sig- nals & systems (2nd ed.). Prentice-Hall, Inc., USA, 1996. ISBN 0138147574, (page l).
http::://o...org3...1137 2).
(page
Orvieto, A., Smith, S. L.., Gu, A., Fernando, A., Gulcehre, C., Pascanu, R., and De. S. Resurrecting recurrent neural networks for long sequences. In International Confer- ence Machine Learning, pp. 26670-26698. PMLR, 2023. (pages 17, 25).
Ljung, L. System Identification: Theory for the User. Prentice Hall information and system sci- ences series. Prentice Hall PTR, 1999. ISBN 9780136566953, URL https: tttes::/booker htt///ookks.oooooo google. com/books?idet ?id=nHF oQqAACAAJ. (page 2).
Parnichkun, R. N., Massaroli, S., Moro, A., Smith, J. T., Hasani, R., Lechner, M., An, Q., Ré, C., Asama, H., Er- mon, S., et al. State-free inference of state-space mod- els: the transfer function approach. In Proceedings of the 41st International Conference Machine Learning. pp. 99834-399868 2024. (pages 1, 17, 25).
Loshchilov, I. and Hutter, F. Decoupled weight decay regu- larization. In International Conference Learning Rep- resentations. (pages 29, 29, 31, 31).
Lynn, P. A. and Fuerst, W. Introductory digital signal pro- cessing with computer applications (revised ed.). John Wiley & Sons, Inc., USA, 1994. ISBN 0471943746. (page 1).
Pascanu, R., Mikolov, T., and Bengio, Y. On the dif- ficulty of training recurrent neural networks. In Das- gupta. S. and McAllester, D. (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Re- search, pp. 1310-1318, Atlanta, Georgia, USA. 17-19 Jun 2013. PMLR. URL https:: http///rcceeeeeeeeeeeooeoccceeee iggs. mir press/ v28 pwpppaussanaaa.m html. (page 2).
Martin, E. and Cundy, C. Parallelizing linear recurrent neu- ral nets over sequence length. In International Confer- ence on Learning Representations, 2018. (page 2).
Massaroli, S., Poli, M., Fu, D., Kumbong, H., Parnichkun, R., Romero, D., Timalsina, A., McIntyre, Q., Chen, B., Rudra, A., et al. Laughing hyena distillery: Extract- ing compact recurrences from convolutions. Advances in Neural Information Processing Systems, 36:17072- 17116, 2023. (pages 8, 17).
Penedo, G., Kydlfček, H., Lozhkov, A., Mitchell, M., Raf- fel, C. A., Von Werra, L., Wolf, T., et al. The fineweb datasets: Decanting the web for the finest text data at
12

---

Quantifying Memory Utilization with Effective State-Size
scale. Advances in Neural Information Processing Sys- tems, 57:30811-30844 2024. (pages 9, 32).
Sun, M., Chen, X., Kolter, Z., and Liu, Z. Massive tivations in large language models. In First Conference Language Modeling. (pages 18).
Poli, M., Thomas, A. W., Nguyen, E., Ponnusamy, P., Deis- eroth, B., Kersting, K., Suzuki, T., Hie, B., Ermon, S., Re, C., et al. Mechanistic design and scaling of hybrid architectures. In Forty-first International Conference Machine Learning. (pages 5, 7, 18, 26).
Sun, Y., Dong, L.,, Huang. S., Ma, S., Xia, Y., Xue, J., Wang, J., and Wei, Retentive network: A successor transformer for large language models, 2023. URL https:/ http/aaxivvvo.// wttp://.rggabb/////// 08621. (pages 5, 17).
Poli, M., Massaroli, S., Nguyen, E., Fu, D. Y., Dao, T., Baccus, S., Bengio, Y., Ermon, S., and Ré, C. Hyena hi- erarchy: Towards larger convolutional language models. In International Conference on Muchine Learning, pp. 28043-28078. PMLR, 2023. (pages 2, 5, 17).
Tsai, Y.-H. H., Bai, S., Yamada, M., Morency, L.-P., and Salakhutdinov, R. Transformer dissection: An unified understanding for transformer's s attention via the lens of kernel. In Inui, K., Jiang, J., Ng, V., and Wan, X. (eds.), Proceedings ofe 2019 Conference on Empirical Methods in Natural Language Processing and the 9th In- ternational Joint Conference on Natural Language Pro- cessing (MMNLP-IJCNLP), pp. 4344 4353, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/w1/D199-1433, htLhttps: (pages 1,
Power, A., Burda, Y., Edwards, H., Babuschkin, I., and Misra, V. Grokking: Generalization beyond overfit- ting on small algorithmic datasets, 2022. URL https: //arxivv... org/ abs 2201. 02177.
(page 18).
Ramsauer, H., Schäfl, B., Lehner, J., Seidl, P., Widrich, M., Gruber, L., Holzleitner, M., Adler, T., Kreil, D., Kopp. M. K. al. Hopfield networks is all you need. In Inter- national Conference on Learning Representations. (page 18).
/ //aclant hwc.lantholloy..org//11e9-----9. 17).
Tumma, N., Lechner, M., Loo, N., Hasani, R., and Rus, D. Leveraging low-rank and sparse recurrent connectivity for robust closed-loop control. In The Twelfth Interna- tional Conference on Learning Representations. (page 18).
Romero, D. W., Kuzina, A., Bekkers, E. J., Tomezak, J. M., and Hoogendoorn, M. Ckconv: Continuous kernel con- volution for sequential data. In International Conference on Learning Representations. (page I).
Vandebril, R., Van Barel, M., and Mastronardi, N. A note the representation and definition of semiseparable matrices. Numerical Linear Algebra with Applications, 12(8):839-888, 2005. doi: https://doi. hrtp://ooooog///..
Roy, O. and Vetterli, M. The effective rank: A measure of effective dimensionality. In 2007 15th European Signul Processing Conference, pp. 606-610, 2007. (pages 3, 4).
org//0../////ooo/ (tmp/10.11000 1002/nla.
455. URL. ht https:::tt........... //ol///on//oioo// onl nne ibrrry.rr.rrrr wi ley. hom/doii/ali / 10. /00. 100. /nla. 455, (page 3).
Rumelhart, D. E. and McClelland, L. Learning Inter- nal Representations by Error Propagation, pp. 318-362. 1987. (page 2).
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, E., and Polosukhin, I. At- tention is all you need. Advances in neural information processing systems, 30, 2017. (pages 1, 1, 2, 6, 17, 18).
Shazeer, N. Glu variants improve transformer, 2020. URL https: http///axxvvvvv.a org/ abs/ 2002. 05202. (page 7).
Vig, J. A multiscale visualization of attention in the trans- former model. In Proceedings of the 57th Annual Meet- ing of the Association for Computational Linguistics: System Demonstrations, pp. 37. Association for Com- putational Linguistics, 2019. (pages 1, 18).
Shen, H. Mutual information scaling and expressive power of sequence models, 2019, URL https: //axivv. org/ 042./11211.......... 04271. (page 18).
Wang, Y., Ma, X., Zhang, G., Ni, Y., Chandra, A., Guo, S., Ren, W., Arulraj, A., He, X., Jiang, Z., et al. Mmlu-pro: A more robust and challenging multi-task language un- derstanding benchmark. The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track. (page 19).
Smith, J. T., Warrington, A., and Linderman, S. Simpli- fied state space layers for sequence modeling. In The Eleventh International Conference on Learning Repre- sentations. (pages 1,5, 17, 22).
Su, J., Ahmed, M., Lu, Y., Pan, S. Bo. W., and Liu, Y. Ro- former: Enhanced transformer with rotary position em- bedding. Neurocomputing, 568:127063, 2024. (pages 25, 26).
Willems, 171-269.
C.
Models for Dynamics, pp. Vieweg+Teubreg Verlag, Wies-
baden, 1989.
ISBN 978-3-322-96657-5 doi:
13

---

Quantifying Memory Utilization with Effective State-Size
10. 10.107799733333329965555553 URL https:// (page 2).
Wu, X., Ajorlou, A., Wang, Y., Jegelka, S., and Jadbabaie, A. On the role of attention masks and layernorm in transformers. In The Thirty-eighth Annual Conference Neural Information Processing Systems. (pages 2, 18).
Xia, Chandrasekaran, S., Gu, M. and Li, X. S. Fast algorithms for hierarchically semiseparable matri- ces. Numerical Linear Algebra with Applications, 17 (6):953-976, 2010. doi: www.eii.rrrrwwer.......... wiley. 691. URL
com/ hom//di//e abs www.///.... 1002 /wa./naaa.... 691. (page 3).
Xiao, G., Tian, Y., Chen, B., Han, S., and Lewis, M. Effi- cient streaming language models with attention sinks. In The Twelfth International Conference on Learning Rep- resentations. (pages 1, 18).
Yang, S., Wang, B., Shen, Y., Panda, R., and Kim, Y. Gated linear attention transformers with hardware- efficient training. In Forty-first International Conference on Machine Learning, a. (pages 1, 5, 8, 17, 18, 29).
Yang, S., Wang, B., Zhang, Y., Shen, Y., and Kim, Y. Par- allelizing linear transformers with the delta rule over se- quence length. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, (page 7).
Zancato, I.., Seshadri, A., Dukler, Y., Golatkar, A. S., Shen, Y., Bowman, B., Trager, M. Achille, A., and Soatto, S. B' mojo: Hybrid state space realizations of founda- tion models with eidetic and fading memory. Advances in Neural Information Processing Systems, 37:130433- 130462, 2024. (pages 1, 2).
Zellers, R., Holtzman, A., Bisk, Y., Farhadi, A., and Choi, Y. HellaSwag: Can a machine really finish Korhonen, A., Traum, D., and
your sentence? Märquez, L. (eds.), Proceedings of the 57th Annual Meeting of the Association for Computational Linguis- fics, pp. 4791-4800, Florence, Italy, July 2019. Associa- tion for Computational Linguistics. doi: 10. 18.55////v/ P19-1472. URL https: https::/accleee h/tps////ocaaaooooo..oooooo org/ P19-1472/. (page 19).
14

---

Quantifying Memory Utilization with Effective State-Size
Supplementary Material
Contents
Introduction
2
Related Work
3
Theory
3
3.1
Computing Effective State-Size
4
4
Empirical Validation of Effective State-Size
5
4.1
Cross Task-Model Analysis..
6
4.2
Within Task-Model Analysis
5
Applications of Effective State-Size
5.1 Initialization-PPe Analysis
5.2
ESS-Informed Model-Order Reduction
5.3
State Modulation of Large Language Models
Conclusion
Mid-Training Analysis
17
Extended Related Work
17
Theoretical Background
19
C.1
Notation
19
C.2 Derivations and Proofs
20
C.2.1
The Operator Realization of Linear Recurrences
20
C.2.2 Factorizing The Operator Realization Submatrix H,
20
C.2.3
The Trivial Recurrence Realization
21
C.2.4
Existence of Infinite Recurrent Realizations (Proof of Theorem 3.1)
21
C.3 Minimal Recurrent Realization (Proof of Theorem 3.2)
22
D Methods
22
D.1 Additional Details on Computing ESS
22
D.1.1 Computing ESS for SISO models
22
15

---

Quantifying Memory Utilization with Effective State-Size
D.1.2 Tolerance entropy ESS
23
D.1.3 PyTorch Implementation
23
D.2 Formulation of the Featurizers
24
D.3 Empirical Validation
26
D.4
ESS-Informed Featurizer Selection and Initialization Scheme
28
D.5 ESS-Informed Regularization
29
D.6 ESS-Informed Model-Order Reduction
29
D.7 ESS Analysis for Hybrid Networks
30
D.8
State Modulation of Large Language Models
31
Extended Experimental Results
32
E.1
Empirical Validation
32
E.1.1
State Collapse Continued
32
E.1.2
Entropy-ESS MQAR Results Continued
33
E.1.3
Tolerance- ESS MQAR Results
40
E.1.4
Selective Copying and Compression Results
46
E.1.5
ESS Training Dynamics in MQAR
47
E.2
Initialization-Phase Analysis
48
E.3
Mid-Training Analysis
49
E,4
Post-Training Analysis
50
E.4.1
Model-Order Reduction
50
E.4.2
Hybridization
................................
51
E.5
State Modulation of Large Language Models
53
E.6
Miscellaneous...
57
E.6.1
Effective State-Size on C++ Code
57
E.6.2
How the Number Prompting Shots Affects the Effective State Size Language Models
58
16

---

Quantifying Memory Utilization with Effective State-Size
A. Mid-Training Analysis
ESS
LA ESS
WLA ESS
GLA
||A||
LA
|
WLA |
WLA GLA -- LA
1.50
0.3
ESS/K //v
1.25
Aille
accuracy
1.00
0.2
0.75
0.1
0.50
0.0
20
27
2
20
200
21
seq len
(b)
Figure 10. (a) ESS/kv and |I A:Il function of sequence length. (b) Accuracy of models as function of ESS-based regularizer strength.
To motivate the idea of increasing ESS mid-training, we revisit the concept of state collapse a phenomenon that arises due to trainability issues in recurrent models (Figure 29), as discussed in Section 4.2. Recall that state collapse describes a failure mode of learning GLA and WLA which, unlike LA, have learnable A, matrices (where i denotes the index along the sequence dimension). To see why this contributes to state collapse, we note that the values of the operator submatrices H, are disproportionately influenced by A, due to the presence of terms in the form of A,-1 ... A1 for each i. Hence, the closer A, lies to the O-matrix, the faster these terms decay, reducing the numerical rank of H;. We demonstrate this empirically in Figure 10a, which shows that for both GLA and WLA, ESS/kv and II, A:ll decrease as function of sequence length. In contrast, for LA, whose A matrix is given by the identity, ESS/Kv remains large as the sequence length grows.
Given this insight, one approach to addressing state collapse in GLA and WLA is pushing the A matrices towards the identity by adding the following term the loss function: A||A Fl where À denotes the strength of the regularizer and I denotes the identity. In doing so, are effectively decaying the model towards LA, increasing its ESS and giving the following:
GLA and WLA trained using the ESS-based regularization scheme described above outperform LA. When trained with- out it, they perform worse than LA (Figure 10b).
For more commentary on this result, please refer to Section E.3.
B. Extended Related Work
Causal sequence models. From classical linear recurrences to modern sequence models like Transformers, a vast array of causal model architectures have emerged (Vaswani et al., 2017: Tsai et al., 2019; Katharopoulos et al., 2020; Poli et al., 2023; Yang et al., a; Gu & Dao, 2024; Dao & Gu, 2024; Sun et al., 2023). recent years, the ability to process sequences in parallel has become increasingly critical, largely due to advancements in hardware accelerators such as GPUs. This need for parallelism likely explains the growing popularity of models like attention, Mamba, and $4.
We observe that all of these models, which support parallelization across the sequence dimension, can be formulated using a linear system representation (y Tuu as detailed in the introductory and theoretical sections (Sections 1 and 3). For this work, we categorize these models into two types: linear systems and systems with input-varying linear operators. The key distinction between these two frameworks is that the former, the operator T in input-invariant models is composed of fixed system parameters, whereas the latter, the parameters are dynamically generated from the input. Linear systems encompass both linear time-narying (LTV) and linear time-invariant (LTI) systems. Although LTV systems have been relatively unexplored in deep learning, several LTI models have been studied (Gu et al., 2022a;b; Smith al.; Orvieto et al., 2023; Parnichkun et al., 2024). Convolutional models use kernels h to construct Hankel matrices H, whose rank corresponds to the minimal state-size of the model (DeWilde & van der Veen, 1998). Massaroli et al. (2023) explored
17

---

Quantifying Memory Utilization with Effective State-Size
methods reduce the order of models by leveraging the Hankel matrix. Notably, the submatrix H, (defined in Theory 3.2) exhibits Hankel structure in LTT models and provides per-sequence-inde information. In this work, however, we do not explore Hankel matrices further, as they are not easily generalizable to LTV systems.
In contrast, systems with input-varying linear operators, which we abbreviate as LIVs, are characterized by an operator T that is dynamically constructed through a featurizer and is defined by T = =u)). Examples of such models include softmax attention (Vaswani et al., 2017), linear attention (Katharopoulos et al., 2020), Liquid-S4 (Hasani et al., 2023), Mamba (Gu & Dao. 2024; Dao & Gu, 2024), and gated linear attention (Yang et al., a). Although these models may appear nonlinear, they can still be represented like a linear system, enabling the application of linear analysis techniques. This forms the basis for the effective state-size metric.
Interpretability. Analysis tools for sequence models can be categorized into two types: extrinsic and intrinsic. Extrinsic tools focus solely on the imput and output, treating the model's internal processes as black boxes. This approach is highly generalizable, it can be applied to any model, including those with non-linear recurrences. A notable example by Shen (2019) uses statistical measures such mutual information to compute metrics that capture model "expressivity". While these methods are versatile and applicable to various datasets, their generality makes them less effective at capturing the inner workings of causal sequence models, which is the primary focus of this work.
Intrinsic tools, conversely, directly visualize the model's internal mechanisms. A recently popular framework known as mechanistic interpretabilityy provides one such example (Power et al., 2022). Mechanistic interpretability involves dissect- ing complex models to understand how specific components contribute to the model's overall behavior (Cammarata al., 2020). Unlike work, mechanistic interpretability does not target the operator view of the model but instead emphasizes the functional roles and interactions of individual model components.
For our purposes, we are primarily concerned with the visualization and analysis of classical and modern causal sequence models through the unifying lens of linear operators. Most analyses of these operators rely on visualization techniques (Olsson et al., 2022a; Vig, 2019; Abnar& & Zuidema, 2020; Ali et al., 2024; Xiao et al.; Sun et al.) to gain insights into the model's s internal processes. Visualizing the operator T is advantageous, as it reveals important features like the formation of induction heads, strong activations, diagonal and block-diagonal patterns, and Toeplitz structures. However, raw visualizations are largely qualitative and oftentimes do not provide the quantitative metries necessary for effectively evaluating model's internal mechanisms gap we aim to address in this work.
Other, more quantitative, intrinsic methods perform some form of spectral analysis on the full operator (Dong et al., 2021; Min & Li, 2024; Tumma et al.; Bhojanapalli et al., 2020). A limitation of these approaches is that they often disregard the causal masking of T, which significantly impacts the model's s rank and singular values (Wu et al.). As a result, the rank of the causal operator T alone lacks a clear interpretation.
The proposed effective state-size metric is a method applicable both linear systems and LIVs. As quantitative proxy for memory utilization, it offers insights into the inner workings of causal sequence models, ensuring generality, usability, and interpretability.
Synthetic and language benchmarks. In this work, we build on synthetic tasks from the mechanistic architecture design (MAD) framework introduced in (Poli al.). MAD defines set of small-scale tasks designed to evaluate key model capabilities, such as in-context recall (Akyürek et al., 2024; Bhattamishra et al., 2024; Elhage et al., 2021; Olsson et al., 2022b). Training models on these tasks is efficient, making them well-suited for exploring a large space of tasks and models, as demonstrated in several prior works (Dupont et al., 2019; Arora et al., 2024b; Fu et al., 2023). In this work, we investigate the effective state-size across a subset of the MAD tasks: multi-query associative recall (MQAR), selective copying, and compression, varying the difficulty of each to gain nuanced understanding of how effective state-size evolves across these task landscapes.
Among the synthetic tasks we examine, MQAR stands out in particular. Proposed by Arora et al. (2024a), MQAR was designed to bridge the gap between synthetic and real language tasks explained by associative recall - the ability of a model to retrieve information based on relationships between different elements in its memory. This capability has long been sought after in the construction of sequence model architectures (Ramsauer et al.; Ba et al., 2016); as such, we evaluate the performance of our models on MQAR to measure the benefits of using effective state-size to iterate on canonical frameworks used in sequence modeling.
18

---

Quantifying Memory Utilization with Effective State-Size
One notable aspect of MQAR observed in Arora al. (2024a) is that the size of the model cache needs scale with the difficulty of the task maintain performance. While this observation holds, work demonstrates that model cache size is an imperfect measure in this context due to the discrepancy between memory capacity, as measured by theoretically realizable state size, and memory utilization, as measured by effective state-size. At higher level, this demonstrates how our work provides a new perspective on analyzing memory-intensive synthetic tasks.
While the MAD framework and synthetic tasks have shown correlations with model performance on large-scale language tasks, language itself poses a unique challenge. Models are tasked with predicting the next token given previous tokens simple yet general objective. New tasks can be created simply by altering the prompts, thereby expanding the range of possible task domains.
Although numerous language evaluation tasks such as those in (Hendrycks et al., 2021; Wang et al.; Zellers et al., 2019) have been proposed, they often probe narrow task space and tend to be brittle. For example, shuffling the order of multiple choices in MMLU can drastically change model rankings (Alzahrani et al., 2024).
Unlike narrow benchmarks, perplexity scores can be computed across an entire pre-training dataset, covering much broader task domain. However, small perplexity gaps between models make it challenging metric for evaluation. Re- cently, Arora al. found that much of the difference in perplexity between models can be attributed to bigram perplexity measure of model's s ability to utilize the context and predict successor token (second token of a bigram) given repeated context token (first token of bigram) within sequence. They demonstrate that most of the average perplex- ity difference between a gated convolution model and an attention model stems from differences in bigram perplexity, suggesting that recall is a key capability for language models.
The effective state-size analysis presented in this work reveals that strong recall performance as measured by bigram perplexity language modeling tasks depends not only on memory capacity, but also model's ability modulate its state-size within a given context.
C. Theoretical Background
C.1. Notation
We adopt the following notation in this paper:
Inputs, outputs, and operators follow flattened notation. i.e., y E Réd and T E Reddxéd In particular, the original inputs and outputs with shape are flattened in row-major ordering, resulting in T having x sub-blocks, each of which is of size d x d.
Tensor subscripts index sequence indices (time-step) and superscripts index channel/hidden dimensions. Lee., for an input E Ed uy E Rd denotes the input vector sequence-index and u T denotes the input vector for channel a. Similarly, Tiş E Rdd denotes the linear weighing of uj to yi.
Indices within square brackets indicate matrix indices void of semantics (sequence index, channels, etc.). Ie., A indexes row a and column 3 of matrix A,-
Semicolons within subscripts denote product over ranges (A133 = A14)
Tensor slices are denoted with colons and are inclusive over the ranges. I.e., uo:2 = Un
19

---

Quantifying Memory Utilization with Effective State-Size
C.2. Derivations and Proofs
C.2.1. THE OPERATOR REALIZATION OF LINEAR RECURRENCES
Unrolling the recurrence in Equation unveils the following formulation:
80
81 Bo
82 Bi + A (Bouo)
53 Bru2 + A2(Biu + Ai (i
8,
A& Byu
(C.2.1)
14 C.
Av Byuj
+ Dzuis
(C.2.2)
Lk=i-1
which corresponds to the operator:
0
i<j
Tij = D:
i=
(C.2.3)
CAI-1 i>
C.2.2. FACTORIZING THE OPERATOR REALIZATION SUBMATRIX H;
Factorizing the strictly lower triangular submatrices the operator (T into causal and anti-causal factors, unveils that n, (i.e. the TSS) upper bounds the dimensionality of the inner product between the factors, and thus, also the rank of the submatrix (n; 2 ra)k(H:)):
C
At-1:1
Bo
T = H
Cr-1
Ar-2:1
Az-2;1
B
I
C
A
Bo
A+11a
[A--1
A-112
A r]
Ct-1
:
B
A
(C.2.4)
Ci CitiAi C422A411
[A--1 Bo
A
Bi-1] A-1
O,C,,
C 1 A8
cousal
d(t 4)
anti-cousal
Besides unveiling the relationship between the rank of the realized operator and the original state-size n., the following insights can be drawn from the decomposition:
The causal portion C, is the input-state projection matrix at time-step (i.e., = C,u.i-)) corresponding to Equation (C.2.1).
ESS (rank(H), (H,)) is simply the minimum rank between the causal and anti-causal projections.
conjunction with Theorem 3.2, we observe that the causally determinable minimal state-size (causal ESS) is
20

---

Quantifying Memory Utilization with Effective State-Size
equivalent the rank of the causal projection. This insight allows us to construct a more efficient realization of the recurrence:
We can minimally factorize the causal projection C, L,R,, where L, E R and R, E R with rank(C,).
-The right factor R; becomes the new input-state projection matrix for H,, effectively reducing the state dimen- sion to the causal ESS.
- and 1 can be determined from R, using the process outlined in Theorem 3.1, and C = C.Li-
C.2.3. THE TRIVIAL RECURRENCE REALIZATION
Any input-varying and input-invariant causal operator can be trivially realized with the following recurrence:
[(a))
((ss)
84+1*
((a))
8t +
(d))
Us
(C.2.5)
y4 [T,0
T2
T 1] 8: Tu
In simple terms, the state 8, stores each input from E - 1], which then mapped to the output with operator features row i. Note that in the case where the operator is input varying, the trivial realization upholds the causality of the featurization process (i.e. the features (A,, B;; C. Di)je( of the trivial realization causally determined). Moreover, the causally determined ESS (see Section C.2.2) for the trivially realized recurrence is equivalent to its TSS, as C; Ias.
C.2.4. EXISTENCE OF INFINITE RECURRENT REALIZATIONS (PROOF OF THEOREM 3.1)
Theorem 3.1 Given any causal input-ing input-invariant operator T, there exist infinite variations of linear recurrences in the form of Equation (1) that realize an equivalent input-output operator.
Proof. We first categorize the operator into two portions: the memoryless portion, where = j. and the dynamical portion, where i j. The memoryless portion can be trivially realized by setting D; T,. For the dynamical portion, we draw inspiration from (DeWilde & van der Veen, 1998) and approach the proof of existence by ansatz. The following steps outline the proof:
1. Section C.2.2 demonstrates that, given linear recurrence the form of Equation (1), the operator submatrix can be factorized into causal and anti-causal parts, where the causal part represents the input-state projection matrix. Therefore, we proceed by making the following ansatz: for any operator submatrix T = H, H, can be arbitrarily factorized into O, E (d((-i) and C₁ Xdi and that C, represents the input-state projection at time-step (i.e., 8, = 1a
2. Construct the dynamic features (A,, Bd, Ci)se(t) such that the assumption above holds. Note that we additionally assume the initial and final states to be 0 without loss of generality, therefore the realization of Co, Ag, A and B could be ignored.
(a) Set C Ojp--11 obtain (C:)se(I.e). as given the assumptions above, the first set of rows of O linearly projects s; onto yr D which is identical to C in Equation (1).
(b) Set B;-1 C.J.-4:) to obtain (B.)vete-1). .). for which the identity can be obtained by deconstructing the input-state projection matrix C, and equating its assumed state s; with Equation (1).
8; A + Bi-14--1
(C.2.6)
C= + Cjja. Cita-diUUi11
(c) Using the same state-dynamics equation, we could equate the assumed stoje-projection matrices with each other,
21

---

Quantifying Memory Utilization with Effective State-Size
obtaining (At)uE(1,e-11-
8/+1 A,s; + B,u
C4114: A,C,u;-1 + C+11;:-d:44
(C.2.7)
1d A,C,u;i-1
A; = C1-11::--11(++
3. Verify that the realized recurrence maps back to the original operator Tij. proving that arbitrary factorizations (of which there are infinite variations) of the operator submatrices can be used to construct equivalent operators.
Tij = C:A-- As+1 = Oijd- Ospd-1 -
= -)
(C.2.8)
= Oipd-1]Gip.a 1)a-1 = Hildd- (iaaa-ijj(((+)aa--- = -Tg.
As an example, H; can factorized minimally with SVD as follows:
= = (U (U,) (UG= Vo))
where UG) E Rmxr, (r)) E R'Xr, Vir) E R are the r-truncated SVD decompositions, and r rank of H2 E mmxx Theorem 3.2 demonstrates that these factors realizes recurrence with minimal state-size.
C.3. Minimal Recurrent Realization (Proof of Theorem 3.2)
Theorem 3.2 The rank of the operator submatrix (H; = 1 determines the minimal state size required represent
the causal operation (y- Tuu as recurrence.
Proof. The proof of Theorem 3.1 shows that the operator submatrices H; can be decomposed arbitrarily into two state- projection matrices, O and C, whose inner product dimension defines the state size of its recurrent realization at sequence index i. By the rank-nullity theorem, rank(H,) represents the minimum inner product dimension of any such state- projection matrices and thus corresponds to the minimally realizable state size of the operator T at sequence index
D. Methods
D.I. Additional Details on Computing ESS
D.1.1. COMPUTING ESS FOR SISO MODELS
In this section, provide additional details the distinction between computing ESS in single-inputsingll-uupppp (SISO) models like S4 and multi-input multi-output (MIMO) models like S5.
Recall in Sections 3 and C.I, we introduced flattened notation, as it offers a general framework for formulating wide range of operators and recurrences. Namely, MIMO layer like S5 (Smith et al.), which mixes both the channels and sequence simultaneously, can be formulated as = Tu (with the operator realization outlined C.2. the same way a SISO layer like S4 (Gu al., 2022a) can, which only mixes the sequence. The difference between these two models lies in the structure of T: for models that only mix the sequence, such as $4, Tıj is diagonal; for S5, is dense.
Note that since all of the models in our experiments are SISO, we can compute the ESS independently for each channel using the standard operator formulation Tea E éxxt where a E [d]. This approach is significantly more efficient than computing ESS for the multi-channel (flattened) representation. Furthermore, in the case of attention layers, the computation can further reduced to only the h independent heads, as the operator (i.e. the attention matrix) is shared across channels within the same head.
22

---

storch.no_gras Quantifying Memory Utilization with Effective State-Size
D.1.2. TOLERANCE VS ENTROPY ESS
Regarding the distinction between the entropy and tolerance-based forms of ESS, we that entropy-ESS is a valuable summary metric because its computation is independent of any specific tolerance value chosen. However, it potentially be mislcading when comparing ESS across sequence indices due the unequal normalization applied the singular values. Conversely, when comparing entropy-ESS across different operators, it can be useful as the normalization removes the effect of the norm of the operator. Moreover, in contrast to the tolerance-based metric which is discrete, entropy-ESS can assume continuous values ranging from to -
In most of our experiments, we observe consistent trends between entropy-ESS, and tolerance- ESS when the metrics are marginalized over the sequence length. Therefore, unless stated otherwise, our figures presented using the entropy-ESS. In cases where we require ES$ comparison across the sequence dimension, instead plot ESS for multiple tolerance values. Finally, note that for the analyses presented in Section 4, when computing ESS, we average across samples (batch-size), For the rest of the ESS analyses, we average across 32 samples.
D.1.3. PYTORCH IMPLEMENTATION
Below, we provide a PyTorch implementation of various ESS metrics and helper functions that were leveraged in our analyses:
import torch
def T2H_i (T, i, d=1):
Extract i from
Args:
T: Flattened operator with shape [..., dxL, d+L].
i: Index of (I_i) to retrieve.
10
Block size for multi-channel flattened operator representation (default is 1)
11
Returns:
Hi: submatrix of the operator at index 1.
14
15
return TI. d+i;
16
18
def 1258 (T, d-1):
19
***
20
Converts an operator into list of singular values (Ss).
Args:
23
Flattened operator with shape [..., d+L, d+L]
24
d: Block size for multi-channel flattoned operator representation (default is 1)
Returns:
27
Ss: 1ist singular values for each sequence index
28
seqien T.size
Ss
[]
for range (1, seqlen)
32
Hi (T, i, d)
33
i, torch. svd (H
34
Ss.append (S
return Ss
36
37
Shorch.no_gradi.
38
def $:2TolcranceEs5((5, tol=ic-4):
Computes the tolerance-RSS from the list of singular values.
41
23

---

Quantifying Memory Utilization with Effective State-Size
42
Args:
43
Ss: List of singular values.
44
tol: Telerance value.
Returns:
47
tolerance-ESS
48
19
ranks = []
for SV in 3s:
51
rank terch, sum (SV>=tol, dim -1) appeed (rank)
52
53
ranks = torch.stack stack (ranks, dim--1)
return ranks
56
@torch. no
no_rad)) grad()
57
def Ss2EntropyESS (Ss, clip=le-12):
58
Computes the entropy-E55 from the list of singular values.
60
61
Args:
62
Ss: List of singular values.
clip: clips probabilities below this value avoiding numerical instabilitie when the probabilities too numerically close to 0.
64
65
Returns:
66
entropy-KSS
ranks []
69
for SV Ss:
70
p =
[..., None]
= torch. clip (p, clip)
-torch, succh.sumt (p torch.leg dim=-1)
73
rank terch. (H)
74
ranks. append
15
ranks = terch. stack (ranks, dim--1)
return ranks
Example usage (Python-pseudocode):
>>> out model (u, attput_attenticns-Trree
>>> shape: [bs, layers, heads, len, len]
>>> T out. out.attention_matz
>>> Ss T28s (T) # List of singular values
>>> # ESS shape [bs, layers, heads, len-1]
ESS Ss2TcleranceESS (Ss, tol=le-3)
>>> mean
torch, mean (ESS)
We note that calculating the effective rank may cause numerical instability when p. approaches 0 due to the logarithmic term. This is partially mitigated by clipping the normalized singular values, as shown above.
D.2. Formulation of the Featurizers
In this section, we first establish the equivalence between linear attention and state-space models, then proceed with for- mulating the LIV' discussed in this paper.
Linear attention and state-space model equivalence. We begin by demonstrating that linear attention models state- space models, serving the foundation for the subsequent formulation of featurizers for other models, such gated linear attention and weighted linear attention.
A single linear attention head with dimension d/h, typically formulated as
y = qk² v,
(D.2.1)
24

---

Quantifying Memory Utilization with Effective State-Size
in which q, u Rexd/h input features. They can be reformulated recurrences with matrix- states 8 Rd/hxd/h follows (Katharopoulos et al., 2020):
8, 8:-1 11-1 + k
(D.2.2)
y: q 8/
Without loss of generality, applying column-major flattening to the matrix-valued state and treating et the input u; the recurrence can be formulated like Equation (I), by setting A; [(d/n)a and:
hh
q
q
d/h
4
Bi-1
Ci
(D.2.3)
0
q
d/h
d/h
Notice that each individual channel forms single-input-sinll--uure (SISO) recurrence (like many of the SSM architec- tures including S4, S6, Mamba2, and more [27; 26; 16; 47; 48)), as there is no mixing across channels. Additionally, each of these SISO recurrences has state-size of d/h.
Formulation of the featurizers. For the sake of completeness, we additionally characterize the "values" feature in attention-like models, with fu( as follows:
8/+1 A4s; + (.fu(u;)
(D.2.4)
V/ C. 8; + Difa(us)
The following lists the formulations of the recurrent featurizers based on the results above, along with their per-channel (i.e. average) TSS and total TSS:
Linear Attention (LA):
A I, Bh Ro)) C (o))
fu(us) = W uc.
(D.2.5)
where C: W. W E Rd/hxd; d and h represent the number of channels and heads, respectively. A is fixed identity matrix. Per-channel TSS n. = d/h, and the total TSS is d² /h. Each channel c [d]] is grouped into heads, where the head index corresponding to the channel is given by A = [ch/d], and within a head, all corresponding projection matrices (W&, W etc.) are weight tied (shared). Moreover, a rotational positional encoding (RoPE) is by default applied to the and projections (Su et al., 2024).
Gated Linear Attention (GLA):
A = digg(sigmoid A, u, 1/8, B W
(D.2.6)
C² Whur, fa(u:) = W
where besides having projections identical to those LA, WA, E IR1xdd and W Rd/hx16 Like LA,, per- channel TSS is nţ = d/h, and the total TSS is d² d//h. By default, 3 is set to 16.
25

---

Quantifying Memory Utilization with Effective State-Size
Weighted Linear Attention (WLA))
A diag (oig(middd(A"" A4)2/3))
(D...7)
_ WB,, C Wou, fu(u;) = Wa
where Wc. WB, and Wi are identical to those LA, and A E Rd/h explicitly parameterized and initialized to 0. Like LA, per-channel TSS n d/h, and the total TSS is d? /h.
Softmax Attention (SA);
B RoPE(W
O RoPP(WW
(D.2.8)
T softmax (aaimaxxxC((88 (Bk))) fu(u,) W
where Wc. WB. and Wa are identical to those in LA, and can converted into a recurrence using the trivial realization in Equation C.2.5. Therefore, the per-channel TSS i, and total TSS is id. Like LA, a rotational positional encoding (RoPE) is by default applied to the B and Ö projections (Su al., 2024).
S6 (Gu & Dao, 2024):
A softplus W4 + b))
A diag(exp( (osp(exp(-AA)))) AA*),,
(D.2.9)
B AWW, C2 Wou,
is the channel index, We, Wa Rxdd and Wg Riddd Here, per-channel TSS is and total TSS is nd. Note that by setting d/h, S6 resembles GLA, with the following exceptions:
where Å E T is initialized to [1
2
S6 has only one (not h) different projection matrices for B and C.
S6 has channel-wise projections for the input-varying discretization applied to B and C.
S6 has an explicitly parameterized vector valued Á.
S6 has some minor differences in the non-linearity applied to keep 0 < A < 1. Similar to GLA, S6 also has diagonal A matrices, whereas Mamba2 (Dao Gu, 2024), the A matrix is scalar-valued. However, the B and C projections Mamba2 more closely resemble that of GLA.
GLA-S6:
A = diag(exp(-111 (11//a 2/a n/a]" o softplus(W3), WA, uz))),
(D.2.10)
B Whu,, Cº wau, fu(u.) W u4
GLA-S6 is a combination of S6 and GLA such that the B and C projections are identical to that of GLA, while A is featurized similarly to S6. Namely, it has identical WB. Wc, and W,, those found in LAA which means that the per-channel TSS is d/h h and total TSS is /h. The A matrix is featurized with the aragge term ([1 2 ... nn] T like S6. We additionally added a normalization hyperparameter a, which controls the rate at which elements of A decay to O.
D.3. Empirical Validation
Here, provide details the task-model sweep presented in Section 4. Table lists the hyperparameters that were exhaustively swept across to generate the task-model space. Note that the hyperparameter controlling the task difficulty is task-dependent (for more details, see Poli et al.).
For the MQAR and selective copying tasks, a default vocab size of 8192 (Arora et al., 2024a) was used for all models. For the compression tasks, the vocab size was varied modulate task difficulty, as shown in Table Any other task settings specified here are defaulted to those presented Arora al. (2024a). Two important constraints on the tasks from Arora et al. (2024a) which we also utilize in experiments follows: MQAR task requires that
num pairs S seq len
26

---

Quantifying Memory Utilization with Effective State-Size
and the selective copying task requires that
2 * num tokens to copy + 1 < seq len
Any of the task configurations from Table I that violate these conditions were not trained. This is why the SA plot in Figure has empty spots in the grid.
Finally, we note that all architectures analyzed here consist of 4 layers: 2 sequence mixing layers (i.e one of GLA, LA, WLA or SA) and 2 channel mixing layers (i.e. MLPs).
Configuration
Value(s)
Tasks
MQAR, selective copying, compression
Num. key-value pairs
8, 16, 32, 64, 128 8, 16, 32, 64, 128
Num. tokens to copy
Vocab size (compression)
16, 32, 64, 128
Vocab size (MQAR and selective copying)
8192
Sequence length
64, 128, 256, 512, 1024, 2048 GLA, LA, WLA, SA
Model (featurizer)
Model width
64, 128, 256, 512
Number of heads
4, 8
Optimizer
AdamW
Learning Rate
0.002
Weight Decay
0.1
Batch Size
64
Epochs
70
Steps Per Epoch
2000
Num. Training Samples
128k
Num. Testing Samples
6.4k
Table I. Set of hyperparameters for the task model sweep.
Regarding the post-hoc analysis performed on the sweep, we note the following:
model width Since the average TSS computed over the channels (which equals number of heads for GLA, LA, and WLA) explains more meaningful variation with respect to memory utilization than model width and number of heads individually, we consolidate those two dimensions into one by analyzing across the average TSS axis. For SA, since average TSS is a function of the task rather than model hyperparameters (see Equation C.2.5 and Section D.2), we instead compute the of TSS over all d channels, given by the total TSS per layer d In any cases where the average/total qualifier is not specified, note that are referring the average ESS TSS.
Since we analyze the recurrent models across the average TSS dimension, we compute average ESS in the plots presented in Scction 4.1 in order to compare ESS and TSS as proxies for performance. Similarly, since we analyze the SA models across the total TSS dimension, we compute total ESS for those plots. However, we note that plots for both the average/total ESS and TSS presented in Section E.I.
When we marginalize across dimensions, we average across all models in that bucket of task-model space. For example, in Figure 4, for each (TSS, kv) pair, we average over the correlations of all models that correspond to that pair. Note, however, that we never average across tasks (i.e. MQAR, selective copying, compression) or featurizers (i.e. GLA, LA, WLA, SA).
27

---

Quantifying Memory Utilization with Effective State-Size
When compute cross-model correlations (Figure 3) for SA, we filter models which have accuracy 0.95. This is done in order observe meaningful variation function of (total ESSy/v and (total / since many of the SA models obtain an accuracy of I.
When compute within-model correlations (Figure for MQAR, drop epoch 0 from the computation since observe phase at the start of training in which ESS tends to decrease, but accuracy does not change. We elaborate on this phenomenon in Section E.1 and hope to characterize it further in future work.
Regarding the task-adjusted forms of ESS and TSS which, in the case of MQAR, are computed by normalizing the raw ESS value by the number of kv-pairs in the task, we note that this normalization factor is critical for observing the cross task-model correlations presented in Figure In particular, in Figure 1 11, we find that correlations across the task-model space break down when examining the unnormalized ESS, This points to the higher-level notion that ESS is expected scale with the memory demands of the task.
We interpret the state utilization of model, which given by ESS/TSS, as proxy for what portion of the memory capacity of the network is realized in practice. By definition, state utilization takes on values ranging continuously from 0 to 1. Recall that state utilization near is indicative of state saturation.
While for most of the ESS analysis conducted on the sweep the entropy-ESS, note that for the state utilization plot presented in Figure we find that entropy-ESS fails to capture the state collapse phenomenon. This is because state collapse is primarily dictated by the magnitude of the singular values, as opposed to the relative decay rate of the entire spectrum. In particular, if all of the singular values are close to 0, the layer is likely failing to learn an expressive state, resulting in poor performance. Due to the normalization applied to the spectrum, the entropy-ESS metric may potentially present this state as having high effective rank; however, in practice, we know that this misrepresentation of the true dynamics. Tolerance-ESS, in contrast, appropriately captures the dynamics of the state with respect to the norm of the operator. Because of this, whenever we analyze ESS as it pertains state collapse (e.g. Figure 10a), present the tolerance-ESS instead.
the tolerance-ESS with tolerance level set at 1e-3. We do this because
D.4. ESS-Informed Featurizer Selection and Initialization Scheme
Configuration
Value
Model width
128
Num. heads
8
Configuration
Value
a arngge Norm. (a)
1000
Logit Norm. (1)) -exxpnnsoon
16
Model width
128
1
State expansion (d state)
16
Table 2. Default GLA hyperparameters.
Table 3. Default S6 hyperparameters.
"For GLA-S6.
$ K "e-epansioon is used to vary TSS the featurizer experi- ments.
28

---

Quantifying Memory Utilization with Effective State-Size
Configuration
Value
Sequence length
2048
Num. KV Pairs
128
KV Dist. Const.
0.1
Optimizer
Adam W
Learning Rate
0.002
Weight Decay
0.1
Batch Size
Epochs
Steps Per Epoch
2000
Num. Training Samples
128k
Num. Testing Samples
6.4k
Vocabulary Size
8192
Table Default MQAR task settings employed throughout the
featurizer and initialization experiments in Section 5.1.
"(Loshchilov & Hutter)
D.5. ESS-Informed Regularization
We use the following MQAR configuration for the regularization experiments presented in Section A.
Configuration
Value
Sequence length
4096
Num. KV Pairs
128
KV Dist. Const.
0.1
Optimizer
Adam W
Learning Rate
0.002
Weight Decay
0.1
Batch Size
64
Epochs
70
Steps Per Epoch
2000
Num. Training Samples
128k
Num. Testing Samples
6.4k
Vocabulary Size
8192
Model width
128
Num. heads
Table MQAR settings and model hyperparameters em-
ployed throughout the mid- mid-training experiments in Section A.
"(Loshchilov & Hutter)
Regarding the regularization scheme itself, since examine models with two sequence mixing layers, we explore the following strategies: regularizing both layers, only regularizing the first layer and only regularizing the second layer. Empirically, we find that only regularizing the second layer performs the best and is thus the result presented in Figure 10b. We claborate on why this is the most successful strategy in Section E.3.
D.6. ESS-Informed Model-Order Reduction
The teacher models used in the distillation experiments are 2-layer GLA models (Yang et al., a) with dimension 1288 and TSS = 256 (num 8 and expand.k= = 16). We checkpointed the models every 10 epochs while training on MQAR
29

---

Quantifying Memory Utilization with Effective State-Size
across different task difficulties. The task ranges are given follows:
Sequence length: [512, 1024, 2048)
Number of Key-Value Pairs: [64, 128]
Other settings follow the defaults shown in Table 4. For each task difficulty pair, we repeated the training run with three different seeds. For each teacher model checkpoint, both layers were distilled independently with student models of different state-sizes (16, 32, 64, and 128). Distillation settings are shown in Table 6.
The ESS metric in Figures 6, and 32 was computed by taking the minimum across the batch (32 samples) and channels (128 dimensions), evaluated at the mid-point of the sequence (f/2). Using the mid-point of the sequence statistic was done to reduce compute. The midpoint in particular was chosen as it is the point in the sequence at which H; has the largest dimensions, retaining the largest amount of information from the original operator. Note that in practice, where a large task space isn't being tested to validate our approach, a more thorough computation of ESS across sequence length is feasible. Other marginalization approaches such as taking the maximum or average across the batch and channels also show similar trends, but we found that taking the minimum demonstrates the trend most clearly.
summary
Configuration
Value
Optimizer
AdamW
Batch Size
Learning Rate
0.001
Weight Decay
0.0
Training Steps (Operator)
800
Dropout (Operator)
0.2
Training Steps (Activation)
3200
Dropout (Activation)
0.2
Table 6. Distillation settings used for the results presented in Section 5.2.
D.7. ESS Analysis for Hybrid Networks
In our ESS analysis applied to hybrid networks, we restrict our scope to GLA-SA hybrids. In particular, we explore the following two settings:
8-layer hybrid networks which 4 layers are sequence mixers (i.e. one of GLA or SA) and 4 layers are channel mixers (i.e. MLPs). We exhaust all possible hybrid networks (of which there are 16) and perform post-training, per- layer ESS analysis on the networks. We train these hybrid models on MQAR with task-model settings given below in Table 7.
16-layer hybrid networks in which 8 layers are sequence mixers (i.e. one of GLA or SA) and 8 layers are channel mixers (i.e. MLPs). Here, we explore all combinations of hybrid networks that follow the Jamba hybridization policy (Lieber et al., 2024) and perform post-training, lay-layer ESS analysis on the networks. We train these hybrid models settings given below in Table
on MQAR with
30

---

Quantifying Memory Utilization with Effective State-Size
Configuration
Value
Sequence length
2048
Num. KV Pairs
512
KV Dist. Const.
0.1
Optimizer
Adam W
Learning Rate
0.002
Weight Decay
0.1
Batch Size
Epochs
Steps Per Epoch
2000
Num. Training Samples
128k
Num. Testing Samples
6.4k
Vocabulary Size
8192
Model width
64
Num. heads
Table 7. Default MQAR task settings employed throughout the hybridization experiments conducted in the first setting described above.
"(Loshchilov & Hutter)
Configuration
Value
Sequence length
4096
Num. KV Pairs
1024
KV Dist. Const.
0.1
Optimizer
AdamW
Learning Rate
0.002
Weight Decay
0.1
Batch Size
64
Epochs
70
Steps Per Epoch
2000
Num. Training Samples
128k
Num. Testing Samples
6.4k
Vocabulary Size
8192
Model width
16
Num. heads
Table 8. Default MQAR task settings employed throughout the hybridization experiments conducted in the second setting de- scribed above.
"(Loshchilov & Hutter)
Results for these experiments can be found in Section E.4.2.
D.8. State Modulation of Large Language Models
State modulation of weenhweight models. The following randomly generated sentences were used to study the effects of separator tokens state modulation in open- beinht pre-trained language models.
<bos>>angoes are rich in vitamin and can be blended into a refreshing smoothie<sep> Giraffes are the tallest mammals on Earth due their long necks and legs<sep> She collects vintage typewriters from the 1940s<sep>
31

---

Quantifying Memory Utilization with Effective State-Size
Jupiter's Great Red Spot is a giant storm that has been raging for hundreds of yearscsep>
State modulation on custom-trained IB models. For our custom-trained IB language models, we used longer sen- tences, as state modulation patterns were less discernible with shorter sequences. collection of randomly generated sentences shown below:
<bos>The deep blue ocean, teeming with an extraordinary array of marine from the smallest plankton to the largest whales, stretches out infinitely towards the horizon, a vast and mysterious expanse that has captivated the imaginations of explorers, scienfists, and poets for centuries, hiding within its depths secrets yet be discovered and stories yet to be told<sep> In a bustling city where skyscrapers tower over narrow streets filled with the constant hum of cars and the chatter of pedestrians, a small café, nestled between two imposing buildings, offers a quiet refuge for those seeking a moment of peace, with the comforting aroma of freshly brewed coffee and the soft sound of jazz music playing in the background, creating a cozy ambiance that feels like a world away from the urban chaos outside<sepp oak tree, with its gnarled branches stretching wide and its thick, sturdy trunk standing firm against the passage of time, has witnessed generations of families grow, seasons change, and countless stories unfold beneath its expansive canopy, becoming a silent guardian of the park, offering shade those who seek solace and sense of continuity in rapidly changing world<sep>
The ancient
We note that the specific sentences and their order are not crucial this analysis. Similar patterns have emerged with various sentence arrangements, provided the sentences sufficiently long.
Training settings are outlined in Table 9.
Configuration
Value
Batch Size
16
Max Sequence Length
32k
Training Steps
160k
Optimizer
Adam
Learning Rate
0.001
Weight Decay
0.1
Num. Layers
24
Dimension
2048
Table 9. 1B LLM settings.
The perplexity scores shown in Figure 9b were computed on 16k randomly sampled sequences over the FineWeb (Penedo al., 2024) dataset. The raw perplexity samples were smoothed via kernel density estimation method.
E. Extended Experimental Results
E.1. Empirical Validation
In this section, we provide additional results and commentary from the sweep detailed in Section D.3 that were not pre- sented in the main portion of the paper. One thing to note is that most of the ESS results presented in Section 4 were computed using the entropy-ESS, However, we also computed ESS using the tolerance-based approach affirm that both forms of ESS showcase similar trends. In particular, examined tolerances of le-1, 1e-3 and Ie-5. Since we observe similar trends across tolerances, we provide plots for tolerance of 1e-3 below and omit the others for the sake of brevity.
E. STATE COLLAPSE CONTINUED
Here, we continue our discussion on the state collapse phenomenon presented in Section 4.2. In particular, while we assert that state collapse is observable across all TSS in the high kv bucket for GLA/WLA, Figure 7 shows that accuracy differences between LA and GLA/WLA are only evident in the high TSS/high kv bucket of the tosk- space. This is because state saturation is acting as confounder, worsening performance in LA (see Figure 7 when TSS 8). Therefore,
32

---

Quantifying Memory Utilization with Effective State-Size
although collapse in GLA/WLA does not result in worse performance than LA in this specific task-model setting, remains issue even for models with smaller states when trained sufficienlly difficult tasks. This the motivation behind the task-model setting explored in Section A.
E. 1.2. ENTROPY-ESS MQAR RESULTS CONTINURD
(a)
(b)
WIA
T.A
WILA
ESSyyler
(c)
(d)
5
(e)
(f)
60
(h)
Figure 11. (a) TSS/kv vs accuracy across featurizers. This demonstrates that TSS/kv (i.c. memory capacity) is worse proxy for model performance than ESS/kv as discussed in Section 4. (b) (total TSSykv accuracy across featurizers. This demonstrates that (total TSS)/kv is worse proxy for model performance than (total ESSy/v. (c) ESS/Kv vs accuracy across featurizers. (d) (total ESS)/kv vs accuracy across featurizers. (e) ESS/TSS (i.e. state utilization) accuracy across featurizers. We note that models that saturate their state tend to perform worse on the task, which evidence of the state saturation phenomenon discussed in Section 4.2. The models that do not saturate their state but still perform poorly are the models that undergo state collapse. (f) (total ESSY/(ttal TSS) accuracy across featurizers. (g) ESS vs accuracy across featurizers. Note that without normalizing by kv (i.e. the task memory), the correlation with accuracy breaks down substantially. (h) TSS vs accuracy across featurizers.
33

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
54
(c)
(d)
-1-11
-4-51
Sinue (GLA/LA/71A))
WILA
2.10
11000
11.00
AANISS (831
0.30
1000
(g)
IITIee (BLAILA/MMA))
(CLA CELAAAAA.00 S...//TT..
ACTEAAA (8A)
65-
0.00
030
(h)
Figure 12. MQAR accuracies marginalized across different dimensions.
34

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
(c)
(d)
038 10 NVEA
C.S. -
1388 (OL (OLA/AA (900.
- 27.1 1ire -ae S
CLA
WIAA
(985 (8A)
(5A)
0.6
0.30
1..00
22
(g)
1880 1e (SLAAAAAL))))))
ESS (CLA/ (LA/WA)))))
0.0
15
(h)
(i)
Figure 13. MQAR ESS/KV marginalized across different dimensions.
35

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
WTA
VEA 106
(c)
15 111 an Somee -
LA
WLA
(s6al ESSO (S-0))
(
(8)
(lead ESSi (GLA/AA//BA))
(GLA/IA (91AA)
(boak CO00)
(h)
(i)
Figure 14. MQAR (total ESSy/v marginalized across different dimensions.
36

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
(c)
(d)
18857188 (GLA/A (9110)
1S8-TTS (GLA/LA (WLA)
- 21.1 -10 -4- -90-11
CLA
WILA
ESS TSS (58.
301
(g)
138/138 (RLA (RLATA//IAA AA/WIA
O1.A
1080 188 (OLA)) (GLAA/////A (LA/WAA)
S4
E380 1305 (8-))
00.1
(8.4)
0.3
(h)
Figure 15. MQAR ESS/TSS marginalized across different dimensions.
37

---

Quantifying Memory Utilization with Effective State-Size
200
SAFTTE
Buid
(a)
(b)
1501
TOAA 1200
Te4 190
(c)
(d)
-ph-11 - -16
10100
2H4 ES0 MIM 1300 204
CLA
WLAA
-0.001
0
0,022
32-
(g)
Daa 080 (NAA 7000 014400AA
Desul COCA VIAA TAA 30A4
Dend EOO E35
1000
0.000
(h)
(i)
Figure 16. MQAR (total ESS)/(toal TSS) marginalized across different dimensions.
38

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
(C)
(d)
ESS Me ent (BLA (LA/WAA))
out WLLAAA//A
-1 -
ESS-Ar ouE (8A)
GLA
WIA
SIA
11.25
0.70
0.00
10:00
DS5
135-40 oen. BLA/AA/VLLA
ES8-m CBLAAAA (W1.))
(84)
10g
(h)
Figure 17. MQAR ESS- accuracy correlations computed over training marginalized across different dimensions.
39

---

Quantifying Memory Utilization with Effective State-Size
1.3. TOLERANCE-ESS MQAR RESULTS
Below are plots from the MQAR sweep using tolerance-ESS (tol=le-3) instead of entropy-ESS, We note that all of the prevailing trends remain the same.
GLA
LA
WLA
GLA
LA
WILA
(888 Ft
588 kmx
[ooa ES8//k
Note EsSy/tr
(butal (S8)
(a)
(b)
GLA
WLA
CLA
LA
WLA
2-1
2-1
2-
PSS 73SS
(bota) (to0/ 198
(real PSSj/teal (S3)
(bital B0S (oonal T903
(c)
(d)
GLA
1.A
WLA
0.5
ESS
ESS
(c)
Figure 18. Accuracy vs various forms of tolerance-ESS across task-model space. Plots are entirely analogous to those shown in Figure
40

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
(c)
(d)
CLA
WIAA
(g)
F98/99 (OLA)) A.A.ANAA (ILA
8389 tee (CLA (NLA))
1005 Aer (54)
(943
50--
65
(h)
Figure 19. MQAR ESS/KV marginalized across different dimensions.

---

Quantifying Memory Utilization with Effective State-Size
(a)
(
GLA
9A
(c)
(d)
- - 2111
(b6ea ESS0 20 (B.))
CLA
LA
WLA
5A
-00
20-
-00
(g)
(lead ESSyle (GLA)AA/9RA;;
Ostad 1859 (CLLA) A,A, (4144
(boaa (0s0) ter (34)
(344)
2.3
6.--
(h)
(i)
Figure 20. MQAR (total ESSy/v marginalized across different dimensions.
42

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
1e4 190
(c)
(d)
188/138 (CLA) AOLLAAAAWWLL (WLA)
21.1 - -
WIAA
E395 (T5S5 6583
1E K1
200 000139
FSS
138/719
8.56
301
(g)
188/188 (GLA (OLA/AA/WAA)
6:25
138/ (8-3)
E380 155 (9A)
6:00
(h)
Figure 21. MQAR ESS/TSS marginalized across different dimensions.
43

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
E000 3301
TAAe 1200
120
Tea 108
(c)
(d)
NOAL re0y en NIAA/MAA
2-44 D00 100 (844 244
CLA
WIA
0.0072
000000
0.55
(g)
010000
0:70
10000
020
0.50
0.0000
0.20
(h)
Figure 22, MQAR (total ESS)/(toal TSS) marginalized across different dimensions.
44

---

Quantifying Memory Utilization with Effective State-Size
(a)
(b)
120
THe 100
(c)
(d)
ESS- ent (OLARA//WI))
ESS me out (OLA IGLA (LAA/A/
- - -1
ESS 18-1
DSS AVC ACCE 1932)
GLA
WIA
SAA
+0.0
000
-02
138-4m our. (QLA/TA/VLLY (VLA)
ES8-m ALAA
OXT. 184)
C.0
-00
(h)
(i)
Figure 23. MQAR ESS-ane accuracy correlations computed over training marginalized across different dimensions.
45

---

Quantifying Memory Utilization with Effective State-Size
E. SELECTIVE COPYING AND COMPRESSION RESULTS
Below, we present results for the selective copying and compression tasks, analogous to the ones presented in Section on MQAR.
GLA
LA
WLA
SA
GLA
LA
WIA
SA
B.--
Iste
(bonal (880
foea 1oS0
(botal 18S5) pake
(a)
(b)
GLA
GA
WIA
GLA
LA
WIL
158/
(sood 1851
Josd 1881/ate
(sctal 135//an
(boaal 1ii
(c)
(I)
Figure 24. Selective copying results. Note that ESS here refers to entropy-ESS and we abbreviate num, tokens to copy as ntc in plots above. (a) ESS/nte vs accuracy across featurizers. (b) (total ESSYnte vs accuracy across featurizers. (c) TSS/nte vs accuracy across featurizers, (d) (total yutte vs accuracy across featurizers. (e) ESS- accuracy correlation computed over the course of training in (TSS, kv) buckets. (f) ESS-accuracy correlation computed over the course training in (total TSS, kv) buckets.
1800
(8)
GLA
LA
WIL
GLA
758/AM
T88)
(onal TSSSAA
Steed TSRL/v
7755
GLA
WIA
SA
GLA
WLA
tocah
Picah 1bx
Touah dtc
(c)
Figure 25. Compression results. Note that ESS here refers to entropy-ESS and we abbreviate vocab size as vs in plots above. (a) ESS/vs accuracy across featurizers. (b) (total ESS)/vs accuracy across featurizers. (c) TSS/vs accuracy across featurizers. (d) (total TSS)/vs vs accuracy across featurizers. (c) FSS-accuracy correlation computed over the course of training in (TSS, kv) huckets. (f)
ESS acc-accu correlation computed over the course of training in (total TSS, kv) buckets.
We note that with respect the cross task-model trends, find that in both selective copying and compression, task-
46

---

Quantifying Memory Utilization with Effective State-Size
adjusted ESS is better proxy for model performance than task-adjusted TSS (Figures 24a, 24c, 25a, 25c). This is substantial demonstrates the utility of the ESS metric beyond just MQAR.
Regarding within task-model trends, we observe similar patterns for selective copying those seen in MQAR (Figure 24e), with one notable distinction. Namely, ESS and accuracy are positively correlated across a larger portion of the task-model space selective copying than in MQAR. For compression, however, the within task-model trends look a bit different from what we observe in selective copying and MQAR (Figure 25e). One potential reason for this is that the compression task significantly more difficult than the MQAR and selective copying tasks (as noted by the lower accuracies in Figure 25a), leading to more instabilities over the course of training. But any case, this does highlight the fact that the strength of ESS proxy for model performance changes function of the task. The precise nature of this relationship is something we hope to explore in future work.
E. 1.5. ESS TRAINING DYNAMICS IN MQAR
As mentioned Section D.3, we observe a phase at the start of training in MQAR in which ESS tends to decrease. This shown in Figure in which we select an arbitrary task-model configuration from the sweep and plot its ESS and accuracy over the course of training on a per featurizer basis.
Normalized Entropy-ESS
-X- Normalized Tolerance-ESS
Accuracy
SA
GLA
LA
WLA
1.0
1.0
1.0
1.0
Metric value
Metric value
Metric value
Metric value
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
50
50
50
50
Epoch
Epoch
Epoch
Epoch
Figure Training dynamics of select models (dmodel=256, heads=8) trained on MQAR (seqlen=2048, kv=64). We min-max normalize the ESS curves over the course of training to emphasize the shape of the curve opposed to its magnitude. Note that the tolerance-ESS shown here computed using a tolerance of Ic-3.
We find that at the start of training (i.e. in between epochs and 10), even the accuracy not evolving, the ESS is. In particular, the recurrent frameworks (GLA, LA and WLA), we note a sharp decrease in the ESS before begins rise later in training (and along with it the model accuracy). In contrast, in SA observe the opposite: sharp increase the start of training followed by a steady decrease (even after it has solved the task). This points level of nuance in the training dynamics of MQAR ESS that we have yet to characterize and is something we hope to explore in future work.
47

---

Quantifying Memory Utilization with Effective State-Size
E.2. Initialization-Phan Analysis
val/acc
val/loss
modeldd_saal state: model.d. state: modeldd state: 64 model.d. state: 32
staate: 128 sttatet modeld__satee 16
modeldd_stte: model.d_state 128
model.d state: model.d state:
model.d state:
0.0002
8.37
1.5e-4
8.36
8.35
0.0001
5e-5
train/steps
8.33
train/steps
20k
40k
60k
80k
100k
120k 140k
20k
40k
60k
80k
100k
120k
140k
(a) Validation accuracy of S6
(b) Validation loss of S6
val/acc
val/loss
TSS-256
TS5-128
75S-64
TSS-256
TSS-128
TSS-64
TSS-32
TSS-16
TSS-16
0.0003
8.38
2.5e-4
8.37
0.0002
8.36
1.5e4
0.0001
8.34
5e-5
8.33
train/steps
train/steps
20k
40k
60k
80k
100k
120k
20k
40k
60k
80k
100k
120k
(c) Validation accuracy of GLA-S6
(d) Validation loss of GLA-S6
Figure 27. Loss curves of S6 and GLA-S6 showing that the models unable to improve beyond random guessing MQAR, across various state-sizes.
GLA
25
*
0.75
Accuracy
ESS
0.50
24
0.25
Acc. (Mean)
Act. (Min-Max)
0.00
25
27
TSS
Figure and MQAR accuracy as function TSS custom task regime (sequence length 1024, num. pairs 256). This figure illustrates strong correlation between MQAR accuracy, ESS and TSS.
48

---

Quantifying Memory Utilization with Effective State-Size
E.3. Mid-Training Analysis
Normalized Entropy-ESS - Normalized Tolerance-ESS
Accuracy
SA
GLA
LA
WLA
Metric value
1.0
Metric value
1.0
Metric value
1.0
Metric value
1.0
0.5
0.5
0.5
0.5
0.0
0.0
0.0
0.0
0
50
50
50
50
Epoch
Epoch
Epoch
Epoch
Figure 29, An example of the training dynamics of ESS in select models (dmodel=512. heads=4) trained MQAR (seqlen=2048. kv=128) that undergo state collapse (i.e GLA and WLA). We min-max normalize the ESS curves over the course of training to emphasize the shape of the curve as opposed to its magnitude. Note that the tolerance-FSS shown here is computed using tolerance of lc-3.
First, we provide some additional commentary on the ESS-based regularization results discussed in Section A. Recall we showed that decaying the A matrices in GLA and WLA towards the identity matrix enables these models to outperform LA in the state collapse regime. Our intuition for this result is that by ameliorating state collapse, GLA and WLA can better leverage their increased expressivity, which stems from their learnable A matrices feature absent from LA.
Next, mentioned in Section D.5, provide some intuition behind the efficacy of regularizing only the second layer of the network as opposed to the first or both layers.
K-GAA Layer 0- WLA Layer
* GLA Layer * WLA Layer
x
0.8
ESS kk
0.5
0.2
26
27
2
2
21
seq
len
Figure 30. Per-layer ESS/Kv as a function of MQAR sequence length for the GLA and WLA featurizers. ESS shown here is computed using a tolerance of 1e-3. Layers are O-indexed.
Using O-indexing for the layers, Figure 30 shows that layer realizes lower ESS/kv than layer 0, particularly in the case of WLA. This suggests that layer contributes disproportionately the observed state collapse (Figure 29); consequently, it makes sense that layer would need be regularized more heavily. Now, this begs the question the second layer leads to better performance than regularizing both layers (results of which were not shown). We have two possible hypotheses for this outcome. First, introducing regularization terms for both layers may complicate optimization by creating potentially conflicting objectives. Second, excessive decay of the A matrices towards the identity matrix may cause the model to revert to the LA regime, which as shown in Figure 10b performs worse than GLA and WLA (when sufficiently regularized). Nonetheless, we hope to further explore this intuition and investigate other ESS-based forms of regularization in future work.
why only regularizing
49

---

Quantifying Memory Utilization with Effective State-Size
E.4. Post-Training Analysis
E.4.1. MODEL-ORDER REDUCTION
teacher: 156, Ac9.=0.90
teacher: TSS=256, A 0.90
student: TSS-128, Acc.=0.91
student: TSS=128, Acc. 0.88
student: TSS=64, Acc.=0.87
student: TSS=64, Acc.=0.74
student: TSS=32, Acc.=0.75
student: TSS=32, Acc.=0.37
student: TSS=16, A0c.=0.73
student: TSS=16, Acc.=0.05
Layer 0
Layer 1
80
Mean Effective-Stan Seive--Sttt-Sattt Statee-iize Size
60
40
20
128
256
640
768
896
0
128
256
384
512
640
768
Figure This figure compares MQAR accuracy and ESS across reduction scales for layers and 1. lower ESS in layer 0 of the teacher model leads to better downstream performance after distillation compared to distilling layer 1.
TSS (student) = 16
TSS (student) = 32 TSS (student) = 64
TSS (student) = 128
26
Equal ESS
ESS (student)
2
2²
21
23
25
21
23
25
2
23
2
21
23
2
ESS (teacher)
ESS (teacher)
ESS (teacher)
ESS (teacher)
Figure 32. Teacher ESS vs distilled student ESS. As expected, we observe clear trend: an increase in the student TSS results in the student's ESS more closely matching the teacher's ESS. Plots like these can help provide additional context during the distillation process.
50

---

Quantifying Memory Utilization with Effective State-Size
E.4.2. HYBRIDIZATION
SA Layers
GLA Layers
SA Layers: () Accuracy: 0,0
SA Layers: (0.) Accuracy: 0,0
SA Layers: (1.) Accuracy: 0.98
Layers: (2.)
Accuracy: 1.0
SA ESS
200
GLA ESS
SA ESS
200
GLA ESS
200
200
ESS
SA
2
100
100
100
100
2
2
3
Layer Index
Layer Index
Layer Index
Layer Index
SA Layers: (3,)
SA Layers: (0. 1) Accuracy: 1.0
SA Layers: (0, 2) Accuracy: 0.98
SA Layers: (0, 3)
Accuracy: 1.0
Accuracy: 0.96
300
300
3.00
300
SA ESS
200
GLA ESS
SA ESS
200
GLA ESS
SAESS
200
GLA ESS
SAS
200
ESS
2
100
1
0
2
0
1
2
0
2
Layer
Layer
Layer Index
Layer
SA Layers: (1, 2) Accuracy: 0.37
SA Layers: (1. 3)
SA Layers: (2, 3) Accuracy:
SA Layers: (0, 1, 2)
Accuracy: 0.97
Accuracy: 0.99
300
300
300
300
SS
GLA ESS
ESS
200
GLA ESS
SA ESS
GLA ESS
ESS
200
GLA ESS
100
100
100
100
1
Layer Index
Layer Index
Layer eX
Layer Iindx
SA Layers: (0, 1, 3) Accuracy: 1.0
SA Layers: (0, 2, 3) Accuracy: 0.97
SA Layers: (1, 2. 3) Accuracy: 1.0
SA Layers: (0, 1. 2,
Accuracy: 1.0
300
300
300
300
GLA ESS
ESS
GLA ESS
ESS
100
100
100
100
3
1
2
3
Layer Index
Layer Index
Layer Ind EX
Layer Index
Layers
Layers
SA
Accuracy: 0.0
Accuracy, 0.0
Accurscy: 0.6
Acturacy.0.3
Accurscy: 1.0
000
600
50.555
GLA.S55
CL4215
200
202
200
Lager idde
Layer 11:000
L Endinn
(b)
Figure 33. All results presented here are computed using tolerance-based ESS with a tolerance set at 1c-1. Network layers are 0-indexed. (a) Per-layer ESS of all possible 4-yeyer GLA-SA hybrid networks. Experimental settings can be found Section D.7. (b) Per-layer ESS of all possible 8-layer GLA-SA SA Jamba-ine camba-inspired hybrid networks. Experimental settings can be found in Section D.7.
51

---

Quantifying Memory Utilization with Effective State-Size
-acuuaayyy Avg SA
Max SA ESS
accuracy Avg GLA
1.0
1.0
10
W
0.8
250
0.8
accuracy
accuracy
0.6
200
0.6
ESS
0.4-
150
0.4
0.2
100
0.2-
50
0.0
0,0
('none))
(nmme:)
SA Layer Indices
SA Layer Indices
(a)
(b)
Figure 34. (c) Model accuracy and max/average ESS of SA layers in the 4-layer GLA-SA hybrid networks. (d) Model accuracy and max/average ESS of GLA layers in the 4-layer GL.A-SA hybrid networks.
Recall in Section 5, we demonstrated an application of post-training ESS analysis through the lens of model distillation. Here, we provide another example of post-training analysis that leverages ESS, this time in order to gain intuition into learned network dynamics hybridization. Hybridization is the process of arranging different operators in multi-layer sequence model (Lieber et al., 2024; Glorioso et al., 2024; De et al., 2024). Specifically, we measure the per-layer ESS across various hybrid networks and find that the precise ordering of layers significantly influences ESS dynamics, offering intuition to why certain hybrids outperform others.
In this section, we present results from a post-training ESS analysis applied to GLA-SA hybrid networks to demonstrate the ability of ESS to capture differences among hybrid networks with varying topologies. In the first experimental setting, we train all possible 4-layer GLA-SA hybrid networks and compute the per-layer ESS on each model. We use the tolerance- based ESS since we want to analyze failure modes of learning in hybrid networks. In Figure 33a, we first note that in the pure GLA model, many of the layers fail to learn expressive states (as evidenced by the tolerance-ESS being 0), offering intuition to why the model performs so poorly. Moving on the hybrid networks with single attention layer, we note that all of them perform quite well, except for the network that has attention in the first layer. Interestingly, we find that when attention is placed as the first layer, it suffers from state collapse. At higher level, this substantiates why many state-of-the-art hybrid networks (such as Jamba) do not place attention as the first layer of the network. However, such hybrids are typically constructed purely on the basis of performance: here, ESS is able to provide a distinct perspective. Next, examining the hybrids with SA layers, we find that the only poor performing topology is with attention placed in the second and third layers. Again, we find that the ESS of the attention layers is lower than what we observe in the hybrids that solve the task, indicating its usefulness as a proxy for performance beyond the 2-layer non-hybrid networks we explored in Section 4.
To clarify this, we examine the maximun/average ESS (computed across layers) of the SA and GLA layers separately to understand how each relates to model performance. Notably, we find that maximum ESS across attention layers best correlates with accuracy (Figure 34a). Interestingly, the average SA layer ESS is a worse proxy for performance, potentially indicating that having a single layer with high memory utilization in hybrid networks is more important than having many layers with lower memory utilization. This offers support as why hybrid networks like Jamba have 1:7 ratio between attention and non-attention layers. Regarding the GLA layers, we find that despite both the maximum and average SS varying across models, they do not correspond to changes in accuracy. One possible explanation for this is that since the attention layers are responsible for driving the total ESS of the network up due to their unbounded state size, the role of non-attention layers in hybrid networks may not be captured entirely by the magnitude of their ESS. Nonetheless, this is something we hope to explore in future work.
In the second experimental setting, we move beyond l-yer GLA-SA hybrids to 8-layer GLA-SA hybrids. Here, instead of iterating over all possible topologies, we restrict the space of networks to those constructed via the hybridization policy
52

---

Quantifying Memory Utilization with Effective State-Size
proposed by Jamba. The Jamba hybridization policy takes in the number of layers as input and provides particular hybrid topology output (refer to Lieber al. (2024) for more details). Since most topologies explored in the 4-layer setting solved the task, we both reduce the model dimension of the network and make the task more difficult to see if we can observe performance differences across the architectures (model settings can be found in Table 8). Unsurprisingly, we find that the pure GLA network is unable to solve the task and also realizes a tolerance-hasee ESS of O in all layers (Figure 33b). However, more interesting is the fact that while the 2 SA-layer Jamba hybrid partially learns the task, the SA-layer does not. Examining the ESS shows that the attention layers in the SA-layer hybrid suffer from state collapse, which we know is highly correlated with poor performance on MQAR. This points to deficiency of fixed-topology hybridization policies like Jamba which do not take into account factors like network trainability which significantly influence model performance. Furthermore, this suggests that the ESS metric can be used to better inform the construction of hybrid networks. We hope to further elucidate these per-layer ESS trends and leverage these insights to construct novel ESS-in infomme hybridization policies in future work.
E.5. State Modulation of Large Language Models
State modulation patterns on various open- gpen-weiggt models are illustrated in Figures 35, 36, 37, 38, 39, and 40.
10
The Effect of Separator Tokens on the ESS of Mamba 2.8B
1.0e03
<Ceep' peciod
tol 1.De-06
(total ESS))
3-
Rmmin
Into
zefreshung
Farth
gollacts
typ
1940
Grent
Figure 35.
The Effect of Separator Tokens on the ESS of Pythie-2.811
1.0e01
<app>
tol
1.0e 03
<sepp>
peciod
tol 1.0e-6
(toial
ricb
into
weitees
frem
dhe
Jopite
Rist
Sood
Figure 36.
53

---

Quantifying Memory Utilization with Effective State-Size
x10
The Effect of Separator Tokens the ESS of Phi-2 2.7B
[o) 1.0e
01
<eep2 <ete>>
1.0-03
<> periv
1285).
2-
nich
lnto
mooth
or
sile
Earth
gellect
viu vidaaee
typew
the
1840
barites
Orent
Tuee
Figure 37.
X10
The Effect Separator Tokens the ESS Phi3-Mini-IBs Mintrrrrc--t--trrcttt 3.8B
1.2
Tol
11
<cepp> <eos
1.0-03
<tepp: period
1.0
1.0e05
ESSV
0.8
0.2
0.0
Figure 38.
x10F
The Effect of Separator Tokens the ESS of MistralIB-v0..1
1.50
1.0e-01
<sep.>
1.0e-0
<Cepp
period
1.25
tol
1.0e 05
1.00
[toaa]
Wape
Bum
yoite
Oweest
Sove
g
Figure 39.
54

---

Quantifying Memory Utilization with Effective State-Size
X10"
The Effect of Segurator Tokens on the ESS of MistealiB-v0..2
txl
1.0601
<Mp2 <00>>
<sep>> pectod
1.0e005
0.25
0.00
todo
jpitee
pigting
Figure 40.
The Effect of Separator Tokens on the ESS of Jlama-S-Instruet SB
1.
<eepp>
1.2
100-00
<cepp> period
1.0e00
1.0
0.4
0.2
Jnto
g cnfrebbing
cood
galleccka
typ
1944
Juniter
Grema
Hlgndeuus
Figure 41.
The figures above reveal significant cross-architectural differences in context processing. The attention-based model at a similar 7B scale (Figure 39, 40, and 41) shows minimal change in its ESS pattern when an EOS token is replaced with a period (*."). In contrast, the limited cache-size state-space model (Falcon Mamba 7B, Figure 9a) exhibits a substantial reduction in state modulation under the same token substitution.
We attribute this difference to a phenomenon we term "preemptive state modulation" in limited state-size models, which stems from fundamental architectural differences. State-space models (SSMs) with limited cache must efficiently manage their finite memory capacity and learn to preemptively modulate state-size to optimize information retention, relying on explicit signals like EOS tokens to trigger context resets. In contrast, attention models with linearly increasing cache can store all past information without the need for selective forgetting, do not require preemptive state modulation, and show less sensitivity to explicit demarcation tokens. This distinction highlights the different strategies employed by various model architectures in managing context across diverse inputs, potentially influencing their performance on tasks requiring long-range recall context separation.
However, subset of attention models demonstrated varying state modulation patterns in response to different separator tokens, with this effect being more pronounced in smaller model sizes (see Figure 36, 37, and 38). This phenomenon, while not consistent across all attention architectures, merits deeper exploration.
Figure 42 illustrates the state modulation patterns at different tolerance levels for the four 1B language models (LA, WLA, GLA, SA), trained under identical conditions.
55

---

Quantifying Memory Utilization with Effective State-Size
Separator Token vs. FSS (tol=1 00-033
Separator Token vs. ESS (tll=1.0e02))
GLA
<eppor <000>
2.5
GLA
capps <002
caeps: period
WLA SA
coeps period
Total Eilletive State
2.0
Elltctive
1.5
1.0
0.5
0.0
Context
Context
Contect
Context
Context
Context
(a)
(b)
X10
Separator Token ESS (tol=1.0e-01
X10P
Separator Token ESS (tal= 1.0++00)
LA
GLA
Kaepo: <e00>
Totall Elleciive Stade Size
LA
GLA
<sepp: <ceu>>
WLA
SA
<eep2: period
WLA
SA
coeple period
Eltective
Context
Content
Context
Context Context
Context
(c)
(d)
Figure 42. An illustration of the effect of different separator tokens over different layers across different tolerances. Softmax attention exhibits the most pronounced state modulation, beginning at a tolerance level of 1e-2, followed by gated linear attention with significant modulation starting at tolerance of le-1. Weighted linear attention shows minimal modulation, only detectable a tolerance of 1.0, while linear attention displays no discernible separator token-induced state modulation. Here ESS is summed across channels and layers.
Notably, GLA exhibits a substantial variation in state modulation depending on the separator token, consistent with our earlier observations in Falcon Mamba regarding preemptive state modulation. contrast, SA shows a smaller, yet non- trivial effect. WLA and LA show no discernible differences across separator tokens, which may be attributed their overall limited ability modulate state size.
56

---

Quantifying Memory Utilization with Effective State-Size
E.6. Miscellaneous
E. EFFECTIVE STATE-SIZE ON C++ CODE
Beyond sentence delimiters such as periods and end-of-speech tokens (discussed in Section 5.3), we observe similar "dips" in effective state-size where there scope delimiter tokens such ")"
The following plots demonstrate the ESS pattern of Llama3-8B processing the C++ code of a fast inverse square root algorithm and a Fibonacci sequence generator algorithm.
Quake fast inverse square-root algorithm:
106
ESS of Quake Fast Inv sqrt Algorithm Llama-3-8B-Instrucc
tol 1.0c-01
Total Elfective State Size
3-
2
Figure 43. Effective state-size quake fast inverse square root algorithm's code. Here ESS is summed channels and layers.
#include <icstream>
#include <cmath>
// Quake Fast Inverse Square Root function
float quakeFast (float number)
long i; float x2, y7
fleat threehalfs 1.58;
10
x2 number * 0.5F;
= number;
(longx) 6y;
// Bit-level hacking: convert float long
13
0x5f3759dE (i 1); // Initial magic number and bit shift
14
y 81;
// Convert back from long to float
15
16
// Newton's method step for refining the result
17
y y (threehalfs (x2 y y)); // First iteration
18
19
return y7
20
22
int main()
23
float number;
24
// Input: Get the number from the user std::cout "Enter number: ";
27
std::cin number;
28
// Output: Display the result using the Quake fast inverse sgrt float quake_result quakeFastInvSqrtt (number) std::cout << "Quake Fast Inverse Sqrt: quake_result std::endl;
12
57

---

Quantifying Memory Utilization with Effective State-Size
// Compare with standard sqrt function
34
float std_result 1.0f / :td:t (number);
35
scd::cout << "Standard Inverse Sqrt: << sesur << sed:;
return 0
38
Fibonacci sequence generating algorithm:
x10
ESS of Fibonacci Algorithm on Llama-3-8B-Instaue
tol
1.0c-U1
1.75
tol
1.0c-12
Total Edfective State Size
1.50
1.25
1.00
0.75
0.50
0.25
Figure 44. Effective state-size over Fibonacci sequence generator algorithm's code. Here ESS is summed across channels and layers.
#include <iestream>
int fibonacci (int n) {
it (n <= 1) I
return n;
return fibonacci
fibonacci 2); Recursive case
int main() int n;
std::cout "Enter positive integer:
11
std::cinn >> n;
12
std::cout << "Fibonacci number at position << n << is: << fibonacci (a) << std::
endl; return 0;
14
E.6.2. How THE NUMBER OF PROMPTING SHOTS AFFECTS THE EFFECTIVE STATE-SIZE OF LANGUAGE MODELS
Here, we explore how varying the number of shots when prompting large language models affects their effective state-size patterns. We use Phi-2 as candidate attention model and Mamba-2.8B as the state-space model. The task we tested this on MMLU (elementary mathematics).
At the start of the Q&A section for the attention model, there is noticeable difference in state size between 0-shot and 1-shot prompts. Beyond 1-shot, the difference in ESS appears minimal. For the state-space model, varying the number of shots has minimal impact on the effective state-size.
Although these sparse experimental results require further investigation, we note the stark difference in the effective state- size patterns between these two architectures, which provides additional insights into understanding the fundamental dif- ferences in the way prompts are processed across models.
58

---

Quantifying Memory Utilization with Effective State-Size
x10e
Phi-2 2.7B on MMLU (Elementary Mathematics)
Total Effective State Size (tol (tol=0.001)
Osshots (Ace:: 0.35)
1-shot (Acc: 0.41)
S-hhol (Acc: 0.44)
Gehuut (Axc:: 0.43)
6-
2-
Eg
Eg.
Eg. 3
Eg.
Eg.
Q&cA
Figure 45. The variation in effective state-size with varying number of shots (2.7B Attention). Here ESS is summed across channels and layers.
X106
Mamba-2.8B MMLU (Elementary Mathematics)
3.5
Total Effective State Size (t0=00.001) (.000)
O-shote (Ace:: 0.27)
1-shot (Acc: 0.21)
3-abots (Acc: 0.23)
Srabots (Aco.: 0.26)
2.5
2.0
1.0
0.5
0.0
Eg.
Eg
Eg.
Eg.
Fg.. 5
Q&A
Figure 46. The variation in effective state-size with varying number of shots (2.8B State-Space Model). Here ESS is summed across channels and layers.
59


Reviews:
[{"review_id": "8uGwpl591V", "reviewer": "Reviewer_Z9Uz", "summary": "This paper claims to propose a new metric \"effective state-size (ESS)\", which can not only\nevaluate the memory-utilization of different models, but also brings instruction to the selection\nof initialization/regularization, and distillation strategies. Several empirical results are presents to\nsupports the effect of ESS.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your detailed feedback. Below, we address the main concerns.\n\n## Mischaracterization of Our Claims\n\nYour review states that our work involves measuring the rank of parameter matrices or hidden states using SVD—a practice that has existed for years. However, we would like to clarify that this is a misrepresentation of our contribution, and believe that the subsequent criticisms that we are simply renaming an established method while proposing an overly complex approach to analyze memory utilization, do not apply to the actual claims made in our paper.\n\nWe do not simply apply SVD to parameter matrices or “hidden states” (i.e., the input/output activations of a neural network layer), of which, the former capture complexity only across the model’s channels and cannot be used for analyzing memory utilization or capacity across the temporal dimension. Instead, the ESS metric is computed by applying SVD to specific submatrices of the materialized, flattened operator T in causal sequence models. These submatrices capture complexity across both channel and time dimensions, and we show that its rank provably lower bounds the minimal state size of an equivalent linear recurrence—offering a theoretically grounded proxy for memory utilization in causal sequence models.\n\nMoreover, prior approaches that apply SVD over the entire attention matrix (i.e., the operator T) typically ignore causal masking, which greatly influences the rank of T and distorts any interpretation related to memory usage. Our metric is explicitly designed to account for this, and we discuss this aspect of our contribution in the introductory section of the paper.\n\nImportantly, we note that the underlying theoretical framework for computing minimal recurrent realizations stems from classical signal processing and control theory, and to the best of our knowledge, has not been adapted to analyze the memory utilization of modern deep learning sequence models such as input-varying SSMs, linear attention and softmax attention.\n\nWe summarize the core contributions below:\n1. We demonstrate that most modern deep learning causal sequence models (SSMs, convolutions, attention, etc.) can be cast as input-varying linear models (LIVs), which uniformly realize a materialized operator T. In doing so, we are not introducing an overly complex framework; rather, we show that diverse sequence models can be analyzed under a unified lens.\n2. Drawing from classical control theory, we prove that the rank of specific submatrices of T provides a lower bound on the minimal state-size of an equivalent SSM realization. Based on this result, we propose ESS as a principled and model-agnostic proxy for memory utilization.\n3. We empirically validate ESS using a wide range of modern sequence models and synthetic tasks that are explicitly designed to test memory utilization.\n4. Finally, we show how ESS can be used in practice to inform and analyze downstream tasks such as model distillation, initialization, and regularization, as well as provide insight into memory-related phenomena in language models.\n\nWe hope this clarifies our main claims and the novelty of our contributions.\n\n## Improvements on Presentation\n\nWe thank the reviewer for pointing out areas where the presentation can be improved. We will more clearly define metrics such as ESS/kv and TSS/kv (i.e., ESS normalized by the number of key-value pairs in a task), and clarify related notation earlier in the manuscript.\n\nRegarding the experiments, ESS can be used to support a range of downstream analyzes. First, we would like to clarify that the distillation results are in the main body of the paper (see Section 5.1). However, due to space constraints, the initialization and regularization experiments are presented in the appendix along with all relevant experimental details. We will improve the referencing of these results in the main text to ensure clarity and visibility. Also, we are happy to rearrange the paper to fit some of these results in the main body if the reviewer sees fit.", "claims_and_evidence": "- Calculating the actual (as opposed to theoretical) information capacity using Singular\nValue Decomposition (SVD) and the Rank of parameter matrices/hidden states has been a\nwidely adopted metric among researchers and engineers for many years—its origins are\ndifficult to trace.\n- Additionally, the relationship between singular values and the rank of parameter\nmatrices/hidden states appears trivial and intuitive, raising questions about the necessity\nof the overly complex modeling.", "methods_and_evaluation_criteria": "- Some aspects hold significance, yet there are technical details that necessitate further verification and elaboration.", "theoretical_claims": "- However, this paper claims to introduce ESS, an SVD-based metric, for evaluating model\nmemory utilization. This approach appears highly similar to the common practice\nmentioned above. Simply renaming an established method and presenting it as a novel\ncontribution is unlikely to be encouraged.", "experimental_designs_or_analyses": "I have reviewed the relevant experimental sections, and the pertinent information can be found under the section labeled \"Weakness.\"", "relation_to_broader_scientific_literature": "None", "essential_references_not_discussed": "- Wang, Shida, and Qianxiao Li. \"Stablessm: Alleviating the curse of memory in state-space models through stable reparameterization.\" arXiv preprint arXiv:2311.14495 (2023).\n- Qi, Biqing, et al. \"Smr: State memory replay for long sequence modeling.\" arXiv preprint arXiv:2405.17534 (2024).", "other_strengths_and_weaknesses": "## Strenghs:\n- Within the view of ESS, this paper brings lots of comparisons between models with different\narchitecture, such as attention/recurrent/convolution. These comparison may brings insight on\nthe design of new architectures in the future.\n\n## Weacknesses\n### 1、Over-claimed Contribution (which is also their primary claimed contribution)\n- Calculating the actual (as opposed to theoretical) information capacity using Singular\nValue Decomposition (SVD) and the Rank of parameter matrices/hidden states has been a\nwidely adopted metric among researchers and engineers for many years—its origins are\ndifficult to trace.\n- However, this paper claims to introduce ESS, an SVD-based metric, for evaluating model\nmemory utilization. This approach appears highly similar to the common practice\nmentioned above. Simply renaming an established method and presenting it as a novel\ncontribution is unlikely to be encouraged.\n- Additionally, the relationship between singular values and the rank of parameter\nmatrices/hidden states appears trivial and intuitive, raising questions about the necessity\nof the overly complex modeling.\n### 2. Ambiguous Presentation and Inappropriate Typography\n- Certain metrics, such as ESS/kv and TSS/kv, are used before being properly defined.\n- The captions and explanations of figures are too simple, making it difficult for readers to\nfollow the paper’s contributions.\n- The paper claims that ESS is valuable for initialization, regularization, and distillation\nstrategies, highlighting its practical significance. However, no relevant empirical results\nare provided in the main body to support this claim.", "other_comments_or_suggestions": "See weakness.", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {"ethical_review_flag": "Flag this paper for an ethics review.", "ethics_expertise_needed": ["Other expertise"]}}, {"review_id": "UKxhcwKLCL", "reviewer": "Reviewer_KvDY", "summary": "The paper introduces Effective State-Size (ESS) as a novel measure for quantifying memory utilization in causal sequence modeling architectures. ESS provides interpretable and actionable metrics that can enhance initialization strategies, regularizers, and model distillation. The paper develops a unified framework for analyzing systems with input-invariant and input-varying linear operators (LIVs), and demonstrates the correlation between ESS and performance across various tasks and models. Applications of ESS include model-order reduction, predicting model compressibility, and state modulation in large language models. The empirical validation shows ESS's utility in improving performance-efficiency trade-offs and highlights cross-architectural differences in memory utilization.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your constructive review and positive outlook on the work!\n\n## Presentation and Accessibility\n\nThank you for pointing out that some theoretical sections may be inaccessible to many due to technical jargon. We will simplify those parts and add working examples of computing ESS, as Reviewer KxgK suggested. We will also simplify the introduction by stripping some of the more technical motivation in order to make it more accessible. In addition, we will distill Section 3.1 to make room to add the pseudocode for the implementation of ESS, as you have suggested. \n\n## Real-World Benchmarks\n\nWe agree that applying ESS to popular benchmarks could be valuable. However, we note that many real-world NLP benchmarks often test factual recall (i.e., returning accurate information from the pre-training/post-training data) and not working memory utilization. For example, the Mamba paper (Gu et al., 2024) shows strong performance on standard benchmarks (i.e., against Pythia models (Biderman et al., 2023) of the same scale) but performs poorly on tasks that test raw memory utilization, such as the phone book recall task (Jelassi et al., 2024).\n\nFor this reason, we focused on tasks that directly test memory recall and chose bigram perplexity (Arora et al., 2024) for evaluating the language models (Section 5.2). Bigram perplexity evaluates the model's perplexity on repeated bigrams in any arbitrary dataset. This allows us to use the diverse and large scale NLP pre-training datasets (that the model wasn’t trained on), as opposed to a narrow hand-crafted evaluation set to specifically evaluate the model’s ability to utilize past context. In our case, we extracted 16k randomly selected sequences from the Fineweb dataset (Penedo et al., 2024). Therefore, we maintain that the controlled experiments we have conducted effectively conveys the utility of ESS as a diagnostic tool for analyzing memory utilization.\n\n- Gu, A., & Dao, T. (2024). Mamba: Linear-time sequence modeling with selective state spaces. arXiv. https://arxiv.org/abs/2312.00752\n- Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. arXiv. https://arxiv.org/abs/2402.01032\n- Biderman, S., Schoelkopf, H., Anthony, Q., Bradley, H., O’Brien, K., Hallahan, E., Khan, M. A., Purohit, S., Prashanth, U. S. S., Raff, E., Skowron, A., Sutawika, L., & van der Wal, O. (2023). Pythia: A suite for analyzing large language models across training and scaling. arXiv. https://arxiv.org/abs/2304.01373\n- Arora, S., Eyuboglu, S., Timalsina, A., Johnson, I., Poli, M., Zou, J., Rudra, A., & Ré, C. (2023). Zoology: Measuring and improving recall in efficient language models. arXiv. https://arxiv.org/abs/2312.04927\n- Penedo, G., Kydlíček, H., Ben Allal, L., Lozhkov, A., Mitchell, M., Raffel, C., Von Werra, L., & Wolf, T. (2024). The FineWeb datasets: Decanting the web for the finest text data at scale.", "claims_and_evidence": "1. The introduction of Effective State-Size (ESS) as a quantitative measure of memory utilization is a novel and valuable contribution. It provides interpretable and actionable measurements that can enhance initialization strategies, regularizers, and model distillation.\n2. The paper explores multiple applications of ESS, such as model-order reduction, predicting model compressibility, and state modulation in large language models. These applications are well-explained and show the practical relevance of ESS.", "methods_and_evaluation_criteria": "The authors provide thorough empirical validation of ESS across various tasks and models. The correlation between ESS and performance is demonstrated, highlighting the utility of ESS in improving performance-efficiency trade-offs.", "theoretical_claims": "The paper's theoretical sections, particularly the formal definition and computation of ESS, look solid, but may be challenging for readers who are not well-versed in dynamical systems theory. Simplifying these sections or providing additional explanatory material could improve accessibility.\n\nAlso, I understand the implementation details for computing ESS are discussed in the appendix. But I feel that it would be great if pseudocode can be included in the main text.", "experimental_designs_or_analyses": "The authors showed the correlation between ESS and accuracy on some synthetic tasks, which I understand are especially created to conduct experiments in a controlled manner. However, most people in the field are more familiar with other real-world benchmarks. Does the same correlation extend to popular real-world benchmarks, e.g., machine translation benchmarks (e.g., WMT), language understanding (e.g., GLUE, MMLU), question answering (e.g., SQuAD, PIQA), image classification which can use Vision Transformers (e.g., ImageNet)? It would be great if the authors can include some of these common benchmarks in their experiments, which will strengthen this paper.", "relation_to_broader_scientific_literature": "I believe understanding the memory utilization is of interest to many researchers in the field.", "essential_references_not_discussed": "looks good to me", "other_strengths_and_weaknesses": "NA", "other_comments_or_suggestions": "NA", "questions_for_authors": "NA", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "dD90tZaFuc", "reviewer": "Reviewer_AdxY", "summary": "This work introduces the Effective State-Size (ESS) metric to quantify memory utilization in sequence models while previous approaches focus on memory capacity (such as cache size/total memory available). ESS aims to measure how effectively a model uses its available memory. Using this metric the authors analyze 4 kinds of sequence models on a synthetic task and demonstrate how the metric helps to explain their performance, as well as how it predicts performance after distillation for some models.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive feedback! Below, we respond to the points raised. \n\n## Claims and Evidence \n\n1. Yes, this is a good point, we will extend the regularization experiments to include parameters beyond the range in the plot and update the plot in the paper. We anticipate that at some point, the performance will begin to decline since continuing to increase the regularization strength corresponds to increasingly decaying the model towards linear attention, which, as shown in Section B, performs quite poorly. \n2. Yes, your understanding of how the ESS metric correlates with performance of distillation is correct. We will clarify this in the paper. \n\n## Methods and Evaluation Criteria\n\n1. For theoretical justification as to why the rank represents state utilization, we refer the reviewer to Theorem 3.2. For intuition, we offer the following interpretations of rank as it pertains to state utilization:\n    - **Distinct directions of influence**: Through this lens, rank counts how many linearly independent “directions” connect past inputs to future outputs, i.e., how many unique ways inputs can shape the future.\n    - **Minimal internal memory**: Because each independent direction requires its own coordinate in memory, the rank matches the smallest state dimension that can exactly replicate the operator.\n    - **Effective rank**: Of course, in practice we compute effective rank to capture the “dominant” directions of interests where “dominant” is captured by either some threshold (e.g. tolerance ESS) or by the decay rate of the singular values (e.g. entropy ESS). \n2. Here, we address the concern posed by the reviewer regarding the correlation values being too “low”.\n    - **Having sufficient ESS/kv results in a non-linear effect**: As shown in Figure 2, once ESS/kv exceeds a certain threshold for a given model, the majority of task-model configurations achieve an accuracy of 1. At this point, performance saturates, and further increases in ESS/kv do not translate into higher accuracy. This weakens the correlation, as correlation coefficients do not adequately capture such non-linear relationships. We chose not to omit these data points in our reported correlations, since the non-linear trend is clearly visible in the plot. However, we are happy to also report correlations with these saturated points omitted, should the reviewer find that informative.\n    - **On the interpretation of “low” correlation values**: Our claims about correlation are relative rather than absolute. While it is fair to consider correlations in the range of 0.85–1.0 as “high,” our point is that ESS consistently correlates better with task performance than TSS. In this comparative context, we maintain that ESS is a significantly better proxy. This is relevant since most works consider only the TSS when evaluating model memory. Note that we do not claim that ESS is the best possible performance predictor—there may be stronger alternatives—but evaluating such alternatives is beyond the scope of this work.\n\n## Clarifying notation\n\nTo answer your question in particular, $C$ and $B$ are analogous to the $Q$ and $K$ matrices in attention. And in the context of recurrent models, they canonically represent the state-to-output and input-to-state matrices, respectively. However, we agree that this is unclear in the paper and we will clarify this upon revision along with analogous points raised by reviewers KvDY and KxgK.", "claims_and_evidence": "The claims seem to be generally supported:\n1. The appendix details their experiments on trying to improve performance through ESS-informed regularization (although the accuracy only reaches 0.3 - would this keep going up with more regularization or has it saturated at this point? This needs to be shown.)\n2. The appendix details their experiments on trying to improve performance through ESS-informed initialization for GLA vs S6 models.\n3. They show how their metric correlates with the performance of distillation dependent on whether the state utilization of the teacher model was high or low (if it was high then it should be expected that the distillation should not be very successful if I understand correctly).", "methods_and_evaluation_criteria": "1. The method to study sequence models' usage of their state space through a newly defined metric makes sense. It would be good to get some intuition though of why rank represents utilization of state for recurrent operators since I am not from a signal processing background.\n2. The task used throughout the paper is synthetic - MQAR (with additional results in appendix for compression and copying - also synthetic).\n3. The paper reports correlations between their metric and quantities of interest but something that is worrying is that the actual correlation values are pretty low (0.5-0.7 range) even if the trends in changes in correlation (with changes in models etc) make sense.", "theoretical_claims": "I was not able to check the proofs throughly since I am not very familiar with the literature in this area and the paper's description of notation was confusing (for eg. what are C and B in section 3 is not defined anywhere and has to be guessed).", "experimental_designs_or_analyses": "I read the experiments section throughly but I'm not very familiar with prior work that has shown results on the same tasks.", "relation_to_broader_scientific_literature": "Many prior studies focus on measuring memory capacity rather than utilization, ESS refines this by measuring actual usage rather than theoretical capacity.\nAlso, prior work mostly used qualitative measures (e.g., attention visualizations, synthetic task accuracy). The introduction of ESS as a precise quantitative metric allows drawing insights in a more automated way.", "essential_references_not_discussed": "-", "other_strengths_and_weaknesses": "I found the visualizations of ESS around EOS tokens pretty interesting.", "other_comments_or_suggestions": "-", "questions_for_authors": "-", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6EYl4pIyQw", "reviewer": "Reviewer_KxgK", "summary": "The paper proposes to study the memory stored in a wide range of sequential neural network architecture through the notion of effective state-size, which is motivated by minimal realization theory and applicable out of the box to many architecture. The authors empirically validate it as a sound measure by correlating it with performance on memory intensive tasks, as well as use it to better understand which kind of model can be easier to distill and to derive initialization strategies.\n\n**[EDIT 04/02/2025]**: updated score from 2 to 3.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed comments! The paper will be revised to better explain each of the following points:\n## Why is ESS a function of time? Why does ESS decrease before the EOS token?\nESS at each time step $i$ captures the lower bound for the minimal state-size required at that specific step. Intuitively, this bound reflects at least how much information from the past is relevant for future computation. As the sequence progresses, the model may retain more or less information depending on upcoming needs. While we can summarize ESS over time with an aggregate measure (e.g. max or average, as done in Sections 4 and 5.1), examining ESS at each time step can reveal interesting temporal patterns—particularly w.r.t. a model’s ability to store and discard information dynamically.\n\nBecause our formulation of ESS depends on both past and future relevance, it decreases when the model requires less prior context to process upcoming tokens. Near the EOS, fewer future tokens remain to be influenced by previous tokens, so ESS tends to shrink before the actual EOS token. \n\nWe have also identified ESS variants that measure a causally determinable minimal state-size (i.e. it depends only on the past). Under this metric, state-size drops sharply at the EOS token, rather than exhibiting a more gradual decrease. In Section D.2.2, we briefly discuss these “causal ESS” metrics. Instead of computing the rank of $H$, one can decompose $H$ into its causal and anti-causal parts ($H = \\mathcal{OC}$) and then compute the rank of $\\mathcal{C}$ to obtain a causal ESS. However, this metric may fail to capture certain insights, because for models like softmax attention, $\\mathcal{C}$ is effectively an identity matrix that grows with sequence length (substitute Eq. D.2.5 into $\\mathcal{C}$). As a result, the causal ESS for softmax attention increases linearly with sequence length, failing to reflect how state-size might rise and fall near the EOS token. \n## Comparing ESS to an Effective Window-Based Metric \nCounting the number of state derivatives above a certain threshold effectively measures the window of the model’s memory. Intuitively, a smaller window suggests lower memory utilization. This phenomenon is also reflected in the ESS metric: when the operator decays quickly, most rows/columns in $H_i$ are zeroed out, leading to a lower rank and hence a smaller ESS.\n\nHowever, the effective window alone does not capture the complexity of dependencies within the window. For instance, for linear attention, $\\frac{ds_{i+j}}{ds_j}$ is always an identity matrix, regardless of $i$ or $j$. Consequently, a window-based metric that simply counts these derivatives would increase indefinitely with sequence length and fail to reflect the model’s true memory capacity or utilization that is capped. Moreover, because the window-based metric does not account for complexity, it ignores the effects of $B$ and $C$ (input and output projections), which themselves can be degenerate (e.g. zeroed-out rows).\n\nFor these reasons, ESS provides a more accurate and theoretically grounded measure of memory utilization. Nonetheless, we acknowledge that alternative metrics, such as the one proposed by the reviewer, can still offer useful heuristic insights when analyzing a model’s memory.\n## Criticism of Approaches in Wu et al. (2024) and Related Works\nWe are not criticizing Wu et al.; rather, we cited their work as they similarly highlight how some spectral analyses of the attention matrix (i.e. the operator T) overlook the causality inherent in sequence decoders. That said, we do point out drawbacks in approaches such as Min et al. (2024) and Bhojanapalli et al. (2020), which derive their metrics by taking the SVD of the operator T without accounting for causal masking. For instance, Min et al. state: “It should be noted that the training of GPT-2 necessitated the masking of attention matrices to prevent the model from accessing future data. However, this mask was not applied during the rank evaluation.” Consequently, these approaches fail to tie key notions like memory capacity and memory utilization.\n## Further improvements in writing\n1. We will reduce technical jargon or add contextual information where needed in the introduction to improve accessibility.\n2. We will move Section 3.1 to the appendix to shorten the main text, enabling space for a toy example of ESS (illustrating how ESS varies with time/input). Figure 1 will be moved to the appendix if needed.\n3. Regarding your question on what a model's \"ability to dynamically modulate its ESS in response to inputs\" captures, our results in Section 5.2 indicate that recall performance depends on a model’s ability to modulate its ESS, not just total memory capacity. Although WLA, GLA, and LA share the same memory capacity, their recall differs, implying differences in memory utilization. Nonetheless, we will clarify in the text that ESS only **suggests** this dependence, rather than conclusively proving it.", "claims_and_evidence": "The claims are supported through extensive theoretical and empirical verification.", "methods_and_evaluation_criteria": "The focus on memory intensive tasks is particularly relevant for the proposed approach. The distillation part also makes sense. I have some concerns regarding the setting of Section 5.2 that I detail below.", "theoretical_claims": "The theoretical claims are mostly rooted on established results from the theory of recurrent realization. I haven't checked the correctness of the claims though, and I am not familiar with this literature.", "experimental_designs_or_analyses": "No, but I didn't notice anything raising my awareness.", "relation_to_broader_scientific_literature": "The paper does a great job at linking the ideas it introduces / problems it studies with existing work.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths**:\n- the paper studies a very timely problem, that is the one of in-context memory, of sequential models.\n- the approach is theoretically well motivated, looking at the minimal representation is an interesting angle.\n- most of the experiments make sense. The memory experiments are a good test bed and the proposed method brings an interesting angle to the distillation case.\n- the state saturation vs. state collapse argument is an interesting lens on studying failures of models on memory heavy tasks.\n\n**Weaknesses**:\n- the main weakness to me is the presentation of the paper: I find hard to get an intuitive understanding of the method (see next box for more details).\n- as a result of the previous point, I have some troubles judging whether the proposed method makes sense. This justifies my conservative score for now (weak reject), and I am happy to revise my score once I am more confident in my understanding of the method.\n- The authors claim that model performance on language depends on the model's \"ability to dynamically modulate its ESS in response to inputs\". While I agree with the underlying intuition, I fail to see what ESS is bringing here and would appreciate clarification from the authors.", "other_comments_or_suggestions": "**Regarding the writing.**\nThe clarity of the paper could be greatly improved in my opinion. I appreciate the authors' effort in summarizing important results in boxes, but more should be done for the paper to be ready for acceptance. To list some improvement points:\n- I find the introduction overly technical. As a rule of thumb, I would avoid any mathematical formulation and dynamical system jargon in the introduction.\n- Introducing one or two working example while going through the theoretical section could be useful, e.g. a simple linear SSMs (eventually with rank constraint on the A matrix) and/or an SSM with an input-dependent transition matrix. This would help the reader (me included) to get a better feeling of what the proposed measure is capturing.\n- Moving section 3.1 from the main text to the appendix might be a good idea (keeping a high level summary in the main text would be enough). I find it overly technical for the main text and it does not really seem to be needed in the following.\n- [Minor] Using the terms \"memory capacity\" and \"memory utilization\" would make the text easier to read than the repetitive use of ESS / TSS, which I find slightly confusing.\n- [Minor] Figure 1 is not really dense in terms of information, and the current version could likely be removed without hindering the understanding of the paper.\n- [Minor] Suggestion to improve Figure 5: I would find more interesting to see how the loss evolves as a function of the student TSS. In particular I would like to see how much bigger it needs to be than the ESS of the teacher to get a better than chance performance. The current visualization makes this difficult to assess.", "questions_for_authors": "- I am confused by the fact that the effective state-size is a function of $i$. I cannot really make sense of it: from my understanding $i$ indexes time, and I would expect the effective state-size to be something independent of $i$, for example the maximum of all the $\\text{ESS}_i$ values. Could the authors clarify that point? Having a working example, as I mentioned in the previous section, could be a good way to clarify things.\n- I cannot really make intuitive sense of why the ESS decreases *before* the end of the sequence in Figure 7. Intuitively, I would expect the separator to remove information about previous states but not the tokens appearing before, at least in a causal system. Why isn't it the case? Are there some boundary effects of the method?\n- Can the author compare their method to the following approach: for all mentioned models except softmax attention, look at $ds_{i + j} / ds_j$, compute on average for how many time steps it is greater than some threshold an compute the capacity of the model as this value times the number of states. Intuitively, this would correspond to the \"kv\" cache size of the model. I would appreciate if the authors could compare with this type of metric (theoretically / empirically through a comparison in a simple setting / ...). Linked to that point is the last sentence of the last paragraph in page 1: I could not find a detailed criticism of approaches like Wu et al. 2024. Discussing them in more detail would help better understanding the contribution of the present work.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: In light of the increasing popularity of new language model architectures, this paper proposes a new quantitative metric of memory utilization in general sequence models that can be represented as input-dependent linear transformations. This effective state size (ESS) metric relates to the effective state capacity of a recurrent model that can capture the given sequence transformation, and provides interesting new quantitative and qualitative analysis of the performance of recurrent sequence models. The proposed contributions increase the set of tools available to researchers studying and designing modern sequence models.
Decision: Accept (poster)