Provable Length Generalization in Sequence Prediction via Spectral Filtering

Responses to Weaknesses

1. Please check the main comment to all reviewers where we address this.
2. Assumptions on the eigenvalues of A are important. There are several cases: 1) eigenvalues of A are allowed to be both real and complex, 2) eigenvalues of A are allowed to be nonnegative, 3) eigenvalues of A are allowed to be greater than 1. For (1) there aren’t any special methods that provably work to our knowledge. You can apply classical methods of prediction like regression, etc. However, to obtain meaningful results you must assume marginal stability (i.e. the magnitude of the eigenvalues is strictly bounded by 1). For (2) vanilla spectral filtering provably achieves O(sqrt(T)) regret and it is possible to adapt our result to extend to this case. For (3), there is nothing to do because the system is not stable and the values diverge as T gets large (moreover in the regime where the values are bounded given a horizon T, then rescaling them to be bounded by 1 can be done without loss of generality).
3. Please check the main comment to all reviewers where we address this.

Responses to Questions

1. The regret bound in Hazan et al. (2017b) does not make assumptions on the data, however it only compares to the best LDS predictor (which is by virtue noiseless) and therefore it does not imply that spectral filtering predicts well in general on an LDS with noise. In general LDS literature, Kalman filtering is bayes-optimal with gaussian noise (however this is only for stable systems i.e. eigenvalues of A are bounded away from 1 and therefore the memory of the system is bounded).
2. This is a very cool idea and to our knowledge the empirical methods designed to handle length generalization for LLMs in practice don’t currently try to do this. In a more general sense, theoretical analysis of how best to use memory resources for sequence prediction is an interesting question.

Response to Weakness 1

We first note that inspired by the above we have added some extra text in the introduction (highlighted in blue) to address the connection between LDSs and LLMs. We hesitate to write much more since it may make a future reader think the point of the paper is LLMs when it isn’t. Please let us know if it is helpful. We now directly answer the reviewer.

First, in our related work we reference 11 papers that address the connection between the study of online prediction in LDSs and LLMs (see “Motivated by the high memory and computer requirements of transformers, state space models were revisited…"). However, we implicitly assumed the reader would know that SSMs and the work we cited are based on considering online prediction in LDSs and we now make this connection explicit to the reader who is less familiar with this area of work. These works indeed show that methods based on studying LDSs provide empirically strong approaches in LLM applications. Thank you for pointing this out.

Second, your intuition regarding asymmetric regret and how it may imply the STU architecture is helpful in length generalization for LLMs is correct. However, this isn’t the main point of the paper and to make this point fully we are planning to do a fully experimental work. That said, we now make this point clearer in the introduction (see the blue text).

Finally, in response to what broader intuition should be gleaned from the main theorem: the fact that there exists a computable algorithm (the algorithms we introduce in our paper) which can use only T^{1/4} context length to both learn and make its predictions and obtain small regret compared to the best predictor using the full history of size T is very surprising. This result implies that spectral filtering is powerful in its ability to learn the dynamics of a complicated underlying system with long memory since it naturally handles the issue of what aspects in a sequence should be memorized for the future and what aspects can be forgotten, whereas other existing methods are hand engineered depending on the specific task. Thank you for highlighting that it isn’t clear to the reader. We have added a bit of this intuition into the discussion.

Response to Weakness 2

The fact that it isn’t clear to what extent this result is meaningful beyond LDSs makes sense given your question regarding the use of studying LDSs and their connection to language models (see above the discussion to your first point), which we will make clear. Indeed, results concerning LDSs tend to be quite powerful and useful in practice. For the second part of this weakness, please see our main note to reviewers regarding the nontriviality of the theoretical results. Regarding \phi-regret, it was first introduced by Greenwald and Jafari in 2003, and the context in which it was studied (by the references you cite, as well as Hazan and Kale 2007) is convergence to various notions of equilibrium and relationship to computation of stationary points. We are happy to elaborate on the relationship between this notion and asymmetric regret, which deals primarily with information imbalance, and is a different notion.

Response to Weakness 3

In our main note to all authors we have updated our work with an algorithm that uses two autoregressive components and requires no gap in the spectrum of the matrix. However, we have added more intuition on understanding the gap in the eigenvalues (see the blue text after Theorem 5).

Responses to Questions

1. To provide some intuition on the gap in the spectrum (which we have now included right after the statement of Theorem 5), here is the following: first imagine that all the eigenvalues of A are bounded by 1-\delta. Then the “weight” that input u_{t-t_0} has for y_t is roughly (1-\delta)^t_0 (because it is multiplied by A t_0 many times). When is this weight small? When t_0 is much larger than 1/\delta. When t_0 is much larger than 1/\delta then u_{t-t_0} doesn’t really impact y_t and so it may be forgotten. This is why we have that if the spectrum lies in [0, 1-log(T)/(8T^q)], length generalization is possible. Indeed, letting \delta = log(T)/8T^q and t_0 = T^q (which is much bigger than 8T^q/log(T) for large enough T) we see that when the spectrum of A is smaller than 1 - \delta, after t_0 = T^q many steps we can forget about the previous inputs u_{t-t0}. For the second part of the range: i.e. that the spectrum of A can lie between [1-1/(2T^{5/4}),1] we note that this is a special feature we prove about spectral filtering (this is rather technical but it is very important because it implies that systems that require long memory are able to be learned with spectral filters). The region between the first part and the second part is exactly that the range where the eigenvalues aren’t small enough that they don’t impact y_t and can be forgotten, but also aren’t large enough that the autoregressive component is able to capture them.

2. This is a good question, and the empirical evidence sure seems to suggest so. We currently don’t know how to show such theoretical lower bounds.

3. We plotted these error bars as mean plus/minus 1.96 * standard deviation. The data that generated this plot is very bimodal (either the task is solved or not), and in such a setting it is possible for the error bars to extend past the maximum observation. We figured this indicated the high variance of the data well, and we will add some brief information to the experimental section clarifying the bimodality of the data. If you have any suggestions for a better way to plot this information we can try that out as well, but to us this was the simplest and least cluttered plotting method. As a concrete example, consider a random variable that takes value 1 with probability 0.7 and value 0.5 with probability 0.3. Then, the mean is 0.85 and the standard deviation is 0.229 – the error bars we draw for such a thing would be huge, shading in the interval [0.4, 1.3].

Response to Weakness 4

• Regarding the noise in Figure 5, the high variance is due in particular to bimodality of the accuracies – length generalization is achieved sometimes. We emphasize that resolving this is a modeling/deep learning question that ought to be approached in a task-specific way – our goal was to show that the building block for these STU models (ie the spectral filtering algorithm for LDS) length-generalizes, and to hint that it can be taken advantage of in deep learning tasks. We have added a footnote in blue in the induction heads experimental section to elaborate on the nature of the high variance.

• Building on the above point: as a rough first approximation, one expects transformers to excel in precise copy/lookup tasks like induction heads due to architectural biases, whereas SSMs ought to have a more natural representation of long-range dependency and reasoning-type tasks. It is often difficult to make useful modeling decisions based on synthetic tasks, and it is even more difficult to determine which synthetic task is representative of what we actually care about in application (such as language modeling). As an illustrative example, Mamba (and related methods that introduce a pointwise selection mechanism) was created with the hope of getting the best of both of these memorization and dynamical worlds, and they used length generalization on induction heads to demonstrate this (see Table 2 in [1], where they compare with transformers and other SSMs) – however, we have recently seen in the literature Mamba’s length generalization performance is much more nuanced and exhibits different phenomena on seemingly-similar synthetic tasks (compare Figures 5 and 6 in [2], for example) and in a way that gets even more complicated in language [3]. What can be confidently said is that STU models are capable in many domains (including on these synthetic tasks, though length-generalization was not previously experimented with in STU models) [4], but we feel that it could be misleading or unuseful to focus particularly on comparison with e.g. transformers on synthetic tasks. Our goal was to prove that spectral filtering has robust length generalization, even for marginally-stable systems (the first result of this sort for algorithms that one would use as an SSM layer) and to inspire a proper empirical investigation on large-scale language modeling. If one wants to compare different models side-by-side on synthetic tasks like induction heads, citations [1] and [2] do this for transformers with various positional embeddings (the choice of which also plays a huge role) and SSMs at roughly the same scale as our experiment. However, we believe that these synthetic settings are not quite the right place to make the comparison and designate any SOTA performances on length generalization, and we leave the proper comparison (in language models on long-context evals) to future experimental works.

• [1] – Mamba: Linear-Time Sequence Modeling with Selective State Spaces, Gu & Dao, 2023.

• [2] – Repeat After Me: Transformers are Better than State Space Models at Copying Transformers are Better than State Space Models at Copying, Jelassi et al., 2024.

• [3] – DeciMamba: Exploring the Length Extrapolation Potential of Mamba, Ben-Kish et al., 2024.

• [4] – Flash STU: Fast Spectral Transform Units, Liu et al., 2024

Responses to Weaknesses

1. We first note that we have now added an extra paragraph (highlighted in blue) on asymmetric-regret (right after definition 3) to help ensure extra clarity based on this point. Please let us know if you find it to be helpful. Asymmetric regret is a stronger notion because it implies regular regret. The confusion might be with comparator classes and when those are fixed. In sequence prediction literature, regular regret considers comparator classes in which both the learner and the comparators construct their predictions based on the same sequence length. With asymmetric regret, we measure the learner against even stronger comparators which are allowed to construct their predictions using longer sequence lengths (and therefore more information). Your statement that “often regret guarantees hold for an arbitrary comparator” isn’t accurate. Indeed, if you are comparing an algorithm’s performance to an arbitrary comparator then you will often get meaningless results since you can always pick a predictor which predicts perfectly. For this reason, you will find that regret results in the sequence prediction literature have been based on a restricted comparator class. To understand the difference between asymmetric and regular regret better, fix any comparator class in sequence prediction and you can consider regular regret where the comparator class must use the same information as the learner, or you can consider asymmetric regret where the comparator class is allowed to use more information. In other words, competing against an adversary that has more information than the learner is harder than competing against an adversary that has access to the exact same information that the learner has. Therefore, a guarantee against the stronger adversary automatically implies a guarantee against the weaker one. This is a central concept in our paper, so please ask again if we can clarify further.

2. We plotted these error bars as mean plus/minus 1.96 * standard deviation. The data that generated this plot is very bimodal (either the task is solved or not), and in such a setting it is possible for the error bars to extend past the maximum observation. We figured this indicated the high variance of the data well, and we will add some brief information to the experimental section clarifying the formula for the error bars as well as the bimodality of the data. If you have any suggestions for a better way to plot this information we can try that out as well, but to us this was the simplest and least cluttered plotting method. As a concrete example, consider a random variable that takes value 1 with probability 0.7 and value 0.5 with probability 0.3. Then, the mean is 0.85 and the standard deviation is 0.229 – the error bars we draw for such a thing would be huge, shading in the interval [0.4, 1.3].

Responses to Questions

1. First, notice that our method is indeed improper - it doesn’t do system identification to create an LDS predictor, but rather predicts using spectral filters. We must have a benchmark class in regret, since otherwise the problem is too hard, this is a consequence of Kolmogorov’s theorem. To answer your question more clearly: 1) When the sequence is a linear dynamical system, then the best improper predictor can’t outperform the best LDS predictor since the sequence is in fact an LDS and therefore the best LDS predictor can predict it exactly. 2) However, you are right that improper learning algorithms are a good tool and indeed we use an improper learning algorithm. 3) Perhaps the reviewer has some confusion regarding the guarantee afforded by the Vovk-Azoury-Warmuth (VAW) forecaster( see for example Theorem 2 in this blogpost https://parameterfree.com/2019/10/17/follow-the-regularized-leader-iii-more-logarithmic-bounds/). Please note that the VAW forecaster makes improper predictions by itself but it affords a regret bound only against the class of fixed (but potentially unbounded) linear predictors, i.e. the VAW forecaster only provides a guarantee against any predictor of the type that predicts as the rule hat{y}_t = <u_t, x> where x is held fixed at every time step. In comparison, the class of LDS predictors is significantly larger as every prediction can take each of the preceding tokens into account as opposed to only the current token. To see this consider an LDS with no state, i.e. A = 0: this extremely special case of an LDS captures every fixed linear predictor. Thus while the VAW forecaster provides a logarithmic regret bound on a potentially unbounded class of predictors, these predictors are restricted to be history-less linear predictors which is a much more restricted class of predictors than LDS predictors against which our algorithm provides a regret guarantee.

2. Inspired by this question, we have now added a discussion of the problematic eigenvalues (highlighted in blue) directly after Theorem 5. We note that we now have an algorithm that is robust to these problematic eigenvalues. However, to provide some intuition on the gap in the spectrum (which we have now included right after the statement of Theorem 5), here is the following: first imagine that all the eigenvalues of A are bounded by 1-\delta. Then the “weight” that input u_{t-t_0} has for y_t is roughly (1-\delta)^t_0 (because it is multiplied by A t_0 many times). When is this weight small? When t_0 is much larger than 1/\delta. When t_0 is much larger than 1/\delta then u_{t-t_0} doesn’t really impact y_t and so it may be forgotten. This is why we have that if the spectrum lies in [0, 1-log(T)/(8T^q)], length generalization is possible. Indeed, letting \delta = log(T)/8T^q and t_0 = T^q (which is much bigger than 8T^q/log(T) for large enough T) we see that when the spectrum of A is smaller than 1 - \delta, after t_0 = T^q many steps we can forget about the previous inputs u_{t-t0}. For the second part of the range: i.e. that the spectrum of A can lie between [1-1/(2T^{5/4}),1] we note that this is a special feature we prove about spectral filtering (this is rather technical but it is very important because it implies that systems that require long memory are able to be learned with spectral filters). The region between the first part and the second part is exactly that the range where the eigenvalues aren’t small enough that they don’t impact y_t and can be forgotten, but also aren’t large enough that the autoregressive component is able to capture them. It’s crucial that the problematic eigenvalues are the non-extremal ones, and that they can be solved with a constant amount of extra computation (ie the second autoregressive term). In most algorithms, the problematic eigenvalues would be the marginally-stable ones (i.e. the ones that approach 1), and there is nothing that can be done about those. This is why the length generalization of spectral filtering is nontrivial at all.

We apologize for our message stating "this is a reasonable summary", this was wrong and we do not agree on regret variants. In our paper, we ask the question of whether there are algorithms which form predictions based on a class of functions which use the previous $L$ inputs and compare nicely had they used the previous $T$ inputs (for $L << T$). Dynamic regret looks at the best arbitrary sequence of comparitors in hindsight. For $x_t \in \mathcal{D}$ we consider: $\textrm{Dynamic Regret} = \sum_{t = 1}^T f_t(x_t) - \min_{(w_1, \dots, w_t) \in \mathcal{D}} \sum_{t = 1}^T f_t(w_t).$ Let’s see why dynamic regret cannot capture asymmetric regret. Let’s suppose the input (note that this is a fixed input vector that comes from the data) $u_t$ is the $t^{\textrm{th}}$ basis vector $e_t$ (all zeros vector of length $T$ with $1$ at the $t^{\textrm{th}}$ element) and suppose we consider linear predictors of the form (constrained by context length $L$): $\sum_{i = t-L}^T V_{i}^{\top} e_{i}.$ Therefore the domain is $x_t = (V_{t-L,t}, ..., V_{L,t})$ where each $V_{i, t}$ a diagonal $TxT$ matrix (note if $t-L <0$ then we set it to $0$). What is the standard (static) regret in this case? First you fix what $L$ is and you consider: $\textrm{Regret} = \sum_{t = 1}^T f_t\left( x_t \right) - \min_{ w} \sum_{t = 1}^T f_t(w).$ We assume that $f_t(x_t) = \ell_t\left( \sum_{i = t-L}^T V_{i}^{\top} e_{i} \right).$ Thus we have $\textrm{Regret} = \sum_{t = 1}^T \ell_t \left( \sum_{l = t-L}^t V_{i,t}^{\top} e_{i}\right) - \min_{ (B_{1}, ..., B_T)} \sum_{t = 1}^T \ell_t(\sum_{l = t-L}^t B_{i}^{\top} e_{i}).$ What is dynamic regret? Now the choice of matrices B can depend on t: $\textrm{Dynamic Regret} = \sum_{t = 1}^T \ell_t \left( \sum_{l = t-L}^t V_{i,t}^{\top} e_{i} \right) - \min_{ (B_{1,1}, ..., B_{T,1}), \dots, (B_{1,T}, \dots, B_{T,T})} \sum_{t = 1}^T \ell_t \left( \sum_{l = t-L}^t B_{i,t}^{\top} e_{i} \right).$ What is asymmetric regret? We now want to compare to predictors which are allowed to use length T>>L , so the sum in the comparison starts at 1 instead of t-L:

(Notice that this CANNOT be captured by dynamic regret since $\sum_{t = 1}^T \ell_t \left( \sum_{l = t-L}^t V_{i,t}^{\top} e_{i} \right)$ is supported only on the last $L-T$ coordinates (the matrices are diagonal and the inputs are the basis vectors) whereas $sum_{l = 1}^t B_{i}^{\top} e_{i}$ can be supported on all the coordinates. Notice that we could also consider asymmetric and dynamic regret: $\textrm{Asymmetric and Dynamic Regret} = \sum_{t = 1}^T \ell_t \left( \sum_{l = t-L}^t V_{i,t}^{\top} e_{i} \right) - \min_{ (B_{1,1}, ..., B_{T,1}), \dots, (B_{1,T}, \dots, B_{T,T})} \sum_{t = 1}^T \ell_t \left( \sum_{l = 1}^t B_{i,t}^{\top} e_{i} \right).$

The fact that asymmetric regret allows a different domain for the predictor and the comparator class means there are instances (depending on the comparator class) where linear regret can always be incurred (as you noted!). However, when you constrain to the same class of predictors then (as we show in our paper) there can surprisingly be hope that you get $O(\sqrt(T))$ regret. Indeed, in our paper we consider predictors of the form: $y_t = \sum_{i = T-L}^T M_{i} u_{i}$ and the comparator class is $y_t = \sum_{i = 1}^T B_{i} u_{i}$, where $B_{i}$ has a special form based on spectral filtering that is shown to be (1) computable and (2) have optimal performance compared to the best LDS predictor in hindsight.

Finally, we hope this response made everything clear but if you are still skeptical we recommend looking at the algorithm in the Baby & Wang (2022) paper and formally spelling out what the domain is and what the algorithm would be outputting and how this would NOT capture our notion of asymmetric regret and therefore their results DO NOT apply to our work (indeed either the predictors would be allowed to base their predictions on full context length or the comparator class would have shortened context length using their definitions and depending on how you want to define the domain, neither of which are acceptable for the purposes of our paper). We appreciate this reviewer's carefulness in understanding the distinctions between these things and we recognize the subtlety.

Responses to Weaknesses

1. This is no longer the case, please see the main note to all authors.
2. In this work, we use the experiments mainly to validate our theory (on LDS) and to highlight the factors at play (context length, memory of the system, etc.). As mentioned in the response to another reviewer, LDS-based models are growing in popularity and perform competitively in practical and diverse real-world applications. It is not the goal of our experiments to convince readers that LDS methods have broad utility and behave well as layers in deep models, but rather to exemplify that a particular LDS method (spectral filtering) has an intrinsic efficient length generalization ability. The simple deep learning task is included as a proof-of-concept to suggest that this natural length generalization appears in such a setting. We consider it an important future work to perform a thorough empirical investigation into the right way to take advantage of this and integrate on real-world tasks, but this is a modeling question and is complementary to our results. In short, the practical utility of LDS-based models is already known – we prove length generalization for a particular LDS prediction algorithm and experimentally hint that its benefit carries over to LDS-based deep models, and leave the integration into real-world solutions for future work.
3. Please see our main note to reviewers regarding the nontriviality of the theoretical results.

Responses to Questions

1. For more general sequence prediction tasks, such as nonlinear dynamics or language modeling, it has been empirically shown that one can design spectral filtering-based models that perform very competitively (see Flash STU paper https://arxiv.org/abs/2409.10489). Our results here demonstrate an inherent length generalization property in the way that spectral filtering predicts (in the sense that the optimal predictor on a short context is still optimal given the full one, which is not at all true for general sequence prediction algorithms). While we are unable to prove that these spectral filtering-based deep learning models achieve the same asymmetric-regret bounds on language modeling, for example, it is natural to expect that using a layer that can generalize results in an overall model with improved length generalization capability. We leave it to future work to ask the serious deep learning modeling questions and investigate the length generalization effects in large models, and we included a synthetic deep learning task (induction heads) as a minor first step in this direction.
2. Indeed, we have designed a new algorithm which uses two autoregressive components that provably is robust to this bad range. Not only do we provide a theoretical guarantee of this, we show experimentally that this approach is able to perform well on those bad regions.
3. The current upper bound for spectral filtering using the full context is O(sqrt(T)). This is why we say we achieve “length-generalization”-- we mean that we achieve the same regret bound despite using much smaller context length.
4. We absolutely agree that a full-scale investigation into these length generalization behaviors on real-world tasks (such as language modeling or dynamics prediction) is crucial. Answering this fully is out of the scope of our current theoretical work: instead, we tried to get a first-pass indication of these behaviors on induction heads, which is synthetic and approachable yet considerably outside the assumptions of our theory. However, as mentioned in a response to reviewer xrE3, the topic of evaluating length generalization empirically gets unclear quickly and varies with modeling decisions and choices of datasets/tasks. At the end of the day, what matters is the behavior at the largest scales when the model is pushed to the limit (such as in large language models), on modes of length generalization that matter in practice (such as summarization of very long documents during deployment). There are prior works (e.g. the Flash-STU paper) that indicate STU models are competitive at such scales, and our current work is complementary in that it proves a precise theoretical reason to expect these already-successful models to length-generalize. Properly verifying this at the useful scales is a third contribution to the story, which we leave to future research – without a large experimental setup, it would be difficult to test this conjecture in any meaningful way, and so we treat the synthetic experiment strictly as a proof-of-concept.
5. Hazan’s work just considers the notion of regret that is standard in the literature. We have added a bit more on this after Definition 3 of asymmetric regret in Section 2.1. Please let us know if you have questions about asymmetric regret or standard regret since these definitions are key to understanding the work.

Since the authors were willing to upload a revised draft, could you post one with the proper ICLR 2025 stylefile? the margins on this paper were already modified to allow significantly more space than the submissions were afforded, and now the current draft even exceeds the page limit on top of that. I think it would be unfair to other submissions to change my score based on revisions that would far exceed the page limit guidelines that all other submissions had to abide by.