Provable Length Generalization in Sequence Prediction via Spectral Filtering

Thank you for your thoughtful and detailed review!

Connection to LLMs Response

You are correct - we have zero theory for LLMs, and unfortunately this is not uncommon, our theoretical understanding of LLM is very limited. We see this as a start of a theory for length generalization for any sequence, and the easiest sequence we could start with is that for linear dynamical systems. It is rich enough to have been used in LLM applications, please see papers [2,3,4], and yet amenable to analysis. In section 4.2 we show that despite our analysis only holds for the basic mathematical model of linear dynamical systems, our techniques actually do apply for more sophisticated signals, showing evidence that the theory can be enhanced in the future to incorporate the language signal.

Empirical Validation Response

Our experiments are proof-of-concept, and are not immediately applicable to LLMs, although we certainly hope they will in the future. We assume you mean comparing the performance of Algorithm 1 and 2 (not 5 and 6?). In fact, we do compare them on page 7, showing that Algorithm 2 is much more robust to the spectrum of $A$. That's a great suggestion to include plots for $L = T^{1/4}$ and smaller in Figure 4. We will do this.

[2] Resurrecting Recurrent Neural Networks for Long Sequences Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, Soham De

[3] Efficiently modeling long sequences with structured state spaces A Gu, K Goel, C Ré

[4] Spectral state space models N Agarwal, D Suo, X Chen, E Hazan

Thank you for a thoughtful and detailed review!

Q1 (Realistic Systems)

This is a good question. On the face of it - linear dynamical systems are a toy mathematical model, and it is unclear if real dynamics are linear or symmetric. However: 1) our bounds do not depend on the hidden dimension, so one can imagine very complex dynamics that we only see their projection. 2) Some previous papers such as [2,3,4], argue that LDS have good modeling capacity for language, even with symmetric transition matrices. 3) Our work extends to very recent advances in spectral filtering by [5], such that it extends to more general linear dynamical systems.

Q2 (Induction Heads Task)

This is only for simplifying the setting, there is no need for this assumption, and we could have experimented with random extensions just as well.

Q3 (Figure 5)

This is just a notation, you are 100% correct that above 1 has no meaning, but the standard mathematical notation is mean+std, which can be larger than 1, so we went with that. Our preference would be to cut off at 1. Missing footnote 3 said "Even though the accuracy cannot go above $1$, the error bars are still well defined above this value." We will make sure this compiles correctly.

References

[2] Resurrecting Recurrent Neural Networks for Long Sequences Antonio Orvieto, Samuel L Smith, Albert Gu, Anushan Fernando, Caglar Gulcehre, Razvan Pascanu, Soham De

[3] Efficiently modeling long sequences with structured state spaces A Gu, K Goel, C Ré

[4] Spectral state space models N Agarwal, D Suo, X Chen, E Hazan

[5] Dimension-free Regret for Learning Asymmetric Linear Dynamical Systems Annie Marsden, Elad Hazan

Thank you so much for your thoughtful and detailed review!

Q1 (noisy setting)

Our theorem extends to the case of stochastic noise with bounded variance. To extend to the case of adversarial noise, the norm of the allowed noise must be bounded with inverse to $T$. We present our results in the noiseless case for simplicity of the presentation and proof, but we can happily provide a corollary or note on this. Note that our theorem is presented in the online regret-minimization setting with respect to the best spectral filtering predictor in hindsight. This predictor gets $O(\sqrt{T})$ regret wrt the best linear dynamical predictor for any signal, regardless of the noise (shown in Hazan, E., Singh, K., and Zhang, C. Learning linear dynamical systems via spectral filtering. Advances in Neural Information Processing Systems, 30, 2017b).

Q2 (Section 4.2)

We think of 4.2 demonstrating that while our theory applies to LDS, it can be more widely applicable even to LLM tasks. It is a proof of concept, and we don't make a strong claim here, we just thought this will be useful for future scientists to explore. We are willing to remove it, the paper has a lot of results. We are also willing to include benchmarking from other models.

Thank you so much for your thoughtful and detailed review!

Q1 (L vs T)

In the paper we distinguish between L and T, where L is the context length and T is the overall sequence length. It is possible to take L=T. You are right that $\phi_i$ has size $T$ but we only allow ourselves to look at $L$ previous inputs and so there is indeed a dimension mismatch if $L < T$. The way we deal with this is by zero-padding the inputs. Specifically, the matrix which stacks $u_t, \dots, u_{t-L}$ would instead stack $u_t, \dots, u_{t-L}, 0, \dots, 0$. We state this on page 1 under ``Our Contributions" but we can make it a larger/more-clarified point.

Q2 (Broader Scientific Literature)

This is a great point, and we can give an answer, which we think will improve the paper. The reviewer observes that often in length generalization, the predictor is only allowed to learn/train on short context length, and then predicts on long context length and, importantly, in order for the prediction to be correct, the algorithm must make use of the full long context. In the online learning/regret setting the analogous setting would be to restrict the learning algorithm to context length $L<T$ when making gradient updates (ie keep the algorithm in our paper as is) but to measure its performance when making predictions on the full context length. We agree that we should include this notion of regret in our paper. Note that our theoretical results immediately apply to this setting, since the spectral filtering predictor is only given more power. Indeed, the regret bounds provided in our paper would hold for this notion of regret and, furthermore, they would be stronger-- they would hold without any assumption on the spectrum of $A$ with respect to $L$. Of course -- the above works only if the signal comes from a linear dynamical system with the properties assumed in the paper. We will also include a small discussion on which signals we believe would distinguish between these two notions of regret (i.e. when is our original notion of regret too pessimistic for an algorithm to succeed and when must we move to the notion where predictions use the full context?).

Essential References Response

We will cite [1] as a practically relevant paper, notice our treatment is theoretical and in a broader setting of sequence prediction for linear dynamical systems, not necessarily language.

Main Weaknesses Response

We suspect the reviewer will have a different opinion in light of our clarification of the above theorem and the fact that our results apply to that setting. However, we are happy to emphasize further that our paper is a theoretical one and considers length generalization for general sequence prediction in dynamical systems, rather than an empirical paper on methods that are useful for LLMs to length generalize.

Other Comments

1. We completely agree and will change this phrasing.
2. Algorithm 1: projection can be applied to any convex set K, but the most common one (and the one we use) is with $K_r$ for some Euclidean diameter bound r.