Quantifying Memory Utilization with Effective State-Size

Thank you for your detailed comments! The paper will be revised to better explain each of the following points:

Why is ESS a function of time? Why does ESS decrease before the EOS token?

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

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

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

Comparing ESS to an Effective Window-Based Metric

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

However, 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).

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

Criticism of Approaches in Wu et al. (2024) and Related Works

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

Further improvements in writing

1. We will reduce technical jargon or add contextual information where needed in the introduction to improve accessibility.
2. 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.
3. 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.

We thank the reviewer for the positive feedback! Below, we respond to the points raised.

Claims and Evidence

1. 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.
2. Yes, your understanding of how the ESS metric correlates with performance of distillation is correct. We will clarify this in the paper.

Methods and Evaluation Criteria

1. 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:
   • 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.
   • 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.
   • 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).
2. Here, we address the concern posed by the reviewer regarding the correlation values being too “low”.
   • 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.
   • 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.

Clarifying notation

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

Thank you for your constructive review and positive outlook on the work!

Presentation and Accessibility

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

Real-World Benchmarks

We 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).

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

• Gu, A., & Dao, T. (2024). Mamba: Linear-time sequence modeling with selective state spaces. arXiv. https://arxiv.org/abs/2312.00752
• 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
• 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
• 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
• 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.

We appreciate your detailed feedback. Below, we address the main concerns.

Mischaracterization of Our Claims

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

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

Moreover, 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.

Importantly, 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.

We summarize the core contributions below:

1. 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.
2. 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.
3. We empirically validate ESS using a wide range of modern sequence models and synthetic tasks that are explicitly designed to test memory utilization.
4. 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.

We hope this clarifies our main claims and the novelty of our contributions.

Improvements on Presentation

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

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