State space models can express $n$-gram languages
Authors: Vinoth Nandakumar
Qiang Qu
Peng Mi
Tongliang Liu
Venue: Accepted by TMLR ()

Abstract: Recent advancements in recurrent neural networks (RNNs) have reinvigorated interest in their application to natural language processing tasks, particularly with the development of more efficient and parallelizable variants known as state space models (SSMs), which have shown competitive performance against transformer models while maintaining a lower memory footprint. While RNNs and SSMs (e.g., Mamba) have been empirically more successful than rule-based systems based on $n$-gram models, a rigorous theoretical explanation for this success has not yet been developed, as it is unclear how these models encode the combinatorial rules that govern the next-word prediction task. In this paper, we construct state space language models that can solve the next-word prediction task for languages generated from $n$-gram rules, thereby showing that the former are more expressive. Our proof shows how SSMs can encode $n$-gram rules using new theoretical results on their memorization capacity, and demonstrates how their context window can be controlled by restricting the spectrum of the state transition matrix. We conduct experiments with a small dataset generated from $n$-gram rules to show how our framework can be applied to SSMs and RNNs obtained through gradient-based optimization.

Published in Transactions on Machine Learning Research (02/2025)
State space models can express n-gram languages
Vinoth Nandakumar
Sinoth..0@@mmall.ooo com com
Sydney AI Centre, University of Sydney
Qiang Qu
edu.
Sydney Al Centre, University of Sydney
Peng Mi
mipengg01 @gmaiil com
Sydney Al Centre, University of Sydney
Tongliang Liu
tongliang, liu @sydney. edu
Sydney Al Centre, University of Sydney
Reviewed on OpenReview: /bppsrev/pperevvieee net/ pet/forunts id=Q1BaDKb-370
Abstract
Recent advancements in recurrent neural networks (RNNs) have reinvigorated interest in their application to natural language processing tasks, particularly with the development of more elficient and parallelizable variants known state space models (SSMs), which have shown competitive performance against transformer models while maintaining a lower memory footprint. While RNNs and SSMs (e.g., Mamba) have been empirically more successful than rule-based systems based on n-gram models, a rigorous theoretical explanation for this success has not. yet been developed, as is unclear how these models encode the combinatorial rules that govern the next-word prediction task. In this paper, we construct state space language models that can solve the next-word prediction task for languages generated from n-gram rules, thereby showing that the former are more expressive. Our proof shows how SSMs can encode n-gram rules using new theoretical results on their memorization capacity, and demonstrates how their context window can be controlled by restricting the spectrum of the state transition matrix. We conduct experiments with a small dataset generated from n- n-gram rules to show how our framework can be applied to SSMs and RNNs obtained through gradient-basee basod optimization.
1 Introduction
Over the past two decades, breakthroughs in deep learning have revolutionized natural language processing tasks, including language modeling (Radford et al. (2018). Radford et al. (2019)). question answering (Devlin et al. (2018), Yang et al. (2019)), and machine translation ?. While n- n-gram models were used widely for text-word prediction before the advent of deep learning, it was empirically observed that language models using recurrent neural networks (RNNs) were more effective for complex tasks. Although transformer models, with their attcntion-based layers, have since led significant advancements in these tasks Vaswani al. (2017), recent research has revisited RNN architectures Beck et al. (2024); Feng al. (2024); ; Peng al. (2023 and investigated variants known state space models (SSMs). Notably, SSMs Gu et al. 2022 2021
Wang & Xue (2024); Gu & Dao (2024) can be trained efficiently in parallel like transformers, while requiring less memory than transformers for inference.
Recent advances have largely been empirical, with new deep learning architectures developed through experimentation and validated on large datasets Liu et al. (2019); Radford et al. (2018 2019). A key open problem in natural language processing to provide a rigorous theoretical framework for the success of deep learning models next-word prediction tasks. Such a framework could lead to the design of more efficient.

---

Published in Transactions on Machine Learning Research (02/2025)
... o'clock and was'
s.. and was too'
input sequence
... go back to'
hidden state space
output logits
[0, .... 1, ....
...
3
.... ...
'upset, an- gry, elated
'sleep', 'bed'
next words
'too'
Figure 1: This diagram illustrates our framework for encoding n-gram rules with state space models. For example, the blue dots in the hidden state correspond to input sequences ending in "... go back to". The hidden state embedding vectors have been projected two-dimensional space, and the vectors corresponding the same n-gram form clusters. The output logits represent the next-word probabilitics, and here the non-zero values correspond the words that can occur next in the sequence (see Figure 2 for more details about this example of an n-gram model, and Figure 4 for more detailed cluster plot).
architectures for SSMs and RNNs, and provide additional insights into the inner workings of these black-box models. A thcorctical understanding of how SSMs and RNNs obtained through gradient-based optimization techniques encode the rules governing r-gram languages, and solve the next-word prediction lask in this context, would be a first step towards a more rigorous mathematical framework explaining their success in natural language processing.
To theoretically explain the empirical success of these deep learning models, researchers have drawn upon mathematical models based on languages defined by formal grammars Chomsky (1956 1959); Chomsky & Schützenberger (1963) and universal approximation theory WWng & Xue (2024). While simple examples of formal languages like a"b", parity languages, and boolean expressions have been studied with neural networks in Gers & Schmidhuber (2001}; Sennhauser & Berwick (0088): Hahn 2020), these are often too restrictive to capture the nuances of natural language. Another key result is the Turing completeness of RNNs Chung & Siegelmann (2021), which implies that they can recognize recursively enumerable languages. However, the proof uses weights that have unbounded precision (or unbounded memory), which does not explain their success practical applications. Our work is closely related to Sarrof et al. (2024), which thcorctically demonstrates that state space models can express regular grammars, but does not examine the depth and size required to model languages generated by n-gram rules.
In this paper we study the question of constructing state space models that express languages defined via n-gram rules from a theoretical viewpoint. Given a n-grai language model defined for a fixed language, we construct state space model that represents it and bound its size (see Theorem 4.1 for a more precise statement). Our construction provides additional insight into the interpretability of SSMs and RNNs, and in particular the role of the state transition matrix and the neurons the hidden layer. Our proof tochnique is illustrated in Figure above. The contributions of this work be summarized as follows.
2

---

Published in Transactions on Machine Learning Research (02/2025)
Our first key insight is that, in the context of n-gram models, next-word prediction is equivalent to memorizing the 7-gram rules. Based on the observation that the embedding layer transforms the sequence of input words into sequence of vectors, we reduce the problem to demonstrating that SSMs can memorize a set of input-output vectors under mild assumptions.
Our key technical result provides thcoretical analysis of the memorization capacity of SSMs, extending similar findings in feedforward networks Yun et al. (2019)). This result is of independent interest and offers deeper insights into their behaviour.
We analyze the context window of SSMs, particularly focusing on when the network's output depends solely on the most recent n inputs. We demonstrate that this behaviour is closely related to the spectral decomposition of the state transition matrix.
We conduct experiments with toy data consisting of English sentences generated from n-gram languages, and show how our thcorctical insights can be applied to SSMs obtained via gradient-basod optimization. We also discuss how our framework can be extended to RNNs, which incorporate a non-linearity into the hidden state update.
The paper is organized as follows. In Section 2 we describe related work on the theoretical underpinnings of language modelling with neural networks. In Section 3 we recall the definition of n-gram language models and state space models. In Section 4 we state the main thcorctical results that show state space models are at least as expressive as n-gram language models, and outline the proof. In Section we apply our theoretical framework to SSMs and RNNs that are trained with gradient-based optimization on a simple dataset from r-gram rules. In Sections 6 discuss further directions based on our theoretical results. In the Appendix, we present detailed proofs of the main results.
2 Related work
Comparison between neural networks and n-gram models. The experiments in Bengio et al. (2003 using the Brown corpus and the Associated Press News corpus show that langnage models based on feedforward neural networks achieve lower perplexity scores compared to n-gram models (here the perplexity metric quantifies the model's ability to predict the next word in a sequence). In Chelba et al. (2017p, empirical comparisons between RNNs and n-gram models have demonstrated the superior performance of RNNs for next-word prediction, based experiments conducted using the UPenn TreeBank corpus and the One Billion Words benchmark. In Svete & Cotterell (2024), ; the authors theoretically analyze how transformer models can solve next-word prediction for datasets generated from n-gram rules, but their choice of query-key matrices only retain information about the positional encoding and not the word embodding, limiting their applicability to models trained using gradient-based optimization. Our work complements these empirical findings by providing a rigorous theoretical framework explaining how SSMs encode n-gram rules, which can be used to derive insights about the hidden layers and weights of SSMs and RNNs trained on simple datasets with stochastic gradient descent.
Memorization and learning. While memorization refers to the model's ability to fit pre-specified input- output pairs, learning involves capturing the underlying patterns or distributions present in the data, enabling the model to make accurate predictions on new, unseen inputs. The differences between memorization and learning have been analyzed through an empirical lens with gradient- optimization in Arpit et al. (011 7)
The memorization capacity of feodforward neural networks has been studied thcorctically in Yun et al.
(219));
Zhang et al. (017)); Huang & Huang (1991), but their work does not extend to SSMs or RNNs. Our work builds these theoretical results by analyzing the memorization capacity of SSMs, and illustrates how memorization enables deep learning models to encode n-gram rules for next-word prediction.
On the expressiveness of RNNs and SSMs. In Wang & Xue 2024), it is shown SSMs with layer-wise nonlinear activation can approximate any continuous sequence-to-sen tottttossseneee function with arbitrarily small error, building on analogous results for RNNs. In Jelassi et al. (2024), the authors thcorctically analyze the effectiveness of transformers and SSMs on copying tasks. It has also been theoretically demonstrated that both RNNs and transformers are Turing complete, meaning they have the computational power to
3

---

Published in Transactions on Machine Learning Research (02/2025)
simulate any Turing machine given sufficient time and resources Pérez et al. (2021) ;: Chung & Siegelmann
(2021). The connection between RNNs and formal languages can be viewed through the lens of the Chomsky
hierarchy, Chommsy (1956), which classifies formal languages into four categories - regular, context-free, context-sensitive, and recursively enumerable- - based on their generative power. RNNs, due to their Turing completeness, can theoretically recognize languages across this entire hierarchy Chung Siegelmann (2021) but the construction there uses weights that have unbounded procision or unbounded memory, which are not realistic in practice. Our work complements these studies by presenting theoretical framework showing how SSMs can model languages generated from n-gram rules without imposing any unrealistic assumptions, and we show how our theory can be applied to models trained via gradient-based optimization on simple datasets.
Learning formal languages with stochastic gradient descent. There has been a line of work using special cases of formal languages to investigate how they can be learnt with stochastic gradient descent using commonly used neural network architectures, such as LSTMs and transformers [Ackerman & Cybenko (2200) Liu et al. (2023). The effectiveness of LSTM models on the languages a"8" (and variants thercof) was established in Gers & Schmidhuber (2001), Bodén & Wiles (2000), while similar results for Dyek languages were obtained in Sennhauser & Berwick (2018). In the latter context, LSTM models require exponential memory in terms of the input length, and do not Icarn the underlying combinatorial rules perfectly. Analogous experiments with Dyck languages were carried out with transformer-based models in Ebrahimi et al. (020
and the inner workings of these models were studied. However, the simple examples of formal languages studied these papers do not model the English language effectively, and are less realistic than languages based on n-gram rules. While an extensive empirical evaluation was conducted in Delétang et al. (0023 Poel
et al. (2024) by training a variety of neural network models on more complex formal languages, theoretical
understanding of these models is lacking. Our work complements these studies by presenting thcoretical framework showing how SSMs and RNNs can encode n-gram rules, as demonstrate with experimental results using stochastic gradient descent.
3
Preliminaries on language modelling.
In this section we recall the definitions of n-gram languages and state space models, describing how they can be used for next-word prediction. We then give an overview of the main results in our work.
3.1 Language models
In this section, we formalize the language modelling task, which has been studied extensively in empirical work Radford et al. 2018). Radford et al.] (2019) (it is often referred to as "text generation' or "text completion") "). In empirical settings, the language typically consists of sentences about a given set of topics that are semantically correct. We follow the conventions in Svete Cotterell (2024).
Definition 3.1. The vocabulary W is a finite set consisting of all occurring words. Let d = [WV, and W := WU (0}. For a given vocabulary W, a language is a subset of the set of sequences of words from W, and represents the set of all possible sentences. Here the integer N is the maximal length of a sentence.
cc U (W*)"
OSAEN
Definition 3.2. Let w = (Wi, ... wx) E We be a sequence of words. Given isicj S k, define the truncated sequence Way (w,, w)). We say that w E W is valid starting sequence there exists another sequence (Wa+1,". wi) such that (wi,".,, Wk: Wkt1, ... ( E L. Let Lk C W denote the set of all valid starting sequences with length k.
Given an input sequence (wi,".. wk), we can use a language model ʃ to generate the next word # this process can be iterated to generate an output sentence (w ,wN) (see Goodman (200I), Katz (1987)). ). In the below definition, we assume a language model takes as an input a valid starting sequence, as defined above, and its output probability distribution that can be used to determine which words are likely appear next in the sequence.
4

---

Published in Transactions on Machine Learning Research (02/2025)
A woke > at > B o'clock and was too -> C >o go back > to D
Harry, Ron Sirius
five, six seven, cight
upset, angry excited, clated
sleep bed
Figure 2: An example of an language that can be modelled by n-grams, based on text from J.K. Rowling's
'Harry Potter'. The graph illustrates how sentences are formed, with each of the symbols A, B, C, D being replaced by one of the words in the corresponding box. Two examples of sentences from this language are 'Ron woke at seven o'clock and was too upsct to go back to bod." and "Sirius woke at five o o'clock and was too elated to go back to sleep.' See Appendix B for an expanded example, including a full list of n-gram rules.
Definition 3.3. Given a vocabulary W, a language model is a function f that takes in an arbitrary sequence w E W* for some k, and outputs a vector of probabilities f(w) = (f(wlw))wew-. Here (w(w) denotes the probability that the model f will output the word w given the input sequence w.
3.2 n--ram language models
In this section, we rigorously define n-grum language models that we will use subsequent sections. An n-gram language model generates sequences by estimating the probability of each word based on the n precoding words, and capture local dependencics within sequence Katz (1987); Goodman (2001). These models, widely used in natural language processing, can approximate real-world sentences but struggle with long-range dependencies Goldberg & Orwant (013). They are closely related to the theory of formal grammars, which provide more general mathematical framework for generating strings based on set of combinatorial rules Chomsky (1956 1959); Chomsky & Schützenberger (1963).
Definition 3.4. Given language L, an n-gram language model fng consists of the set (Pg (the integer is also fixed).
Let C. (W+))=1 be the of all sequences of consecutive words appearing in L.
Let fng P + A be function, where A = {u e [o, 1j Σ x; 1} is the + I-dimensional probability simplex. The additional + 1-th entry corresponds to unknown token 0.
Given the above data, the n-gram language model füg can be defined as follows (here w E W is any starting sequence). We use the n-gram assumption, which states that the probability of a given word w occurring after sequence w depends only on the last words in the sequence w. If the last n words w W are not in P, the output by corresponds to the unknown token (1
UnngW"ennt+iee 1 1:kk-1))) if - E P, 1ig
fng (w)
otherwise,
Example. Here we give a simple example n-gram language model for the language L specified in the above Figure, with n 4. For each phrase in P, all valid choices for the next word are specified by the graph above, and stipulate that they all equally likely. The below formulas show how the function fng is defined. The sequence ('and", 'was", 'too") is in P, and the vector U'upeet" denotes the one-hot vector in Rd+1 corresponding to the word "upset".
('and", 'was', 'too') (U'upeet + U'angry" + V'elated" + (excited')
Sng("go' "back', 'to')
2 (Ursteop" + Urboa")
5

---

Published in Transactions on Machine Learning Research (02/2025)
We now introduce notion of equivalence between language models below, using the total variation distance between two probability distributions. This extends the definition of "weak equivalence" from Section 2 of Svete & Cotterell (2024) (their definition corresponds to the = 0 setting here). Below we fix language C over a vocabulary W.
Definition 3.5. Given a vocabulary W and a language L, let g be an n-gram language model, and let ʃ be a language model. For a fixed é > 0, we say that f and g are e-equivalent when restricted to L if the following condition holds for each valid sequence w E Ls with 2 n (here note that f(w), g(w) are vectors in Rd+1 and 1-Il denotes the L1 norm of vector).
- <e
3.3 State space language models
We start by recalling the definition of a state space model with a single hidden layer, following Section of Wang & Xue (2024). Below a denotes the ReLU function applicd to each coordinate of the corresponding vector. While it similar to a recurrent neural network, in this setting the a non-linearity is applied while computing the output, and not during the hidden state update. We note that while the bias term by is often not used in state space models, here we add it to the output to increase the expressiveness of the model.
Definition 3.6. A state space model 8 consists of an input layer, a hidden layer, and an output layer. The network is characterized by no imput neurons, n4 hidden neurons, and ny output neurons. The weight matrices are denoted by A, B, and C, where A E Ma, (R) is the state transition matrix, B E Ma (R) is the input-to-hidder matrix, and C E Mataa my (R) is the hidden-to-outpe matrix. We omit the input-output matrix D, which is somctimes used to implement a skip connection.
It gives rise to a function Ss (R*")?" - (R"E)T; the state space model processes a sequence of inputs (11, (2,,2))..,,77) produces a sequence of hidden states (h2, (u,, he...., hap...,,hp hr) and yields outputs (y1, ....... yr) via the following equations. Here I, R0 is the input vector, hr E ma is the hidden state vector, z/ E Rz is the output vector, and by is bias term.
h Ah.- + Br, ye = Ca((h by)
We recall how state space models can be used for next-word prediction in language modelling tasks (sce also Section 4.3 of Gu et al. (2. The two key components that we need are a state space model and an embedding layer. The input to the state space language model is sequence of words, so we start with a preprocessing step that converts sequence of words to sequence of one-hot vectors. The dimension of the vector 8* (w) is proportional to the number of words in W, which is usually quite large. In order to reduce the dimension of the vector before feeding it into the state space model, we define an embedding layer as follows.
Definition 3.7. Let 7: W {0, 1,... d) be bijective map such that i(0) 0, where d = d- We define the word representation 0 W + Rd+1 as the map which sends w E W to the corresponding one-hot vector in [R++1 (Le. with a 1 in the i(w)-th position and 0-s elsewhere). Given an embedding matrix Wo E Mate,,411BB. (R) with distinct columns, denote by Φw, (d++1) + R the induced function. The embedding function F. is defined below; we say that the embedding layer has dimension e.
(W1," WN);
((w)= [ow,
ow,
After the preprocessing step using the embedding layer, which sends to E'w),, the resulting vector can in turn be fed into state space model, which then outputs a vector of probabilities by applying a softmax layer at the end.
Definition 3.8. A state space language model f& is a function obtained by combining an embedding layer E as above, a state space model S, and a softmax layer. Here S has e input neurons, d+ output neurons and one hidden layer. It takes as input a sequence w = (wi, with variable length k, and outputs a probability vector in Id+1
1' (w) softmax Is F(w)
6

---

Published in Transactions on Machine Learning Research (02/2025)
4 Main results
In this section, we state the main results state space language models are more expressive than n-gram language models. In our proof, construct. state space language model that encodes the n-gram rules which govern next-word prediction in this setting, through rigorous theoretical analysis of their memorization capacity. We show how the context window of the state space model can be controlled by imposing restrictions on the spectrum of the state transition matrix. We also discuss how these results can be extended to recurrent neural networks, which incorporate a non-linearity into the hidden state update. We refer the reader to Appendix A for detailed proofs of all results.
4.1 Statement: overview
We present a framework for showing why state space models can represent any n-gram language model, with arbitrarily small error; our main result is as follows.
Theorem 4.1. Let L be a language over vocabulary W. Let füg be n-gram language model, oblained from the datum (P, fng). Given any e > O, there exists a state space language model fž with the following properties:
The state space language model, I3 and n-gram language model Sig are e-equivalent when restricted to C.
The state space model component can be chosen so that it has a single hidden layer with |P| newrons (here [P] is the number of n- n-gram rules in the language L).
The embedding loyer has dimension e, where e 2 1 is any fixed integer.
In the next section, we state a more precise version of the above Theorem which specifies the properties that the state space model must satisfy in order for the above to hold. We note that languages generated from n-gramm rules are regular languages that can be represented by a finite state automata. In Theorem of (Sarrof al. (2024), it is shown that state space models can represent any star-free regular language using different approach based on Frohn--Rhodes theory, but without explicit bounds on the depth and size of the model required (sce also Liu et al. (0023) for constructions of transformers that represent regular languages). Another line of work establishes the Turing completeness of RNNs Pérez et al. (2021): Chung & Siegelmann (2021), but their networks use either unbounded precision or unbounded memory, which limits their application to real-world scttings.
Our approach does not have these limitations, and the next section we explain how our insights can be translated to RNNs and state space models trained on toy English datasets. We note crucially that the number of neurons in the state space model does not depend on the mumber of sentences in the language C, but only the number of rules in the n-gram language L. For any language L, can construct model that solves next-word prediction by simply memorizing all of its sentences, but such a model would not provide much insight into learning, and the number of neurons would scale exponentially with the number of sentences in L. The key feature of the above result is that the model does not simply memorize all of the input data, but rather learns the patterns in the data that are specified in the n-gram language.
We note that our statement above uses an e >>>>> 0 due to the presence of softmax layer. The probability distribution obtained from state space model with softmax layer contains entries that are strictly positive, though they can be made arbitrarily small. However, the probability distribution obtained from an n-gram language model typically contains many zero entries corresponding to words that cannot logically follow a given n-gram. The parameter is used to account for the discrepancy between the two lypes of language models. While can be replaced by zero if we use a hardmax or sparsemax instead the softmax (see Section of /Sete & Cotterell (2024) for definitions), we opt for the softmax function as it is used more widely empirical settings. We also note that our existence results hold for any value of the embedding dimension e, including 1 see the discussion after Proposition 4.5 for explanation of why this is the case.

---

Published in Transactions on Machine Learning Research (02/2025)
4.2 Proof: construction of the state space model.
First we introduce some additional notation specify the algebraic rules that encapsulate the next word prediction task. Recall that the input to language model is sequence of words w' that is valid starting sequence, and the output is vector that indicates which words can be added to the end, that the resulting sequence can be completed to obtain valid sentence. Note that only the last words are needed for the prediction. We formalize this below, assuming that the sequence w' lies in P, the set of all valid n-grams in the language C.
Definition 4.2. Given sequence of words w² {wi,".. (n--1 E P, recall that the vector fng (w') is probability vector in A (recall that d = [W]). For each w' EP, we choose the vector Swr E gd+1 with all coordinates being strictly positive with the property that !| - fng (w (w)ll < c. We also choose a vector Sn with the property that softmax(3
Sw'
Above, the vector Sur can be constructed by replacing all of the zero coordinates in fng(w) with strictly positive values that smaller than c, and modifying the other coordinates accordingly that the entries have sum 1, ensuring that the vector lies in A The coordinates of w can be obtained by taking logarithms of the entries in Sw', noting that the image of the softmax function is preciscly those vectors in A which have strictly positive coordinates.
Our state space model processes a given imput sequence w by converting it to the vector representation E(w) using the embodding layer. We will construct a state space model that maps the vector E(w), corresponding to given valid starting sequence u', to the vector Fu, 1. specifying the probability distribution. This will ensure that the state space language model f3 will be e-equivalent to the n-gram model fäg: due to the bound from Definition words in the input sequence, are needed to predict the next word. Note that Theorem 4.1 follows as a direct consequence of Theorem 4.3 and Proposition 4.4
4.2 above. We formalize this below, noting that only Wk-n+l:k-1 the last
Theorem 4.3. Let S be one-layer state space model, such that the following holds, where w (W1, :* Wk-1) is any valid starting sequence, and E is any embedding function.
fs(((w)) _w
Let f's be the state space language model obtained by concatenating the state space model S with embedding layer E. Then I3s and fäg are c-equivalent when restricted to L.
Proposition 4.4. There erists state space model with e inpull neurons, one hidden layer containing [P] neurons, and d d+1 output neurons such that the following holds (here w (Wi, wk-1) is any volid starting sequence).
fs(((w)) S |
Example. We revisit the example from Section 3.2 to help clarify the notation used in this subsection. Let w be the valid starting sequence below, noting that k = 10 and = 4.
['Ron', "woke", "up", 'at", 'one", 'o'clock', "and', "was", 'too"] W6g Wk-n+ ['and', 'was", "Loo")
To compute Sy, wtik we start with vector that has values 1 for the positions corresponding to the words "upset", "angry, "elated" and "excited", and zero values elsewhere. Then $ga-ptra-1 is obtained by changing all values by an amount smaller than so that all values are strictly positive, and the sum of all entries is still equal to 1. Finally, the vector B 8m.-nt14-1 is obtained by taking the logarithms of the coordinates of the vector 8 which is possible as they are strictly positive.
To prove these results, we outline the two key Propositions that are needed in Sections 4.3 and
SeC
Appendix A for detailed proofs. Theorem 4.3 follows from the definitions in Section and Proposition
follows as direct. consequence of Proposition 4.5 and Proposition
4.6
8

---

Published in Transactions on Machine Learning Research (02/2025)
4.3 Proofs: on the memorization capacity of state space models.
To construct a state space models which satisfies this properly, we analyse their memorization capacity in the below proposition. This is an extension of existing results on the memorization capacity of feedforward neural networks Zhang et al. (2017), which show that they can fit an arbitrary finito set of input-output vectors (see also Huang & Huang [1991), Yun et al.] (2019)). ). Our approach is constructive in nature, with the parameters constructed algebraically so that the equalities are satisfied. Due to the recurrence used to compute hidden states in this context, the proof is more subtle than the analogous construction for feedforward neural networks from Theorem in Zhang et al. (017). See Appendix A.1 for detailed proof.
Proposition 4.5. Let (ai,,i) be a fonile set of inpul-oulpul pairs, where each impul sequence x; = (21,1,... ,i,n) consists of n vectors with zij E R2, and ys E R2, Assume that the input sequences are distinct, i.e. I2 f zj for i i j. Then there erists state space model S with single hidden layer with K neurons that can memorize this data. Specifically, there crist weight matrices B E R A E R C E RKxg and bias vectors by E TR9 such that for each sequence of inpul Fi, the state space. model produces the corresponding output yi- Further, the state transition matric A can be chosen so that = 0.
The key technical step is Lemma A.2 which show that, for a generic choice of weight matricos A and B, the hidden state vectors corresponding to the inpuls are linearly independent. To prove this, first Lemma A.3 shows that for generic choice of A and B, the j-th coordinates of the hidden state vectors are pairwise distinct for cach S j S K. Using this property, Lemma A.4 constructs bias terms by such that the hidden states obtained after applying the ReLU activation are linearly independent. Using the linear independence property, we complete the proof of Proposition 4.5 by showing there exists a unique fully connected weight matrix C such that the outputs of the SSM are the vectors yı. In order to apply this proposition in our context, we will need to impose the additional condition that the state transition matrix A is nilpotent of order n, and our method can be employed with this constraint. See Section 6 for a discussion of how these techniques can be used to obtain optimal bounds on the memorization capacity of state space models.
We also note that Proposition 4.5 holds for any imput dimension p: this is used to show that Theorem 4.1 holds for any value of the embedding dimension We do not need any restrictions on p in this setting, as our proof shows that Lemma A.3, which states the individual coordinates of the hidden state vectors are pairwise distinct, holds generically (even when p = 1). The bias term by is then used to ensure that the resulting hidden state vectors are linearly independent. We note that our SSM model will be less expressive if the bias term is not used, and restrictions on p and e may be required to prove analogous results in that sctting.
4.4 Proofs: on the context window of state space models.
Our second key result establishes how we can restrict the state space model's memory to a fixed context window by imposing restrictions on the state transition matrix. This is crucial for modeling n-gram languages, where the prediction of the next word only depends on the preceding n tokens. The result. is formalized in Proposition 4.6 below: if the state transition matrix of the state space model is nilpotent of order n, then the state space model's output at any time step depends only on the last input tokens. We note that restrictions on the eigenvalues of the weight matrix has been explored as means of mitigating exploding/vaniihing gradients in Helfrich & Ye (020)) Kerg et al. (2019)
Proposition 4.6. Let S be a state space model such that the state transition matriz A is nilpotent of order n. Then given an input sequence I = ("1,""",TT)) the output fs(x) only depends on the last n imput values (rr-nt1.'". ,TT).
Our proof involves linear algebra calculation with powers of the state transition matrix, obtained by using the convolutional view of the state space model (see Section 2.1 of Wang & Xue (2024) for more details). This formula expresses the hidden state by at time T as sum of past inputs multiplied by power of the state transition matrix. When the state transition matrix A satisfies A = O, the terms corresponding to inputs older than hy-n vanish, and by depends only on the most recent n imputs (IT, IT---11) This shows that constraining the state transition matrix is nilpotent with order n effectively bounds to the context. window of the model to have length n. See Figure 3 below for an illustration of the proof, and Appendix A.2 for detailed proof.
9

---

Published in Transactions on Machine Learning Research (02/202))
21
TT-3
2T-2
2T-1
TT
A
A
A
h
hr
A
A
A
YT
Figure This figure shows how the context window state space model can be bounded using nilpotent state transition matrix A. In this example, A = shown by the arrows which represent the action of A the individual neurons, which map each neuron to a neuron with larger index. These neurons represent the Jordan basis for the nilpotent, and the arrows indicate how all inputs prior to 2T-2 do not influence the output yr-
4.5 Extensions to recurrent neural networks.
Recurrent neural networks (RNNs) and state space models both rely on sequential updates of a hidden state to process input sequences, with the hidden state at each time step being determined by the previous hidden state and the current imput. The key difference is that the hidden state update for RNNs uses non-linearity and bias terms. Unlike the S3 model which can compute multiple time steps simultaneously due to the linear hidden state update, the output a RNN cannot be parallelized during training Gu et al. (2021). We now outline how the key results of this section can be extended to recurrent neural networks.
Our results in Proposition 4.6 on the context window of state space layers can be extended this setting with one key modification. In addition to the hidden weight matrix A satisfying the nilpotency constraint A" = 0, we also impose condition that it is block nilpotent with respect to the standard basis formed by the hidden neurons. This is imposed to ensure that the piecewise linear map o A satisfies the same relation; see Appendix A for more precise definition and proof. We expect that Proposition 4.5 can also be extended to this setting; the key step is to establish the linear independence of the hidden state vectors obtained from each of the input sequences, noting that in this setting the ReLU function is applied during the hidden state update.
5 Experimental results: learning n-gram rules with state space models.
In this section, we conduct experiments with state space models trained on datasets generated synthetically from n-gram rules using the example in Section and show that they capable of predicting the next word with near-perfect accuracies. We study the spectral decomposition of the hidden weight matrices, and use it to analyse the context window of the state space model. We also illustrate how the hidden state vectors in the state space model encode n-gram rules, by visualizing the vectors with a cluster plot using principle component analysis. These results show how our thcoretical framework can be used to gain insights into the inner workings of state space models trained on simple datasets using gradient-based optimization.
5.1 Experimental set-up
We generate data using the language which defined Appendix B using a list of n-gram rules, and explained with a simplified diagram in Section 3.2] To generate sentence from L, see the illustration in Figure 2
(note however that this diagram does not display the full set of rules, but they listed in the
10

---

Published in Transactions on Machine Learning Research (02/2025)
table from Appendix B). The language L consists of hundreds of sentences that are grammatically correct. From the discussion in Appendix B and the example in Section 3.2 an n-gram language model fng with n = 4 and 145 valid trigrams in (P] can be used to solve next-word prediction for this language L. We train a state space language model for next- word prediction on this dataset. Following the setup in Section 4.3, the model consists of an embedding layer with dimension 10, a hidden layer with 145 hidden neurons, and a fully connected layer that produces the output logits. We note that the size of the hidden layer equal to the number of valid trigrams in our model, following the setup in Theorem 4.1
Instead of using pre-trained word embedding (based on algorithms such word2vec Mikolov et al. 20113a by or GloVe Pennington et al. (2014)). ) in our experiments we train the embedding layer jointly with the state space model. While our theoretical construction works with an embedding layer with arbitrary dimension e e2 1, empirically we find that stochastic gradient descent does not converge when the dimension e is very small. The model is trained for 25 epochs with a fixed learning rate of 0.001, using stochastic gradient descent. with the cross-entropy loss function. We evaluate the model's accuracy on an unscen test set; each sentence in the test set is truncated to length of 6, and the model's completion is evaluated by checking if the resulting sentence lies in our dataset. After replicating the experiment 5 times with different random seeds, we find the model consistently achieves near-perfect accuracies when using a training dataset with 40 sentences.
5.2 Spectral decomposition of hidden weight matrices
In this section, revisit the theoretical framework in Proposition 4.6 and apply these insights to state space models that are obtained through gradient-basct optimization. The key feature of our thcoretical construction is that the state transition matrix is nilpotent, i.. that it has zero eigenvalues. The below plot shows the eigenvalues derived from the state transition matrix of the recurrent language model. Fach blue dot represents an eigenvalue, plotted in the complex plane with the real part on the x-axis and the imaginary part on the y-axis, while the red dashed circle indicates the unit disk.
In an empirical setting, we note that the nilpotency condition on the matrix must be relaxed; instead of stipulating that power of A is exactly O, we check whether it is approximately equal to 0. Using the eigen-decomposiiioon of matrix, it is well-known that this will be true for n sufficiently large precisely if all of the eigenvalues have norm less than 1. The diagram shows that the eigenvalues all lie within the unit circle, with only one exception, which supports our hypothesis.
To exlend Lemma 4.6 to RNNs, as described in Appendix A, we impose the additional condition that the hidden weight matrix is block nilpotent. This implies that (ar a A)" is the zero map; note here o A is a piecewise lincar function. Unlike our observations for SSMs above, we find that for RNNs the matrix A has multiple eigenvalues with values greater than 1. However, by computing the norm of vectors (a o B(w) for varying input sequences w. we find that nevertheless the map (a o A)" approaches O as n increases. Using these observations based on the theory in Lemma [.6] bounding the Lipschitz norm of the map o o A allows us to estimate the effective context window of an RNN (here we use Lipschitz norm instead of operator norm as this map is non-linear).
5.3 Hidden embeddings for n-grams
In this section, we revisit the theoretical framework from Proposition 4.4 and apply these insights to the state space models obtained through gradient-based optimization. The key feature of our theoretical construction there is that given an input soquence of words w. the image in the hidden layer depends only on the n-gram obtained from the last words of the sequence; in other words, the n-gram rule is encoded as a vector in the hidden layer. Our findings build on earlier study in Karpathy et al. (2015), which investigates the internal mechanisms of LSTMs and finds interpretable behaviors in individual neurons that track specific attributes, such as quotes or parenthesis.
To evaluate this hypothesis, below we apply principal component analysis (PCA) to compress the 64- dimensional hidden embedding vectors obtained from the state space model to 2-dimensional vectors that we visualize in Figure I above. Each point in the scatter plot represents a hidden state corresponding to a specific input sentence, with the colors of the points corresponding to the n-gram consisting of the last.
11

---

Published in Transactions on Machine Learning Research (02/2025)
Eigenvalues the Hidden-Hlideen Weight Matrix
N-gram embeddings
25
1.00
Figettaaions
Labels
Cride
l'and, "wass, 'too')
0.75
l'eight', 'o'clock', 'and')
15
l'and', 'putt, 'on')
PCA Component 2
['on', 'his', 'jeans') l'o'clock'' 'and', l'put', 'on', 'his')
0.25
5
0.00
-5
-0.50
-0.75
-15
-1. 1..000077 -0.50-00000000 10.00000 -0.505 0:00 0.25 0.50 0.75 1.00
PCA Companent
Real Part
(a) Figenvalues of the matrix A
(b) Embeddings Plot
Figure 4: Plot showing the eigenvalues of the matrix A, and hidden state embeddings corresponding to T-grams.
three words. This visualization illustrates the state space model's ability to encode n-gram rules, vectors corresponding to the same n-gram are clustered together. Since input sequences ending in the n-gram rule have similar representations in the hidden layer, the output logits will also have similar values, enabling the state space model to predict the next word accurately.
5.4 Extension to transformer models
Transformers employ self-attention mechanisms, which enable them to directly attend to all tokens in a sequence, making them highly effective for long-context language modeling. It is known thcoretically that transformers are more expressive than r-gram models Svete Cotterell (2024), but all of the attention heads in their construction use query-key matrices that only use positional encodings and the word embeddings. Empirically it has been observed that while some attention heads are positionally- most of them capture semantic information; trained on the TinyStories dataset. In Elhage et al. (2021), by studying altention heatmaps il has been observed that transformers learn skip-gram rules, rather than 7-gram rules (see also Nguyen 2024h, where skip-gram rules are defined with the character and extracted for transformers trained on TinyStorics). To address this, it would interesting to extend our theoretical framework to analyze how transformer models encode skip-gram rules. Existing results on the memorization capacity of transformer models Mahdavi et al. (2024) can be used substitute for Proposition 4.5 this context. While the nilpotency constraint from Proposition 4.6 isn't applicable to transformers, different attention heads in the transformer could encode different skip-gram rules.
Section 5.1 Eldan & Li 2023) for analysis of transformers
6 Discussion: limitations and further directions.
Tight bounds on the memorization capacity of state space models. In our proof of Proposition 4.5 extend the construction for feedforward neural networks with single hidden layer from Theorem of Zhang et al. (217)) to state space models. The construction uses two-layer network with K hidden nodes and (K3) parameters to memorize K input-output pairs. In Section 3 of Yun et al. (2019), it is shown that a three-layer leedforward neural network with O(K) parameters can memorize K input-output pairs, and that this construction optimal (i.e. the network had less than K parameters, cannot memorize an arbitrary set of K input-output pairs). It would be interesting extend the techniques from Yun al. (2019) to our setting, and construct state space models with two hidden layers and O(K) parameters that memorize any set of K input-output pairs. Due to the recurrence used to compute hidden states for the
12

---

Published in Transactions on Machine Learning Research (02/2025)
state space model, the analogous construction for feedforward neural networks from Section 3 of Yun et al. (2019) does not carry over verbatim, but we expect that it can be adapted to this context.
Generalizations of n-gram languages. The present work shows that using the simplified models based on n-gram languages can provide insights into the workings of state space models and RNNs trained on next-word prediction tasks. However, it well-known that n-gram languages cannot effectively model the English language in its entirety, as the number of n-gram rules required to do so will be prohibitively large Liu et al. (0024). More intricate models based formal grammars Chomsky (1956) 1959) Chomsky & Schützenberger (1963), which have been developed by the mathematical linguistics community, are better suited for complex real-world datascts. While existing results on the Turing completeness of RNNS and transformers show the existence of models that can solve the recognition problem for formal grammars Pérez lt al. (2021); Chung & Siegelmann 2021), the construction uses unbounded precision (or memory), which does not explain their success in real-world scenarios. It would be interesting to extend our construction to formal grammars and use our framework to gain further insights into the inner workings of state space models and RNNs obtained through graseent-based optimization, (see Yang et al. (2024] Bhattamishra et al. 2020): Chiang et al. (2023); Strobl et al. (2024) for analogous results on transformers).
Generalized state space models. While the state space models we study struggle with maintaining information over long sequences, more recent work on generalized state space models extends this by incorporating the HiPPo matrix Gu et al.] (2022), which helps them efficiently represent and update the state information to capture both recent and distant past inputs. While this may not be needed for n-gram languages where the model only needs to consider the most recent tokens, formal grammars require the model to remember recursive relationships between different parts of sentence and maintain consistency across nested clauses. It would be interesting to extend our theoretical framework to generalized state space models Gu et al. (2022); Gu & Dao (2024), and gain a deeper understanding of the interpretability of these models which are obtained through gradient-based optimization.
Our empirical results indicate that state space models and RNNs that are trained on synthetic language data generated from n-gram rules can achieve near-perfect accuracies. This leads us to the question of whether we can obtain theoretical guarantees showing that stochastic gradient descent yiclds state space models and RNNs with near-zero generalization error in this context. There has been existing work on generalization bounds for RNNs using the PAC-Learning framework Chen et al. (2020): Zhang et al. (2018), and theoretical guarantees for the convergence of stochastic gradient descent with RNNs in the overparametrized setting Allen-Zhu et al. (2019)) (building on earlier work for feedforward neural networks Allen-ZZhn Zhu et al. (2019a); Zou et al. (2020) ). However, these theoretical guarantees are obtained for tasks that involve memorizing input-output pairs, and cannot. be directly applied to distributions that model real-world settings, such as data generated from n-grams models. It would be interesting to extend their theoretical guarantees to state space models and RNNs using datasets generated
Stochastic gradient descent and overparametrizatiinn
from n-gram languages, and analyze theoretically whether or not overparametrization is required in this setting.
7 Conclusion.
In this paper, we use languages generated from n-gram rules to theoretically understand how neural networks excel at language modelling tasks. We focus on next-word prediction using state space language models, and show theoretically that they are more expressive than n-gram language models. In our proof, we construct a state space model that encodes the n-gram rules which govern the next-word prediction task, and discuss how our results can be extended to recurrent neural networks. Our key theoretical results inchude a rigorous analysis of the memorization capacity of state space models, and demonstrate how the context window can be regulated through the spectral decomposition of the state transition matrix. We conduct experiments with a small dataset to illustrate how our framework yields insights into the inner workings of state space models obtained through gradient- based optimization. We do not anticipate any negative societal impacts, as the present work is theoretical. Our results shod light on the thcoretical underpinnings of language modelling with neural networks, and we expect that it will provide insights into the inner workings of these black- box models.
13

---

Published in Transactions on Machine Learning Research (02/2025)
Acknowledgments
Vinoth Nandakumar was partially supported by an RTP scholarship at the University Sydney. Peng Mi was partially supported by an OPPO Postgraduate Research Award Computational Science. Tongliang Liu is partially supported by the following Australian Research Council projects: F'T201003318, DP220102121, LP220100527, LP220200949, IC190100031. The authors would also like to thank the anonymous reviewers and the action editor for their detailed feedback.
References
Joshua Ackerman and George Cybenko. A survey of neural networks and formal languages. arXiv preprint arXiv:2006.0133 2020.
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International conference machine learning, pp.
PMLR, 2019a.
Zeyuan Zhu, Yuanzhi Li, and Zhao Song. On the convergence rate of training recurrent neural networks. Advances in neural information processing systems, 32, 2019b.
Devansh Arpit, Stanislaw Jastrzebski, Nicolas Ballas, David Krueger, Enmanuel Bengio, Maxinder Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization deep networks. In International conference machine learning, pp. 233 242. PMLR, 2017.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Günter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024. URL
04517
Yoshua Bengio, Réjean Ducharme, Pascal Vincent, and Christian Jauvin. neural probabilistic language model. Journal of Machine Learning Research, 11117-11555) 2003.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers recognize formal languages. In Proceedings of the 2020 Conference Empirical Methods in Natural Language Processing (EMNLP), pp. 7096 7116. Association for Computational Linguistics, 2020.
Mikael Bodén and Janet Wiles. Context-free and context-sensitive dynamics in recurrent neural networks. Connection Science, 2000.
Ciprian Chelba, Mohammad Norouzi, and Samy Bengio. N-gram language modeling using recurrent neural net-
work estimation. Technical report, Google Tech Report, 2017. URL com/media/reseaacce google.coooooooooooooooooo wwt.com.eeeeee http:///aaaa.aa/a//////////a//aa
Minshuo Chen, Xingguo Li, and Tuo Zhao. On generalization bounds of family of recurrent neural networks. In Proceedings of the 23rd International Conference on Artificial Intelligence and Statistics (AISTATS), PMLR 108, Palermo, Italy, 2020.
David Chiang, Peter Cholak, and Anand Pillay. Tighter bounds on the expressivity of transformer encoders. In International Conference on Machine Learning, pp. 5544-5562. PMLR, 2023.
Noam Chomsky. Three models for the description of language. IRE Transactions on Information Theory, 1956.
Noam Chomsky. On certain formal properties of grammars. Information and Control, 2(2), 1959.
Noam Chomsky and Marcel-Paul Schützenberger. The algebraic theory of context-free langnages. In Computer Programming and Formal Systems, Studies in Logic. 1963.
Stephen Chung and Hava Siegelmann. Turing completeness of A duances in neural information processing systems,
recurrent neural networks. 2021.
14

---

Published in Transactions on Machine Learning Research (02/2025)
Grégoire Delétang, Anian Ruoss, Jordi Grau-Moya, Tim Genewein, Li Wenliang, Elliot Catt, Chris Cundy, Marcus Hutter, Shane Tegg, Joel Veness, and Pedro A. Ortega. Neural networks and the chomsky hierarchy. In International Conference on Learning Representations, 2023.
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. ArXiv preprint 1rXi::8800004000 2018.
Javid Ebrahimi, Dhruv Gelda, and Wei Zhang. How can self-attention networks recognize dyck-n languages? In Findings of the Association for Computational Linguistics, 2020.
Ronen Eldan and Yuanzhi Li. Tinystories: How small can language models be and still speak coherent. english? arviv preprint arX:::335500755, 2023.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, et al. A mathematical framework for transformer circuits. 2021. Transformer Circuits Thread.
Leo Feng, Frederick Tung, Mohamed Osama Ahmed, Yoshua Bengio, and Hossein Hajimirsadegh. Were rmns all we needed?, 2024. URL https:///axxiii wwtpa//wwww....//////...0.......
Felix A. Gers and E. Schmidhuber. Lstm recurrent networks learn simple context-free and context-sensitive languages. IEEE Transactions on Neural Networks, 12(6), 2001.
Yoav Goldberg and Jon Orwant. dataset of syntactic-ngrams over time from a very large corpus of English books. In Mona Diab, Tim Baldwin, and Marco Baroni (ods.). Second Joint Conference on Lerical and Compulational Semanlics (*SEM), Volume 1: Proceedings of the Main Conference and the Shared Task: Semantic Textual Similarity, pp. 241 247, Atlanta, Georgia, USA, June 2013. Association for Computational Ttrg/S3--1055
Linguistics. URL htttps://cclantholo.g
Joshua Goodman. A bit of progress in language modeling. Technical Report MSR-TR-2001-72 Microsoft. Research, 2001.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2024. URL
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and contimuous-time models with lincar state space layers. Advances in neural information processing systems, 34:572-558 585, 2021.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. International Conference on Learning Representations, 2022.
Michael Hahn. Thcoretical limitations of self-attention in neural sequence models. Transactions of the Association for Compulational Linguistics, 8:156 171, 2020.
Kyle Helfrich and Qiang Ye. Eigenvalue normalized recurrent neural networks for short term memory. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 4115 4122, 2020.
Shil-Chi Huang and Yih-Fang Huang. Bounds on the number of hidden neurons in multilayer perceptrons. IEEE Transactions on Neural Networks, 2(1), 1991.
Samy Jelassi, David Brandfonbrener, Sham M. Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying- In Proceedings of the 41st International Conference on Machine Learning, volume 235 of Proceedings of Machine Learning Research, pp. 21502 21521. PMLR, 21-27 Jul 2024.
Andrej Karpathy, Justin Johnson, and Li Fei-Fei. Visualizing and understanding recurrent networks. arXiv preprint Xrxi::1006.00 02078, 2015.
Slava Katz. Estimation of probabilities from sparse data for the language model component of a speech recognizer. IEEE Transactions on Acoustics, Speech, and Signal Processing, (33--3533 March 1987.
15

---

Published in Transactions on Machine Learning Research (02/2025)
Giancarlo Kerg, Kyle Goyette, Maximilian Puelma Touzel, Gauthier Gidel, Eugene Vorontsov, Yoshua Bengio, and Guillaume Lajoie. Non-normal recurrent neural network (nnrnn): learning long time dependencies while improving expressivity with transient dynamics. In Н. Wallach, Н. Larochelle, A. Beygekzimer, F. d' Alchó--Buc, E. Fox, and R. Garnett (ods.), Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019. URL http:tttptpppceeeeeiiieiee 0005://1188a000a77777aaa87777770
wwwriis.cc/pppppe files/paper/
pdf
Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In International Conference on Learning Reeresentatiinss (ICLR), 2023.
Jiacheng Liu, Sewon Min, Luke Zettlemoyer, Yejin Choi, and Hannaneh Hajishirzi. Infini-gram: Scaling unbounded n-gram language models to trillion tokens. In Conference on Language Modeling (COLM). COLM, 2024.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: robustly optimized bert pretraining approach. ArXiv preprint Arxi::1907. 11692, 2019.
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. International Conference on Learning Represenlations, 2024.
Tomas Mikolov, Kai Chen, Greg Corrado, and Jeffrey Dean. Efficient estimation of word representations in vector space. In International Conference Learning Representations, 2013a.
Tomas Mikolov, Iya Sutskever, Kai Chen, Greg Corrado, and Jeff Dean. Distributed representations of words and phrases and their compositionality. In Advances Neural Information Processing Systems, 2013b.
Timothy Nguyen. Understanding transformers via n-gram statistics. In Advances in Newral Information Processing Systems, 2024.
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, and Stella et al Biderman. RWKV: Reinventing RNNs for the transformer era. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Findings of the Association for Computational Linguistics: EMNLP 2023, pp. findings-emnlp..30 December 2023. Association for Computational Linguistics. doi: 10. /21///2222.nnmo
Singapore,
URL
2rw.thooooogyorg/2233..... fendiiiindnssss emnlp.90 .966
Jeffrey Pennington, Richard Socher, and Christopher D. Manning. Glove: Global vectors for word representa- tion. In Proceedings of the Conference on Empiricul Methods in Natural Language Processing (EMNLP),
2014.
Sam Van Der Poel, Dakotah Lambert, Kalina Kostyszyn, Tiantian Gao, Rahul Verma, Derek Andersen, Joanne Chau, Emily Peterson, Cody St Clair, Paul Fodor, Chihiro Shibata, and Jeffrey Heinz. Miregtest: A benchmark for the machine learning of regular languages. Journal of Machine Learning Research, 25
(283):1 45, 2024.
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is turing-complete. Journal of Machine Learning Research, 2021.
Alec Radford, Karthik Narasimhan, Tim Salimans, and Πya Sutskever. Improving language understanding by generative pre-training. Technical report, OpenA1, 2018.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, and 1lya Sutskever. Language models are unsupervised multitask learners. Tochnical report, OpenAI, 2019.
Yash Sarrof, Yana Veitsman, and Michael Hahn. The expressive capacity of state space models: A formal language perspective. Advances in Neural Information Processing Systems, 2024.
Luzi Sennhauser and Robert C. Berwick. Evaluating the ability of Istms to learn context-free grammars. In Proceedings of the EMNLP Workshop BlackborNT.,, pp. 115-124, 2018.
16

---

Published in Transactions on Machine Learning Research (02/2025)
Lena Strobl, William Merrill, Gail Weiss, David Chiang, and Dana Angluin. What formal languages transformers express? a survey. Transactions of the Associalion for Compulational Lánguisties, 12:543-50 561, 2024.
Anej Svete and Ryan Cotterell. Transformers can represent n-gram language models. Tn Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies (Volume 1: Long Pupers), pp. 6845 6881, Mexico City, Mexico, June 2024. Association for Computational Linguistics.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukass Kaiser, and Illia Polosukhin. Attention is all you need. In Aduances in Neural Information Processing Systems, 2017.
Shida Wang and Beichen Xue. Stato-space models with layer-wise nonlincarity are universal approximators with exponential decaying memory. Advances in Newral Informalion Processing Systems, 36, 2024.
Andy Yang, David Chiang, and Dana Angluin. Masked hard-attention transformers recognize exactly the star-free languages. In The Thirty-eine eighth Annual Conference on Neural Information Processing Systems, 2024.
Zhilin Yang, Zihang Dai, Yiming Yang, Jaime G. Carbonell, Ruslan Salakhutdinov, and Quoc V. Le. Xlnet: Generalized autoregressive pretraining for language understanding. In 33rd Conference on Neural Information Processing Systems, 2019.
Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small relu networks are powerful memorizers: tight analysis of memorization capacity. In Advances in Neural Information Processing Systems, 2019.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
Jiong Zhang, Qi Lei, and Inderjit S. Dhillon. Stabilizing gradients for deep neural networks via efficient svd parameterization In Proceedings of the 35th International Conference on Machine Learning (ICML), 2018.
Difan Zou, Yuan Cao, Dongruo Zhou, and Quanquan Gu. Gradient descent optimizes over-parameterizeel deep relu networks. Machine learning, 109:467 492, 2020.
17

---

Published in Transactions on Machine Learning Research (02/2025)
A Appendix A
A.1 On the memorization capacity of state space models
In this section, we prove Proposition 4.5 First we state the below Lemma, which is the key ingredient.
Definition A.1. Given an input sequence I (z1,-... ,.,) to state space model S, recall that the hidden state sequence (hi,""",,, hh) with ha E R are computed as follows (here 1 S S n):
he = Ah + Bx
We define the functions hs, hi as follows, which take in an input sequence and output a hidden state vector. Here recall that a is the ReLU function and by is the bias term.
hg(x)= hnn hs(x) olhn by)
Lemma A.2. Let {xi)e= be finâte set of imput sequences, where each imput sequence 2; = (Ia,1,'.. Finn) consisls of veclors with vijj E RD. Assume that these impul. sequences are distincl, i.e. 34 $ 2j for i j. Then there exist weights for a state space model S with a single hidden layer with N newrons with the following property: the vectors (hs(x;)) are linearly independent. Further, the state transition matriz can be chosen so that =.
We note that above, while the input sequences { are distinct (2, * Ij for * j), some of their coordinates may be equal (e.g. it possible that the first coordinates of x; and Iy are equal, provided that another pair of coordinates are distinct). In order to prove the above Lemma, will use the following two Lemmas.
Lemma A.3. Given set of input sequences { as above, there erists stade space model S with single hidden layer with K neurons with the following property: for each 1 SSj S. K, the j-th coordinates of are pairwise distinct (in other words (,(h((z)) + 6,(h((n)) for all 1 si<i i' < K, where Øy RK + R is the function that extracts the j-th coordinate). Further, the state transition matrir A of the state space model can be chosen so that A = 0.
the vectors {hs(r;))
Proof. Recall that by expanding the hidden state recurrence relation iteratively using the convolutional view of the SSM, we can express hx();) as follows.
n-
hs(x;)
Σ A Bain-- k=0
We claim that for a generic choice of the matrices and B, for each j the j-th coordinates of the vectors (hs(x;))< pairwise distinct. To this, below consider the j-th coordinate of the difference hs(x)) - hs(s)), where i + v:; recall oj R RR is the function that extracts the j-th coordinate.
<s(hs(x;) hs(x:)) = Σ
- =x)
K=
Each of the terms (A B(T,n - xx)) polynomial in the coordinates of the matrices A and B with degree ie + 1. Further, it is sum of monomials, that all the coefficients have value 1. Since i + i', the set S defined below is non-empty.
S= Fi,n -z-n-k 0}
It follows that for each E S, the coordinate (A* B(Ii - FV. -x)) is non-zero for a generic choice of ) is also non-zero for a generic choice of A and B, since the degrees of the multivariate polynomial constituents are distinct. This statement remains true for generic choice of nilpotent matrix A satisfying A" O. This true because any such matrix has the property that all coordinates of 4 are generically non-zero when n.
matrices A and B. It follows that their sum Φ, (z,) - hs(xx))
18

---

Published in Transactions on Machine Learning Research (02/2025)
Lemma A.4. Let {Us) be set of vectors, with v, RR satisfying the following property: for each J S n, the j-th coordinates of the vectors [w:]er are pairwise distinct (in other words, (z(r;) & (p(())) for all S < i' S K). Then there exists a vector by E RK such that the vectors {o + by) are linearly independent.
Proof. The vectors (a(u; by)) are lincarly independent preciscly if the determinant of the matrix A obtained by concatenating these vectors is non-zero. We proceed by induction to construct the vector by E TRK The K = 1 case is clear, given any vector E R we can choose bias by so that the determinant o(11 + by) # O. For the induction step, we assumne that the statement is true for every set of K vectors with the stated property.
Recall that the determinant of the K K matrix A can be computed using the Laplace expansion along the first row. Here di is the (i,j)-th entry of A, and where Aij is the (K - 1) (K 1) submatrix obtained by removing the first row and the j-th column of A.
K
det A)= = E j=1
a1 a (e1(1)))))
Note anj = abyy, + U2,1), where by = (by.1.-.. .y. and uy = (Uy.,,... 2J,N) are vectors in R.V. Since the values 1, are pairwise distinct, we can choose by.1 so that there exists a unique value of for which alyy1 + fj,1) > 0, and (by,1 + uyr,)) = for all j' j With this choice of parameters, the determinant can expressed as follows.
det(A) = (-1)++/
(y.)dee(AA))
By using the induction hypothesis, we can choose the bias terms (by,2) .y.k) so that det(Aig) # 0. Here note that (Ay()1)) corresponds to the set of K-1 vectors in RK obtained by delcting the first coordinates from the vectors (",,-", ; Ug-1, ""+ ,KK. It follows that for this choice of bias by, det(A) O, completing the induction step.
Proof of Lemma A.2 First we use Lemma A.3 to choose the weight matrices A, B for the state space model S so that, for each j, the j-th coordinates of the vectors (hg (x;)) are pairwise distinct (A can also be chosen be nilpotent with An = 0). Noting that (x,) = a (h)(2;)) + by), we can now use Lemma A.4 linearly independent.
choose the bias vectors by that the vectors {
The proof of Proposition 4.5 now follows from the linear independence statement in Lemma A.2 by choosing the weights of the fully connected matrix suitably.
Proof of Proposition 4.5 From Lemma it follows that there exist weights such that the vectors (hs(xr)) lisxx } are linearly independent. In fact, it is shown that this is true for generic choice of weight matrices A, B for the state space model S, and that the weight matrix A can also be chosen to be nilpotent with A" 0.
Using the linear independence, it now clear that there exists a weight matrix C such that the following holds for each i with S i < K, completing the proof.
Chs(z;) = fs(z;)
A.2 On the context window of state space models and recurrent neural networks.
We first describe the proof of Proposition 4.6 which follows from using the convolutional view of state space models.
19

---

Published in Transactions on Machine Learning Research (02/2025)
Proof of Proposition 4.6 Recall that the sequence of outputs produced is denoted (11.-... yr), and the
sequence of intermediary hidden states as (h,.... .r)) By expanding the hidden state recurrence relation iteratively, we can express h, in terms of previous inputs Iz, 11, 21 as follows:
hT Σ
BoT-K
When & 2 n, note that A = 0; it follows that the value ht depends only on the values (œT, " :..,,,,, , (T-++1) (and not Pr-n or any carlier inputs).
A.3 Proof of the main results for state space models.
Theorem 4.1 follows immediately as a consequence of Theorem 4.3 and Proposition
4.4 We now describe the
proof of Theorem 4.3 and Proposition 4.4
Proof of Proposition 4-4 We choose a SSM whose state transition matrix is nilpotent with order n By the above Lemma, fs( (w)) only depends on the last words of wy we denote the result by Js( (s(F(wx-nt1," wk) To complete the proof, we need to choose the weights of S that the following
holds.
fs(E(uk-nt1: wk-1)) = -
We can do this using the above Lemma on memorizing data with SSMs, by setting x fs( E (w wh w)) and yi = -3 -i
Proof for Theorem 4.3 We show that f3 and fig are cequivalent when restricted to below by unravelling the definitions from Section suffices to prove that for any valid starting sequence w E Lks for some k the following holds.
[lf((w) fng (w)ll <e
This follows from the below chain of equalities.
fš(w) = offtmax(fs(E)))
softmax (8 -
Sw -
fig (w) = fng (wi ((
(1f3(w) - fng (w)) = |3 " Sng(W) nt 1:
A.4 Extensions to RNNs.
First we recall the definition of a recurrent neural network. While the hidden weight matrices are often denoted Wa, WA, W,, in the literature, for consistency we use the notation B, C from Section 3
Definition A.5. A one-lager recurent neural network (RNN) consists of an input layer with ng input neurons, hidden layer with nu hidden neurons, and an output layer with nt, output neurons. The weight matrices denoted by A, B, and C, where A E Mat ny (R) is the hidden-to-hiddae weight matrix, B E Ma (IR) is the input-to-hidden weight matrix, and C' E Mat (RR) is the hidden-to-output weight matrix. Further, by E mm and by E Rnr. are the bias vectors for the hidden and output layers, respectively.
It yields a function fR : (RR")T + (R=2 YT; the RNN processes a sequence of inputs (I1, (2,.22..... ØT), produces a sequence of hidden states (h, hz,.... hrr), and yields outputs (V1, .2..... .r) via the following equations. Here < t t T.
h, aa Aht-1 bn)
yt Ch+ by
20

---

Published in Transactions on Machine Learning Research (02/2025)
Now we show how Proposition 4.6 can be extended to RNNs, by imposing the block nilpotency condition the hidden weight matrix A.
Definition A.6. A matrix M M (iSs..ss E Mat (B) is block nilpotent of order n M = 0 and the induced map ΦM : IR + IR* +I* satisfies the following property: for each i, the space ker(d's) has a basis consisting of standard vectors in IR* (here standard vector is vector in R& with exactly one non-zero coordinate, with value 1).
Lemma A.7. Let R be a recurrent neural network such that the hidden weight matrix A is block nilpotent of order n. Then given an imput sequence x (#1.-.. ,r), the output (x( only depends on the last n input values (#T-a+1>'" ,T)).
Proof of Lemma 7 Denote the sequence of outputs as (y11--.- 3rr), and the sequence of intermediary hidden states as (h1,'** ... .tr) Consider the kernel ker( it is easy to see that bylker(A) depends only on the value of or (and not xr-1 or any carlier inputs). It follows that hrrker(42)) depends only on the value of (IT,2T-1) (and not 22-2 or any earlier inputs). By induction, we see that hrlker((')) depends only on the value of (IT, ITT-411) (and not ZT-4 or any earlier inputs). When = n, note that A = 0 and ker(A") = " follows that the value of hr and yr depend only on the values (FT, (T-n+1)) (and not. Tr-n or any earlier inputs).
We now discuss how Theorem can be extended to RNNs. The proofs in Section A.3 can be applied in this setting without any changes once Proposition 4.6 and Proposition 4.5 have been established. The only remaining ingredient is Proposition 4.5 and we now describe how the proofs in Section A.1 can be exlended
4.5 holds for RNNs, we need to show that the linear independence statement in Lemma A.2 holds in this context; the proof of Proposition 4.5 in Scction A.. can then be applicd without any changes. To show Lemma A.2 note that the bias lerm by and the non-linearity both appear in the hidden state update. However, when the bias term b, has entries which are sufficiently large, the non-linearities are redundant as all the values are positive. The proof of Lemma A.3 can then be used to establish the linear independence, by viewing the coordinates polynomials in the bias by, and the weight matrices A and B.
to RNNs, To show that Proposition
A.5 Bounding the memorization capacity of state space models.
In this section, we revisit Proposition 4.5 and investigate the following question: what is the minimum size of a state space model S with the property that it can memorize an arbitrary set of K input-output pairs? Following the setup in Theorem 3.3 of Yun et al. (2019), for simplicity we assume that the outputs y/ E R (i.e. g = 1). Below we extend the statement. that Theorem for two-layer leedforward neural networks to state space models, which shows that Proposition 1.5 is tight in the setting where =
Proposition A.8. Suppose S is state space model with p input neurons, single hidden layer with k neurons, and one oulput newron. Suppose that l < N. Then there exists a set of input-output pairs (.). with x E IR, y/ E IR with the following property: for any choice of weights matrices A, B, C for S, there erists such that fs( + ye.
Proof. Let u RP be any non-zero vector. Following Yun et al. (019)) we choose 24 iu and yi = (-1)s. We define the piecewise linear function (() = fs((u), and assumne for sake of contradiction that. fs(#;) = y for all 1 i N. It follows that g(i) (-1)* for 1 s i < N, and hence that the function g(t) has at least N pieces.
2019 which states if fi (resp. f2) are piecewise linear functions from R to I with at most P1 (resp. P2) pieces, then fi + fo is a piecewise linear function with at most p1 + P2 - 1 pieces. It follows that by induction that if g1.-.. In are piccewise lincar functions with 2 picces, then 91 + ... + In is a piecewise linear function with at most n pieces.
Recall Lemma D.1 from Yun et al.
Recall that hs(x:) can be expressed as follows.
n-1
hs( = Σ
A Bri
0=0
21

---

Published in Transactions on Machine Learning Research (02/2025)
It follows that the functions g(t) can be expressed sum g(t) 91 (t) +... gh(t), where each ga(t) is composition of the ReLa function and linear function. Since each function gi(l) has 2 pieces, it follows from the above that the piecewise linear function g(t) has at most k + 1 pieces. This contradicts the prior deduction that g(t has at least N pieces, since l < N.
B Appendix B
In this section, we elaborate on the example from Section 3.2 which is also used in the experiments Section [
We define language L using the two sentence templates below. Sentences in L are generated by substituting specific words into the placcholders (A, B. J) specified below.
A woke, at, B, o'clock, and, was, too, C, to, go, back, to, D]
He, E, up, and, F, on his, G, because he, didn't, want, to, H, into the, I, in, his, J]
A: [Hrrr,, Ron, Sirius, Hagrid, Fred, George, Neville, Draco, Albus, Snape]
B: [one, two, three, four, five, six, seven, eight, nine, ten, eleven, twelve]
C: [excited, angry, elated, agitated, nervous; troubled, upset]
D: (sleep, bed]
E: [got, stood, rose]
F: [put, pulled]
G: Ijeans, trousers, pants, denims, slacks]
H: [walk, run, jog, stroll, sprint, saunter, amble]
I: [station, platform, house, building, apartment]
J: [robes, pyjamas, shorts, underpants, boxers, trunks]
Below we list sample sentences that can be generated from L.
Harry woke at one o'clock and was too excited to go back sleep.
2. Ron woke at seven o o'clock and was too troubled to go back to bed.
3. He got up and put on his jeans because he didn't want jog into the station in his robes.
He stood up and pulled on his trousers because he didn't want to sprint into the platform in his pyjamas.
The language L can be generated from an n-gram language model fäg; with n 4. The table below has a complete list of all n-grams in P: note that [P] 145. See Section for an example showing how the function Sng P + A is defined.
22

---

Published in Transactions on Machine Learning Research (02/2025)
Table 1: All valid trigrams in from the language L.
Harry woke
Ron woke Neville woke at woke at two woke at seven woke at twelve at five o'clock al len o'clock three o'clock and cight o' o'clock and o clock and was was too agitated too angry to too upset to nervous to go back sleep got up and and put on his trousers his trousers because trousers because he he didn't want want stroll to into to amble into sprint into the into the house the house in house in his his shorts
Sirius woke Draco woke at
Hagrid woke at Albus woke at woke at four woke at nine at two o'clock at seven o'clock twelve o'clock five o' o'clock and ten o o'clock and was too excited was too troubled too agitated to angry to go upset to go He got up rose up and put his on his denims his denims because denims because he want to walk want to saunter
Fred woke Snape woke at woke at five woke at ten at three o'clock at cight o'clock one o o'clock and six o'clock and cleven o o'clock and was too angry was too upset too nervous to elated to go to go back He stood up up and put pulled on his on his slacks his slacks because slacks because he want to run want to amble to sprint into jog into the into the station the station in station in his in his robes in his trunks
George woke at woke at one woke at six woke at cleven at four o'clock at nine o'clock o'clock and seven o'clock and twelve o c'clock and was too elated too excited to too troubled to agitated to go go back to He rose up up and pulled on his jeans his jeans because jeans because he because he didn't want jog to walk into to saunter into stroll into the into the platform the platform in platform in his
woke at three woke at eight at one o'clock at six o' o'clock al eleven o' o'clock four o'clock and nine o clock and and was too too nervous too elated to excited to go troubled go back to bed stood up and and pulled on on his pants his pants because pants because he didn't want to want sprint to jog into walk into the saunter into the into the building the building in building in his in his underpants
stroll into run into the amble into the into the apartment the apartment in apartment in in his boxers
in his pyjamas
23


Reviews:
[{"review_id": "ThNZiHhaIc", "reviewer": "Reviewer_4zFd", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths:**\n1. The paper provides explicit mathematical constructions for what it aims to achieve.\n2. The paper cites the literature well.\n\n**Weaknesses:**\n1. The definitions in the paper are very questionable if not incorrect. For instance, the definition of an n-gram language is just as a list of valid strings. Why is there no underlying probabilistic model (as would be the case for the usual definition for an n-gram model in which p(w | context) is given)?\n\n    Likewise, the definition of risk is quite odd. It is equal to the average probability mass that f assigns to strings not in the n-gram language (thus zero risk means placing mass only on the supported n-gram strings of the n-gram language). But this is rather odd because it does not take into account what the distribution of f looks like on the n-gram language. What one really wants is an a priori distribution on strings and for the risk to measure the deviation from that distribution (e.g. a cross entropy type loss). \n\n     As a degenerate example, consider a language that consists of all possible strings. Then any state space model, however initialized, has zero risk. \n\n    Perhaps the authors had more in mind something like the recognition problem for languages (following [3]), but are restricting to n-gram languages. Then the terminology should reflect that (n-gram models are probabilistic and so are not the appropriate terminology). However I would argue this is not the correct perspective to take, since people are interested in n-gram models not n-gram languages.\n\n2. Assuming #1 is fixed, the title appears quite suboptimal. A more accurate and direct title would be \"State-space models can express\nn-gram models\". It's unclear why there is a question in title when the paper is putting forth a theorem: the title should just reflect that. \n\n3. There is no discussion about the sharpness of the results. Given that, morally speaking, a sufficiently large neural network can model any finite dataset, it's not surprising that some configuration of the state-space model can represent an n-gram model. Moreover, this is more obvious for a state-space model given its linear recurrence, than for a transformer: an n-gram model can be thought of as an (n-1)-Markov model and the latter can be encoded in a sufficiently large matrix; a state-space model which \"implements\" this Markov transition matrix would recover the n-gram language. However, such a transition matrix would in general have as its dimensions $V^{n-1} \\times V$ (where $V$ is the vocab size), i.e. the space of all possible transitions from (n-1)-grams to 1-grams. While it is interesting that the author's main theorem allows of any embedding dimension greater than or equal to 1 (this seems surprising, I have not looked carefully as to why embedding dimension = 1 works, the author should emphasis this wide range), the author's main theorem also emphasizes that their hidden state has size the number of n-grams but not the number of sentences in the language. This however isn't surprising given that the corresponding Markov matrix also has size the number of n-grams. \n\n    If the authors could provide some analysis on (sub)optimality, the paper would be more complete in its analysis. For instance, in the cited reference [1], there are sharpness results showing that a certain capacity is needed for their representation results. It would be meaningful if the authors could do a similar analysis. \n\n4. Summarizing the above, the authors would benefit from referring to [2], which although was cited in their work, did not seem to inform the authors as much as it should have. This is because [2] treats n-grams probabilistically. While [2] also needs hidden dimension that grows as |V|^n (again to encode the Markov matrix associated to an n-gram model), their construction at least discusses this dependence (and how it might be mitigated) and provides a variety of different constructions (varying number of heads/layers). The current paper seems to give a single construction without much context as to how (un)canonical or (not) close to optimal it may be. \n\n\n[1] Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small relu networks are powerful memorizers: a tight analysis\nof memorization capacity.\n[2] Anej Svete Ryan Cotterell. Transformers Can Represent n-gram Language Models\n[3] Perez et al. Attention is Turing complete.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors show that state space models can model an n-gram language in a certain (unconventional) sense.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "**Technical revisions:**\nPlease address or revise the technical concerns mentioned above. \n\nI marked the Audience criteria below as a \"No\" based on my concern #1 above. The notion of risk is very unusual and there is no attempt to model n-gram models as probabilistic models. In my view, this needs to be addressed in order to revise the answer to \"Yes\".\n\n**Notational comments:**\nMany quantities are overlined and underlined and it is unclear (and confusing) what such modifications indicate (if anything). Please fix. *Examples:*\n\nA language is denoted by an $\\bar{L}$, but there is no $L$. Likewise with $E$ for embedding.\n\nWhat is the difference between $f$ and $\\bar{f}$?\n\nA nonlinearity is denoted by $\\sigma$ but a overline $\\sigma$ is softmax.\n\nIn Definition 3.3, the notation $f(w | \\underline{w})$ is used; $w$ is meant to be a dummy variable while $\\underline{w}$ is fixed, yet this conflicts with $\\underline{w}$ being used to denote that $w$ is a vector elsewhere.\n\n**Expository comments:**\n\nThe example after Definition 3.1 could just be explained more straightforwardly and less confusingly. Using the words \"intuitively\", showing an incomplete table of n-grams, and using a partial example of substitutions like \"Ron\" and then finally just showing the full table in the Appendix, is wasting the reader's time when all that's needed is just to say refer to full table in the Appendix for the full list of n-gram generators.  In fact it's unclear why these were the chosen valid n-gram generators (is it just all valid n-grams from two sentences of Harry Potter as well as some formulaic word substitutions?). If so, that can just be explained without writing a very long table.", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "l7CzZIrdOt", "reviewer": "Reviewer_bNnk", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I like very much the topic of the paper and I think it could be of great interest of the community.\n\nStrengths:\n- Topic is of great interest for the community: understanding the power of simple linear recurrences is an unexplored and important topic.\n- The paper is relatively easy to read, proof are quite simple\n- The authors provide a nice introduction to the topic and give all necessary background for people with little background to understand \n\nWeaknesses:\n- Proof is a bit fragmented, and some key information is hidden in short appendix comments. See requested changes section.\n- The main content of the paper is not much: the proof can be condensed in a few key ideas and experiments are quite simple. \n- Notation is a bit heavy at times.\n- Role of some key quantities, such as the epsilon in the main theorem, is not clear and not commented.\n- The paper refers to Jelassi et al. 2024 as a similar contribution: while I agree, what is missing here is a detailed comparison with transformer architectures.\n- Notation can change in the paper: recurrent matrix is sometimes called A, sometimes W_hh, sometimes W_h.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors study expressivity of simple linear RNNs (perhaps trained like SSMs) followed by a nonlinearity. In particular, the authors are concerned with showing SSMs, of increasing size, can represent all rules in an n-gram language. The authors start with a literature review, comprising both an overview of why SSMs are of interest and a summary of related works on n-grams representations. The authors proceed with formal definition of embeddings, n-grams (with examples), and linear recurrences. Section 4 is dedicated to the proof while Section 5 complements with empirical findings. \nThe heart of the paper is section 4, outlining the main result and few key propositions proved in the appendix. Essentially, the authors prove that 1 layer model SSM (without MLP but with final nonlinearity) can represent a number of n-gram rules equal to its hidden neurons.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "**Changes to the proof workflow:**\n\nI suggest making the inner workings of the proof a bit more evident in the main paper.\n- Thm 4.1: it is not clear what role epsilon has. It seems it has very little effect, but then one might ask why its there! You should comment on this and provide an example. This is also related to Definition 4.2, that  is indeed very complicated and perhaps an example can help the reader understand how exactly you feed in input-output sequences.\n- Thm 4.3 and Prop 4.4: here notation is very hard to parse. please put in examples.\n- Relation between Prop 4.5 and prop 4.6: I feel like this is the heart of the proof but this is indeed hidden in the appendix: Prop 4.5 says that there EXISTS (A,B,C) such that... but in prop 4.6 you need A (called W_hh, please adapt) to be nihlpotent. It is unclear to me if you can use 4.5 while setting A to be nihl potent.. digging in the appendix, I found (proof of 4.5) \"this is true in fact for any A,B\". What is it then? does there exist a pair $(A,B)$ or can you work with nihlpotent?\n- \"pairwise distinct\": please move definition from appendix to main paper. Why is this assumption not restrictive? does this impliy something on the possible n-gram rules you can model?\n- Lemma A2 in app: $N$ neurons... do you mean $K$ neurons? same holds for $\\{h_S(x_i)\\}_{i=1}^N$, which I think should be $K$.\n\n**Changes to discussion and experiments**\n\nPlease provide a comparison (empirical and theoretical) with attention-based models.\n\n**Typos**\nPlease update the way to build the main paper: compile all and cut pages, so you avoid the \"??\" for appendix references.\n\n--------------------------\n\nComments mostly solved by the authors", "broader_impact_concerns": "No concern.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vcDFh333XA", "reviewer": "Reviewer_FqGM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Good: The main results are clear and bring more intuition on what SSM can do.\n\nBad: I do not like the paper writing.\n\nThere is so much state wasted on filler, e.g. whole example of n-gram model could be shorter, not use full words from Harry Potter, but just symbols A B C D; Section 4 is full of text, but there is no proof, all is deferred to the Appendix (more on this later).\n\nThere is no simple figure explaining the proof idea. I am not 100% sure about SSM, but I can explain a similar result for RNN with ReLU nonlinearity and one-hot encoded inputs with a simple picture (the hidden state will one-hot encode which n-gram we see, ...). Authors should explain the main proof idea in the main body of the paper (including some pictures!) and refer to technicalities in the appendix. \n\nThe experiment section writing is also not great. It refers to nonexistent Lemma 5.6. Also, it never states the size of the n-gram language (the number of valid n-grams), so it is hard to compare the number of hidden neurons (64) to the size of the language.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper shows that SSMs have more representation power than n-gram models.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Rewrite the paper.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper provides a proof by construction on how an SSM can encode a n-th order Markov chain. The construction seems valid, and the result adds to our theoretical understanding of what these architectures can do. The work could have been improved with strong empirical exploration, but I think that is not for acceptance.
Decision: Accept (poster)