RNNs are not Transformers (Yet):  The Key Bottleneck on In-Context Retrieval
Authors: Kaiyue Wen
Xingyu Dang
Kaifeng Lyu
Venue: ICLR 2025 Poster (2025)
Keywords: rnn
cot
representation theory
Primary area: learning theory
TL;DR: We identify a representation gap between Transformers and RNNs (w/o CoT) due to the incapability to retrieve in context and show by adopting techniques to enhance this capability can bridge the gap.

Abstract: This paper investigates the gap in representation powers of Transformers and Recurrent Neural Networks (RNNs), which are more memory efficient than Transformers. We aim to understand whether RNNs can match the performance of Transformers, particularly when enhanced with Chain-of-Thought (CoT) prompting. Our theoretical analysis reveals that CoT improves RNNs but is insufficient to close the gap with Transformers. A key bottleneck lies in the inability of RNNs to perfectly retrieve information from the context, even with CoT: 
for several tasks that explicitly or implicitly require this capability, such as associative recall and determining if a graph is a tree, we prove that RNNs are not expressive enough to solve the tasks while Transformers can solve them with ease.
Conversely, we prove that adopting techniques to enhance the in-context retrieval capability of RNNs, including Retrieval-Augmented Generation (RAG) and adding a single Transformer layer, can elevate RNNs to be capable of solving all polynomial-time solvable problems with CoT, hence closing the representation gap with Transformers. We validate our theory on synthetic and natural language experiments.

Published as a conference paper at ICLR 2025
RNNS ARE NOT TRANSFORMERS (YET): THE KEY BOTTLENECK ON IN-CONTEXT RETRIEVAL
Kaiyue Wen'* Xingyu Dang" Kaifeng Lyu
1 Stanford University Tsinghu University
3 University of California, Berkeley
ka i yuwwestan edu tanggaa edu.cm
ka i fenglyuubereeeeg edu
ABSTRACT
This paper investigates the gap in representation powers of Transformers and Recur- rent Neural Networks (RNNs), which are more memory efficient than Transformers. We aim to understand whether RNNs can match the performance of Transform- ers, particularly when enhanced with Chain-of-Thought (CoT) prompting. Our theoretical analysis reveals that CoT improves RNNs but is insufficient to close the gap with Transformers. A key bottleneck lies in the inability of RNNs perfectly retrieve information from the context, even with CoT: for several tasks that explicitly implicitly require this capability, such associative recall and determining if a graph is a tree, we prove that RNNs are not expressive enough to solve the tasks while Transformers can solve them with ease. Conversely, we prove that adopting techniques to enhance the in-context retrieval capability of RNNs, including Retrieval-Augee Augmented Generation (RAG) and adding single Transformer layer, can elevate RNNs to be capable of solving all polynomial-ting solvable problems with CoT, hence closing the representation gap with Transformers. We validate our theory synthetic and natural language experiments.
INTRODUCTION
Transformers Vaswann et al. 2017) have become the dominant choice of the backbone for Large Language Models (LLMs). The core component of Transformers is self-attention modules, which allow the model to route information densely across the entire sequence. However, this design leads to high inference costs for long sequences, including memory cost linear in the sequence length to maintain intermediate attention keys and values for each token, and a time cost quadratic in the sequence length to compute the attention score for each pair of tokens.
Recently, Recurrent Neural Networks (RNNs) have become an increasingly popular choice in sequence modeling tasks due to their ability to maintain a memory size constant in sequence length during inference, thus being more memory efficient than Transformers. Katharopoulos et al. (2020) showed that Linear Transformers (Transformers with linear attention) be expressed RNNs.
Gu et al. 2022) took a different path to design RNNs by structuring latent states as State Space Models (SSMs) from control theory. These ideas have led to series of development of modern
RNNs, including RWKV (Peng et al. (023), RetNet (Sun et al.] (023), and Mamba (Gu & Dao 2023). Most notably, Mamba can achieve competitive performance with Transformers on several sequence modeling tasks with linear time and constant memory in sequence length.
Can RNNs replace Transformers yel? The rise of these modern RNNs has led an interest in understanding their limitations. A recent work by Arora et al. (2023) showed that a broad family of RNNs, input-independent gating SSMs, are empirically inferior Transformers in task that has a long history in artificial intelligence, Associative Recall (AR) (Willshaw et al., 1969 Hopfield 1982, Hinton& & Anderson 2014): Given a series of key-value pairs as a string, the model is required to recall the value given a key. On the theory side, [Sanford et al. (2023) and JJlassi et al. (2024) demonstrated that constant
RNNs do not have sufficient representation power to solve
"qqual contribution
1 soonespn author
The work was done while Kaiyue was Tsinghua University and Kaifeng was at Princeton University.

---

Published as a conference paper at ICLR 2025
Transformer + CoT C P
RNN + In-Context RAG 3 P
RNN + CoT D
Vanilla Transformer
Vanilla RNN
Transformer
RNN (This work)
Figure Hierarchy of Representation Power. While RNN with chain-of-thought (CoT) with Oloo n) bit memory provably has strictly stronger representation power than RNN without CoT under mild complexity assumptions (Theorem 4.1). I, it is still exponentially weaker than Transformer with CoT in representing solutions to algorithmic problems (Theorem 4.7 We proceed to show that the incapability of RNNs in In-Context Retrieval is the root cause of the gap and propose two forms of In-Context Retrieval Augmented Generation (In-Context RAG) to close the gap by illustrating their power to simulate any polynomial-time Turing machines (Theorems 5.2 and 5.4).
the tasks of averaging given subset of input vectors (q-sparse averaging) and repeating the input sequence (copying), respectively, while there exist shallow Transformers that can solve these tasks.
However, the above results do not exclude the possibility that enhancing RNNs with additional prompting techniques or minor architectural changes could close the gap with Transformers. In fact, Transformers themselves are not perfect either and sometimes need additional techniques at inference time to perform well. For instance, Transformers may struggle with mathematical and algorithmic reasoning problems if they are forced to produce the correct answer immediately after processing the input sequence. But with Chain-of- of-Thhught (CoT) prompting applicd, a prompting technique that guides the model generate series of intermediate tokens before arriving at the final answer, their performance can be significantly improved. Feng et al. (2023) y Li et al.) (2024) explained the success of CoT from the perspective of representation power: Transformers alone do not have sufficient representation power to solve problems beyond certain circuit complexity class (TC"), but with CoT, they can even simulate any polynomial-tinn -ime Turing machine.
The effectiveness of CoT on Transformers naturally leads the following question:
Can similar enhancements, such as adopting CoT, improve RNNs be on par with Transformers?
Our Contributions. This paper examines potential ways to close the gap in the representation powers of RNNs and Transformers (with softmax attention) on algorithmic problems. Through a series of lower and upper bound results, we show that CoT improves the representation power of RNNs, but to close the gap with Transformers, CoT alone is not enough overcome key bottleneck of RNNs: their inability to retrieve information from the context, which we call in-context retrieval for short.
We further illustrate that addressing this in-context retrieval bottleneck is sufficient close this gap: RNNs can solve all polynomial-ting solvable problems if adopting techniques to enhance the in-context retrieval capability, including involving Petrieval-Augment Generation (RAG) and using hybrid architectures, such as appending a single Transformer layer.
Our main contributions are listed as follows:
1. CoT improves RNNs but cannot close the representation gap with Transformers, (Section
On the positive side, we prove that CoT makes RNNs strictly more expressive under mild assumptions from circuit complexity (PSPACE of P/poly).
On the negative side, we show that adopting CoT is not enough close the representation gap between RNNs and Transformers: the memory efficiency of RNNs fundamentally limits their ability to perform in-context retrieval, even with COT. This point is made concrete by proving that RNNs with CoT cannot solve set of fundamental algorithmic problems that directly ask for in in-context retrieval, including associative recall.
We further exemplify that in co-conext retrieval can be implicitly required in tasks that appear unrelated, by proving the inability of RNNs to solve the classic problem of determining whether graph is tree (Is (IsTree).
2

---

Published as conference paper at ICLR 2025
0.5M TRANSFORMER
1M TRANSFORMER
2M TRANSFORMER
Val Accuracy
1
0.8
0.4
16
32
64
128
16
32
64
128
16
32
64
128
0.5M RNN
1M RNN
2M RNN
Val Accuracy
1-
0.8
0.4
16
32
64
128
16
32
64
128
16
32
64
128
Number of Nodes
Number Nodes
Number of Nodes
Venilla
Herrival
Figure 2: We train RNNs (Mamba) and Transformers (LLaMA Touvron et al. 2023)) with a frozen word embedding and decoding head of three different model sizes (0.5M, 1M, 2M) on IsTree with three different sizes of graph (16, 32, 64) under three different setups. Vanilla means the model directly predicts the label. COT means the model will generate chain-of-thought process based on DFS (see Algorithm|| [ before prediction. Retrieval means the model will generate the chain of search queries and reasoning before prediction (see Algorithm We observe that (1) Both Transformer and RNNs can't tan't solve the Is'Tree question without chain of thought; (2) RNNs' performance with tothohh decays quickly when the number of nodes increase, which consistent with our theory; (3) All models reach almost perfect accuracy when enhanced with retrieval.
Our negative results hold for wide range RNN architectures, including the aforementioned Mamba, RWKV, and even Linear Transformers. Technically, this is because RNNs memory efficient that they can trigger streaming lower bounds (Sanford et al. 2023 problems that require in-context retrieval.
especially for
Enhancing the in-context retrieval capability of RNNS close the gap. (Section
We prove that allowing RNNs invoke function calls to perform primitive of in-context retrieval based on regular expression is sufficient to boost their representation power to solve all polynomial-ting time solvable problems with CoT, hence closing the representation gap.
As one layer of the Transformer is sufficient to perform many in-context retrieval operations, mixing some Transformer layers in RNNs should also narrow the representation gap. We prove that minimal possible change in the RNN architecture can just work: adding one Transformer layer the end of the RNN architecture is sufficient close the representation gap.
Our positive results showing that enhancing in co-conextt retrieval can improve RNNs' representation power hold for vanilla Linear RNNs, and can be easily extended to more complex architectures. The intuition behind these results is that RNN can focus on the local reasoning steps and the in-context retrieval module to adaptively fetch the relevant information from the context.
We validate our theoretical findings through synthetic and natural language experiments on IsTree and HotPor-QA, confirming that while CoT alone cannot close the performance gap between RNNs and Transformers, the proposed two solutions effectively narrow this gap. (Section
Implications. We believe these results could provide valuable insights architecture designs of LLMs: RNNs alone can suffer from many fundamental limitations in representation power, even with CoT; on the other hand, it is promising to explore strategies enhance the in-context retrieval capability of RNNs with little overhead, such using hybrid architecture that mixes in one or more Transformer layers (Ren et al. 2024. Waleffe et al. 2024
Lieber et al. 2024b).
RELATED WORKS
State Space Machines and Linear Transformers. There has been a recent surge of interest in state space machines and (kernalized) linear transformers (Gu et al.
2022:
Katharopoulos al.
2020 Peng et al. 2023: Sun et al., 2023; Gu & Dao 2023) Fu et al. 2023 Poli et al. (023) Luo al. 202 Peng et al. 2021 Wang al. 2020), which are class of models that combine the paralle ciizaallizaaility the Transformer with the memory efficiency of the RNN. These models can process both sequential and recurrent form, and the former for fast parallelizable training

---

Published as a conference paper at ICLR 2025
and the latter for memory -effccienef inference. However, these models are still empirically inferior to the Transformer in terms of performance. Our work investigates the reasons behind this gap and proposes to close it by enbancing the in-context retrieval capability.
Chain of Thought (CoT), Chain of thought (Wei et al. 2023; Nye et al. 2021 Kojima et al. 2023
Wang & Zhou 2024) is augmentation to the Transformer, that allows it to solve more complex reasoning tasks by generating a reasoning process before outputting the answer. It has been shown that Transformers with CoT provably have more expressive power than the original Transformer without CoT (Feng et al. 2023 Li et al. 2024). However, the expressive power of RNNs with CoT has not yet been systematically studied. Theorem F.1 Feng et al. (2023) shows that RNN cannot output a particular format of CoT for evaluating arithmetic expressions and solving linear equations while Transformers with the same amount of parameters can. Concurrent work (Yang et al.) 2024) discovers that linear Transformers, special class of RNNs, are not able to solve some dynamic programming problems with CoT, unless the number of parameters grows with the length of the input. One high-level message our work conveys is similar to theirs: RNNs have limited representation power to perform reasoning with COT. However, we show that such limitation not specific to the output format or architecture and apply tools from streaming complexity to prove lower bounds broader range of tasks and memory memory-efficient architectures.
Streaming Algorithms. Our lower bound leverages the technique in streaming algorithms. Streaming algorithms are algorithms that take constant (typically just 1) pass over the input and use sublinear space, hence including RNNs with fixed state space as a special case. Works in streaming algorithms date back to the 1980s (Munro& & Paterson 1980) and have been formalized and popularized in the 1990s (Alon et al. 1996) due the need to process large data streams. The study of streaming algorithm is closely connected with the concept of communication complexity. The communication complexity of an algorithm is defined by the amount of communication cost required when the algorithm is distributed, and all streaming algorithms can be viewed as distributed algorithms with sublinear communication complexity, whose communication content is the internal state of the algorithm. The lower bound in our work is direct application of this observation to the study of RNNs and we mainly consider the lower bounds for (1) indexing the input (Munro & Paterson 1980) and (2) determining whether the input is tree (Henzinger et al. 1998).
We will defer other related works to Appendix A
PRELIMINARIES
We briefly introduce the definitions that are necessary for understanding our results and defer detailed definitions to Appendix B
Vocabulary and Embeddings. Let n be a parameter so that the input length can be bounded by Onz). We assume finite vocabulary V that includes natural numbers from to and constant number of special tokens. For our theoretical analysis, we fix word and position embeddings and only consider the representation power of the model when the non-embedding parameters can change. See Appendix [B.3]o for the formal definitions. Note that in the common practice, the vocabulary size does not increase with and numbers may be tokenized into a few tokens according to their decimal representations. Nevertheless, our lower bounds can be easily extended to this more practical case, since these results do not rely on the specific form of the vocabulary and embeddings. For upper bounds, our embedding scheme should be easily implemented from other embedding schemes with a few RNN or Transformer layers.
Language Models. We denote the set of all finite sequences and all non-empty finite sequences of tokens in V as V and V+, A language model (LM) function M V - Py that maps a non-empty sequence to probability distribution over the next token, where Pyv is the probability simplex over V. We focus on LMS realized by deep neural networks: mapping an input sequence S to sequence of embeddings Emb(S), then applying a neural network (e.g.. Transformer or RNN) to process the embeddings and output the probability distribution.
Numerical Precision. We consider models with fixed numerical precision. We use p to denote the precision of the number of bits to represent real numbers and use R, to denote the set of all real numbers that can be represented by p-bit floating point numbers. We defer the details to Appendix We assume p = Olog n), which is a common assumption when studying the finite precision neural networks (Feng al. 2023
Merrill & Sabharwal 2023).
4

---

Published as a conference paper at ICLR 2025
Transformers. We study Transformers with standard decoder architecture containing softmax attention and gated feedforward layers. See Appendix [B3 Transformers, where the depth, width and number of parameters are fixed and do not depend on n.
for for definitions. We focus on constant-size
RNNs. Recently there has been a lot of interest in the linear-time Transformer, which replaces the full-attention calculation with linear-time alternatives. These variants mostly special forms of recurrent neural networks (RNNs) that are parallelizable. Here we define a general form of RNNs containing all the common variants the best of our knowledge, including Mamba (Gu & Dao
2023), RWKV (Peng et al. 2023), RetNet (Sun et al. 2023), StreamingLLMM (Xiao et al.) 2023). RMT Bulatov et al., 2022). p. TOVA (Oren al. (024), xLSTM (Beck et al. 2024) etc. An RNN maintains a state St E R. storing A p-bit floating point numbers. At each step, the RNN generates the next state 8t+11 based on the current state 8; and token S and then it makes a prediction based on the updated state $1+1- We characterize the complexity of an RNN with the following three measures: (1) the parameter size P in bits, which is the total number of bits in parameters; (2) the memory size M = Ap in bits, which is the total number of bits in each state; (3) the circuit size C, which is the number of bit-wise operations needed to compute one step of the RNN. A particularly interesting class of RNNS is constant-sizz size RNNs, where the dimension of state and number of parameters are fixed and do not depend on n, i.e., A = O11, p = (log n). P = (log n), M = Olog n), C. = (log n). We do not assume specific structure of the RNN when we want to prove impossibility results for RNNs, i.e., what RNNs cannot represent. But when we want to showcase what RNNs can represent, we focus on Linear RNNs, one of the simplest form of RNNs, so that our results are more likely be generalizable to more complex RNNs. See Appendix B.3 for the definitions of all the above models.
Algorithmic Problems. An algorithmic problem is a problem that may be solved by an algorithm. In this paper, we focus on algorithmic problems f: V - VA that asks for computing (( given a sequence of tokens Snn as the input, where VA the set of possible answers. We say that an LM M can (directly) solve an algorithmic task f if, given the sequence Sin. the probability distribution (S for the next token is peaked at the correct output token /' (S,,, i.e., arg maxjev M (S.)) = ((S)).
Mfain-of-Though Reasoning, Chain-ofTThought (CoTT) reasoning allows the LM to produce intermediate steps before the final output. Following Feng et al. (2023); Li et al. (2024), our paper studies the effectiveness of CoT reasoning in improving the expressiveness of LMs, and we allow the intermediate steps to be an arbitrary sequence of tokens. We say that an LM M can solve an algorithmic problem with CoT if the following process terminates with sequence ended with ((..). First, let S. = Sin- For all i 2 0, decode the next token spext = arg max,ev M (S.(ii) from M, and append it the current sequence S S; @ 8 spext If s next e VA, then the process terminates with S with i steps of CoT,; otherwise the process continues. It is evident that if an LM can solve algorithm problem with 0 steps of COT,, then LM M can (directly) solve the problem. In this case, we also say that the LM can solve the problem without COT.
4 HOW WELL DOES COT IMPROVE THE REPRESENTATION POWER OF RNNS?
In this section, we aim to understand how well CoT improves the representation power of RNNs and how it compares with Transformers.
4.1 COT STRICTLY IMPROVES RNNS
First, we show a positive result that RNNs with CoT can solve tasks that RNNs without CoT cannot solve. This relies on mild complexity assumption PSPACE & P/poly, i.c., not all polynomial space complexity problems can be solved by polynomial-size circuit family. This has been widely believed to be true since PSPACE P P/poly would imply severe complexity consequences (Karp &
Lipton| 1980; Babai et al. 1991).
Theorem 4.1. Assuming PSPACE & P/poly, there exists an algorithmic problem such that (1) there exist constant-size Linear RNNs that can solve the problem with polynomial length CoT; and (2) any constant-size regular RNNs cannot solve the problem without COT.
See Appendix E.S for the proof. The key insight is that the representation power of RNNs without COT is limited to shallow circuits of size poly(log n). but RNNs with CoT can simulate O(log n)- space Turing machines perfectly with (oly(n) steps. This result is consistent with previous works on the benefit of CoT for Transformers (Feng et al. 2023 Li et al. 2024), which also prove based on mild complexity assumptions that Transformers with CoT have the representation power simulate
5

---

Published as a conference paper at ICLR 2025
polynomial-size circuits to solve all the problems in P but Transformers without CoT cannot. Here we rigorously prove that a similar benefit of CoT also applies to RNNs, but a key difference is that CoT cannot boost the representation power of RNNs to simulate every polynomial-siae circuit family.
4.2 COT CANNOT CLOSE THE REPRESENTATION GAP WITH TRANSFORMERS
Now we proceed to understand whether CoT can make RNNs as expressive as Transformers. The answer turns out to be negative: RNNs, even with CoT, struggle to solve very simple algorithmic problems that require the capability of retrieving information from the current context, which we call In-Context Retrieval for short. This limitation is directly related to the memory efficiency of RNNs. For a model with at most ((n) bits in memory. we can involve techniques from streaming complexity to prove impossibility results for problems requiring In-Context Retrieval.
4.2.1 SIMPLE PROBLEMS ON IN-CONTEXT RETRIEVAL
Here we list several simple algorithmic problems that directly test the in-context retrieval capability of the model, which turn out to be a good test-bed for understanding the limitations of RNNs compared to Transformers. We defer discussion on the history of these problems to Appendix [
Definition 4.2 (Index). Index is problem that given sequence of tokens with length n and a query token i [n). requires the model to output the type of the i-th token.
Definition 4.3 (Associative Recall). Associative Recall (AR) is problem that given sequence of tokens with length n consisting of tokens in [m] and a query token E [n]. requires the model to output the next token of q in the sequence.
Definition 4.4 (c-gram Retrieval). An c-gram is contiguous subsequence of e tokens in a sequence of tokens. c-gram retrieval is a problem that given sequence of tokens with length n and query (c - 1)-gram that is the prefix of a c-gram in the sequence, requires the model to output the last token of that c-gram.
Definition 4.5 (Counting). Counting is problem that given sequence of tokens with length n, a query token E [n). and a query number i E N. requires the model to output 0 or 1 to indicate whether the number of occurrences of q in the sequence is greater than f.
The following theorems show that constant-size RNNs cannot solve any of the four tasks when the context length long, while constant-size Transformers can solve them perfectly.
Theorem 4.6. For task T {Index, AR, c-gram retrieval, Counting), ). there exist constant-size Trans- formers that can solve T. On the other hand, any RNN with onn)-bit memory cannot solve T size with any length of CoT for large enough n.
See Appendix E.4 for the proof. The key idea of the lower bound is to view RNNs in the framework of communication complexity and use information- -ttooortccc arguments to prove a lower bound. The computation of RNNs can be distributed to two parties: party A simulates the RNN on the first part of the input and sends the state to party B, then party receives the state and simulates the RNN on the second part of the input (the query string) until it finishes producing the output. Therefore, the RNN can solve the problem with o(n)-bit memory, then the input can be compressed to (n) bits while maintaining all the necessary information to answer all possible queries, which contradicts certain informaaiiontheoreete lower bounds.
We note that Theorem 4.6 does not imply that RNNs are incapable of in-context retrieval at all. Instead, it states that the maximal context length that RNNs can effectively retrieve from is linear in its state size. Although for short context window, RNNs can be trained perform in-context retrieval (see e.g. |Arora et al. (2023); Gu & Dao (2023), this limitation in retrieval capabilities has been empirically observed: for example in Waleffe et al. (2024). both pretrained Mamba and Mamba-2 7B models are shown to have significantly worse Phonebook-retrieval capabilities on IK context length than Transformers with the same size and trained on the same data.
4.2.2 LIMITATIONS OF RNNS BEYOND SIMPLE IN-CONTEXT RETRIEVAL PROBLEMS
A natural question would be if an algorithmic problem does not directly test the in-context retrieval capability, can we hope that RNNs would have the representation power to solve it? Do RNNs and Transformers have the same representation power in this case? Still, the answer is negative. We show

---

Published as a conference paper at ICLR 2025
that the limited memory size in RNNs can still be bottleneck in solving algorithmic problems. Even the retrieval capability is not explicitly tested in an algorithmic problem, it may still be required implicitly for reasoning about the answer.
We demonstrate this gap on a simple algorithmic problem that is seemingly irrelevant to retrieval: the IsTree problem. Given an undirected graph G of nodes, determine whether G is tree, i.e., whether every pair of nodes is connected by exactly one simple path. We convert the graph G into a token sequence, where we input the graph edge by edge. Figure 6 shows example of the IsTree problem. A classical solution to IsTree is running Depth First Search (DFS), which takes O(n) time. See Appendix B for more details.
Our result below shows that constant-size RNNs cannot solve IsTree, even with CoT. But constant- size Transformers can perform CoT reasoning by simulating DFS and perfectly solve the problem.
Theorem 4.7. There exist constant-size Transformers that can solve IsTree with CoT of length O(n)- On the other hand, any RNN with -bii memory cannot solve IsTree with any length of CoT.
See Appendix E.6 for the proof. The key idea of the lower bound is to again utilize information- theoretic lower bounds. This idea lies in the core of streaming complexity literature and investigation on the IsTree problem, dating back to Henzinger et al. (1998). We hereby restate the proof for completeness.
Given any binary sequence x of length n 2 and an index [n 3]. we construct graph as follows: the graph has n nodes, and vertex a is connected to vertex Ha + n for any a E [n 2]. Additionally, vertex is connected to vertex k + 1. The graph is a tree if and only if 24 * Ik+1- Assuming there is an RNN with o(n)-bit memory that can solve IsTree, consider two parties, A and B, each holding the sequence x and the index k. They can construct two parts of the graph using their information: A simulates the RNN on the first part and sends the state to B, who then simulates the RNN (potentially with CoT) on the second part to determine the output of the IsTree problem. equivalent to checking whether Tık + Zk+1. However, since k is never sent to A, B can determine whether F& F Ik+1 for any h E [n- - 3], contradicting certain information -thooretiiit lower bounds.
4.2.3 TRANSFORMERS (WITH CоT) ARE STRICTLY MORE EXPRESSIVE THAN RNNS
The above theorems show that RNNs have significant representation gap with Transformers, even with COT. Finally, we show that this gap is one-sided: there only exist tasks where Transformers can solve them and RNNs cannot, but not the other way around. The proof (Appendix E.7 is inspired by a recent work on CoT for Transformer (Li et al. 2024).
Theorem 4.8. Given input length n, let an RNN with word embedding W(E) E R xd where p = (log n) is the precision, the constant ng is the number of special symbols in the vocabulary, the constant d is the embedding dimension. If each recurrent iteration can be computed by a circuit of size C(n) < 20/2, and if the RNN produces the final answer after ranning most A steps of Co'l' for some constant A >> O, then there exist Transformers with O(log ni-bit precision, (C((n)) parameters and word embedding [W(E) 0(n+ns) xd] that can produce the same final answer after running (C(n) + +1mm steps of COT.
ENHANCING THE IN-CONTEX: RETRIEVAL CAPABILITY CLOSES THE GAP
In Section 4.2 we show that RNNs are deficient at In-Context Retrieval, hence leading to a significant representation gap with Transformers. In this section, aim understand: if we enhance the In-Context Retrieval capability of RNNS, do RNNs remain to have any representation gap with Transformers? We answer this question by examining two representative approaches to enhance the In-Context Retrieval capability, one explicit and one implicit, and show that both ways can close the representation gap between RNNs and Transformers in solving algorithmic problems.
5.1 EXPLICIT RETRIEVAL THROUGH REGULAR EXPRESSION
First, we explore the power of RNNs with Retrieval Augmented Generation (RAG), which gives an LM the capability retrieve relevant information to assist generation. In our context, specifically interested in allowing LMs to call functions to retrieve information from their context, which we call In- in-Context Retrieval Augmented Generation (In- (In-Context RAG).
Perhaps the simplest form of In- In-Context RAG is to allow the model to invoke function calls to associative recall, but show in Appendix E.8 that this is not enough to close the representation
7

---

Appending One Transformer Layer Published as a conference paper at ICLR 2025
Kemaver
Ceggest
Search Query
Soarch Resuk
Scacch
EalSpach
RNW
RNNN
RNN
Allmnno MET
SuntSearch
ExtSearch
In-Context RAG
Figure 3: Illustration of the two approaches to enhance the In- mn-Coneext Retrieval capability of RNNs: In-Context RAG (Ieft) and appending one Transformer layer (right).
gap between RNNs and Transformers. In this light, we need to consider more general form of In-Context Retrieval capability. We then turn to implement In-Context RAG by enabling an LM to perform regular expression matching, since the regular expression is flexible primitive that can be used to describe a wide range of retrieval tasks and can be implemented efficiently on modern hardware. To be precise, when the context is Saring-Starttearrch>p Endtarr <EndSearch>, we evaluate the Python code re. search (pattern, string) group (1) to get the result of the regular expression matching. Then we append the result the end of the context. See Definition 10
The following theorem shows that In-Context RAG with regular expressions is powerful enough for RNNs to solve the In-Context Retrieval problems in Section 4.2. with CoT.
Theorem 5.1.
For task T E {Index. AR, c-gram retrieval, Counting, IsTree). there exist constant- size Linear RNNs that can solve T with In-Context RAG. For T other than IsTree, (()) steps of CoT is required and for IsTree, O(n log steps of CoT is required.
More generally, the following theorem further shows that In-Context RAG empowers constant-size RNNs to simulate polynomial-time Turing machines.
Theorem 5.2. Given A, B > O, for all polynomial-tin Turing machines T E TIME(n with B states and vocabulary size B. there exist Linear RNNs with special symbols, O(Alog n) n)-bit precision and memory, and O(AB²) parameters that can output the result of by running O(n steps of CoT with In-Context RAG.
By introducing In-Context RAG, RNNs are able to use regular expressions to retrieve information from a distance and put relevant information together. Then RNNs can use perform more complex operations locally and output tokens to be retrieved later.
As final note, our focus here to understand the representation power of RNNs given an appropriate RAG, but not to propose method that immediately leads to practical applications. While the above results show that In-Context RAG can close the representation gap between RNNs and Transformers in solving algorithmic problems, limitation here is that In-Context RAG is not an immediate practical solution, as there is no existing training data for this In-Context RAG.
5.2
IMPLICIT RETRIEVAL BY APPENDING JUST ONE TRANSFORMER LAYER
Since Bahdanau et al. (0166), attention mechanisms have been understood as a form of compensation for the fixed memory size of RNNs, allowing the model attend the entire context and retrieve information. More recently, several hybrid Transformer-Ree RoNN architectures, such as Jamba (Lieber
al. 2024 SPADE (Zuo et al.) 2022), and Block state Transformers (Fathi et al. 2023). have shown significant performance improvements over RNNs in practice.
Now we demonstrate that adding attention layers as form of implicitly enhancing the retrieval capability can close the representation gap. We study a simple hybrid architecture that combines RNN and Transformer by appending just one Transformer layer to the RNN output (Appendix B). The following theorem shows that it can solve the In-Context Retrieval problems in Section 4.2
Theorem 5.3. For task T E (Index, AR, c-gram retrieval, Counting, IsTree). there exist constant- size hybrid Linear RNNs that can solve T. For T other than IsTree, no CoT is required and for IsTree, O(n logn) steps of CoT is required.
Further, we show that this hybrid architecture is powerful enough to simulate all polynomial-inn -ime Turing machines with COT. The proof idea is similar to the case of using In-Con Context RAG but now use attention retrieve information instead.
8

---

Published as conference paper at ICLR 2025
0.5M HYBRID
1M HYBRID
2M HYBRID
Val Accuracy
0.8-
0.4 16
128
16 32 64
128
16
128
Figure 4: We train Mamba with additional Transformer layer with frozen word embedding and decoding head of three different model sizes (0.5M, IM, 2M) on IsTree with three different sizes of graph (16, 32, 64) under three different setups. Vanilla means the model directly predicts the label. COT means the model will generate chain-of-thougge process based on DFS (see Algorithm I before prediction. Retrieval means the model will generate the chain of search queries and reasoning before prediction (see Algorithm 2
Theorem 5.4. Given A, B O, for all polynomial-timg -ime Turing machines T E TIME(n with B states and vocabulary size B, there exist constant-sizz hybrid Linear RNNs with B special symbols, O(A log n.) n-iiit precision and memory, and O(AB²) parameters that can output the result of T by running O(nA) steps of Co'.
EMPIRICAL VALIDATION
We validate theoretical findings through synthetic and natural language experiments: CoT alone cannot close the performance gap between RNNs and Transformers, but enhancing the in-context retrieval capability narrow the gap.
6.1 VALIDATION ON SYNTHETIC TASK: ISTREE
First, we validate our theoretical findings on the IsTree task introduced Section 4.2.2) We generate synthetic dataset for this task, where each data point is a tokenized graph concatenated with an answer of YES or NO, potentially with a CoT inserted in between. We then train RNNs and Transformers as language models this dataset autoregressively.
Experiment Details. To generate the graph, we follow the procedure described in the proof of The- orem 4.7 (see Figure (6) The CoT data is generated using Algorithm [ and the retrieval data is generated using Algorithm| 2 For the CoT model, we decode the reasoning path during inference time until we reach the first YES or NO up to max token limit greater than the length of all ground truth COT. For the data points that the model fails to give a prediction, we assume the model gets it correct with 0.5 probability. For the retrieval task, we omit the explicit format of the regular expression and only ask the model to generate the vertices and special tokens in the regular expression to shorten the length of the input sequence. The reported accuracy is calculated over validation set of 5000 samples using the last iteration of the model. We defer the experiment details to Appendix D The results are shown in Figures 2 and
Results with COT. Without CoT, both the Transformers and RNNs model cannot learn to solve the IsTrre problem. CoT improves the performance of both Transformers and RNNs but the RNNs per- formance degrades sharply the graph size increases and the Transformers consistently outperforms RNNs in this case. This consistent with our theory that CoT can improve the expressiveness the RNN models but the expressiveness still not enough solve IsTree (see Theorems and
Results with In-Context RAG. In-Context RAG allows all the models to reach near- near-perfect accuracy. This is consistent with theory that retrieval augmentation via regular expression improve the expressiveness of the RNN models to solve algorithmic tasks (see Theorems 5.2 E.33 The hybrid model Mamba model with one additional Transformer layer on the top) shows performance on par with the Transformer, which is consistent with our theory (see Theorems 5,4 and E.36
6.2 VALIDATION ON HOTPOT-QA
We further conduct experiments open-source LLMss based Transformer, Mambda, and hybrid architecture to show that for tasks that require stronger in-context retrieval capability, RNNs suffer more performance degradation compared to Transformers.
Experiment Details. We use Phi-1.5 1.3B (Li al. 2023) as our Transformer model. We further use two Mamba and Transformer- Mamba hybrid models distilled from Phi-1.5 with approximately the

---

Published as a conference paper at ICLR 2025
Retrieval Required
Retrieval Free
1.0
1.0
0.9
0.9
Accuracy
0.8
Accuracy
0.8
0.7
0.7
Transfommer
Transforme
0.6-
0.6
Hybrid
Hybrid
0.5
RNN
0.5
RNN
0.4
0.4
2
4
6
8
2
6 Number Paragraphs
Number of Paragraphs
Figure 5: Evaluation of Multi-Document Reasoning on Hotpot-QA. paragraphs given the model before the question and 2 among which contains required information to answer the question. (Left). When provided paragraphs are given in random order, the task requires stronger in-context retrieval capabilities when n increases. RNNs' performance drop more sharply, as predicted by theory. (Right). When the required paragraphs
always at the end of the provided paragraphs, the retrieval difficulty significantly decreases and all the models have stable performance with respect to n.
same parameters' sizes (Bick et al. 2024) as our RNN and hybrid models. To test our theory, we use the Hotpot-QA (Yang al. (0188) 8) dataset. In the validation set of Hotpot-QA, each question is accompanied by a set of different related paragraphs from Wikipedia. 2 of these paragraphs contain useful information answer this question. The model needs retrieve the correct paragraphs and reason based on these paragraphs to get the answer to the question.
HotpotQA with Controlled Retrieval Difficulty. We consider the following version of Hotpot-QA with enhanced retrieval difficulty: We choose n paragraphs containing the correct paragraphs and order them randomly before the question. The model needs to answer the question based on the given paragraphs after CoT reasoning. This design allows us to use the hyperparameter n to control the difficulty of in-context retrieval in this task, with larger n corresponding to higher difficulty. We test our models under 4-shot setting with Chain-of-Thougge performance even if n = 2. This is mostly due to their different capabilities to follow instructions and in in-context examples. To mitigate this effect and highlight the impact of retrieval, on a subset of 350 samples of the validation set where all the models can answer correctly given the correct paragraphs. This ensures perfect accuracy when = 2 in both settings. We vary in {2, 4, 6, 8) and the result shown in Figure
The model we tested has varying
only test
RNNs' performance drops sharply with increased retrieval difficulty. While all the model's performance drops with increased retrieval difficulty (Figure left), the RNN model has the largest drops. The hybrid model with only 4 attention layers performs significantly better than the RNN model. This validates our theory that RNN architectures' limited in-context retrieval capabilities will impact their performance in reasoning and hybrid architecture a potential solution this limitation.
Deconfouding the effect of context length. The context length also increases with the number of provided paragraphs and can be a potential confounder. We then experiment with controlled group: after choosing the same set of n paragraphs, we always order the correct paragraph at the end, which significantly simplifies the in-context retrieval process. In this case, all the models have stable performance when the number of paragraphs increases.
7 CONCLUSION AND DISCUSSION
This paper studies the representation gap between RNNs and Transformers on algorithmic problems. We prove that while CoT improves RNNs, insufficient to close the gap with Transformers due to the key bottleneck of RNNs' inability perform in-context retrieval. We show that adopting In-Context RAG or appending a single Transformer layer to RNNs can enhance their in-context retrieval capabilities, thus closing the representation gap with Transformers.
Considering the inherent limitations of RNNs in in-context retrieval, it is crucial to explore strategies that enhance this capability. This includes developing more effective in-context RAG methods tailored for RNNs, beyond the theoretically simple but impractical approach using regular expressions discussed here. Further research should also focus architectural modifications that strike balance between representational power and memory efficiency. In this study, we demonstrate that adding just single Transformer layer significantly enhances the representational capacity of RNNs. Future work could investigate the optimization and generalization of various hybrid architectures to achieve a more comprehensive understanding.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Ekin Akyürek, Bailin Wang, Yoon Kim, and Jacob Andreas. In-context language learning: Architec- tures and algorithms, 2024.
Silas Alberti, Niclas Dern, Laura Thesing, and Gitta Kutyniok. Sumformer: Universal approximation for efficient transformers, 2023.
Noga Alon, Yossi Matias, and Mario Szegedy. The space complexity of approximating the frequency moments. In Proceedings of the twenty-eighth annual ACM symposium Theory of computing, pp. 20-29, 1996.
Simran Arora, Sabri Eyuboglu, Aman Timalsina, Isys Johnson, Michael Poli, James Zou, Atri Rudra, and Christopher Ré. Zoology: Measuring and improving recall in efficient language models. arXiv preprint vrXiv:2312.04922 2023.
Jimmy Ba, Geoffrey E Hinton, Volodymyr Mnih, Joel Z Leibo, and Catalin Ionescu. Using fast weights to attend to the recent past. Advances in neural information processing systems, 29, 2016.
L. Babai, L. Fortnow, N. Nisan, and A. Wigderson. Bpp has subexponential time simulations unless exptime has publishable proofs. In [1991] Proceedings of the Sixth Annual Structure in Complexity Theory Conference, pp. 213-219, 1991. doi: 1t..1009/CT.199.1660000
Dzmitry Bahdanau. Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate, 2016.
Maximilian Beck, Korbinian Pöppel, Markus Spanring, Andreas Auer, Oleksandra Prudnikova, Michael Kopp, Glinter Klambauer, Johannes Brandstetter, and Sepp Hochreiter. xlstm: Extended long short-term memory, 2024.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages, 2020.
Aviv Bick, Kevin Li, Eric Xing, J. Zico Kolter, and Albert Gu. Transformers ssms: quadratic knowledge to subquadratic models, 2024. URL
Distilling
10189
Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Mil- lican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre. Improving language models by retrieving from trillions of tokens, 2022.
Aydar Bulatov, Yuri Kuratov, and Mikhail S. Burtsev. Recurrent memory transformer, 2022.
Nelson Elhage, Neel Nanda, Catherine Olsson, Tom Henighan, Nicholas Joseph, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Nova DasSarma, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodde Danny Hernandez, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https:/transformercciccitts wtt.:/aiiss..b//222//aaaaeoooooo html.
Mahan Fathi, Jonathan Pilault, Orhan Firat, Christopher Pal, Pierre-Luc Bacon, and Ross Goroshin. Block-state transformers, 2023.
Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain thought: theoretical perspective. In Thirry-seventh Conference on Neural Information Processing Systems, 2023. httptpss///ppeeeeeww...o.. forum ld gHrADgAdYu
Daniel Y. Fu, Tri Dao, Khaled K. Saab, Armin W. Thomas, Atri Rudra, and Christopher Ré. Hungry hungry hippos: Towards language modeling with state space models, 2023.
11

---

Published as a conference paper at ICLR 2025
Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401 2014.
arXiv preprint
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023,
Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. In International Conference on Learning Representations, 2022. URL https //openreview. net/ wttp://ooormmmiuuoooooooo.
Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang. Realm: Retrieval- augmented language model pre-training. 2020.
Michael Hahn. Theoretical limitations of self attention in neural sequence models. Transactions of the Association for Computational Linguisticx, 8: 8:156-171, December 2020. ISSN 2307-387X.
doi: 10.1 20.11622acl_a00032000 URL http:: /tt:://dxxxo doi www.://ooo.ooooooo/ooo/o//////
Jie Hao, Xing Wang, Baosong Yang, Longyue Wang. Jinfeng Zhang, and Zhaopeng Tu. Modeling recurrence for transformer. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 1198-1207, Minncapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v1/ N19-1122 1122. URL htLtttts::/ http:///caaattooooooyoo
9-122
Yiding Hao, Dana Angluin, and Robert Frank. Formal language recognition by hard attention transformers: Perspectives from circuit complexity, 2022.
Monika Rauch Henzinger, Prabhakar Raghavan, and Sridhar Rajagopalan. Computing on data streams. External memory algorithms, 50:107-118, 1998.
Geoffrey Hinton and James Anderson. Parallel models of associative memory: updated edition.
Psychology press, 2014.
John Hopfield. Neural networks and physical systems with emergent collective computational abilities. Proceedings of the national academy of sciences, 2554-2555-5555 1982.
Samy Jelassi, David Brandfonbrenerr, Sham Kakade, and Eran Malach. Repeat after me: Trans- 2024.
formers are better than state space models at copying. arXiv preprint
Zhengbao Jiang. Luyu Gao, Zhiruo Wang, Jun Araki, Haibo Ding, Jamie Callan, and Graham Neubig. Retrieval as attention: End-to-end learning of retrieval and reading within a single transformer. In Yoav Goldberg, Zornitsa Kozareva, and Yue Zhang (eds.), Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pp. 2336 2349, Abu Dhabi, United Arab Emirates, December 2022. Association for Computational Linguistics. doi: 10. 10.18653/1120022. emnlp- manlp-main.1ie main. 1499. URL https: http::://aaaaaaataaaaaee org 2002 emnlp-man 149
Richard M. Karp and Richard J. Lipton. Some connections between nonuniform and uniform complexity classes. In Proceedings of the Twelfth Annual ACM Symposium Theory of Computing, STOC '80, pp. 302-309, New York, NY, USA, 1980. Association for Computing Machinery. ISBN 0897910176, 76. doi: 10. 01445//00141.046677 URL http:::////oooo//////// 0tt../..1/....1//000.. 2145/800101 804678
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers RNNs: Fast autoregressive transformers with linear attention. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 5156-5165. PMLR, 13-18 Jul 2020. URL html
https: https://proceecingg mlr http://pe.eeee/////a//aaa/oppp/p
Takeshi Kojima, Shixiang Shane Gu, Machel Reid, Yutaka Matsuo, and Yusuke Iwasawa. Large language models are zero-shot reasoners, 2023.
Yuri Kuratov, Aydar Bulatov, Petr Anokhin, Dmitry Sorokin, Artyom Sorokin, and Mikhail Burtsev. In search of needles in 11m haystack: Recurrent memory finds what lmms miss, 2024.
12

---

Published as a conference paper at ICLR 2025
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report, 2023. URL hatts:://axxvv.org abs/2305 05463
Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems, 2024.
Zhong Li, Jiequn Han, Weinan E, and Qianxiao Li. On the curse of memory in recurrent neural networks: Approximation and optimization analysis. In International Conference on Learning Representations, 2021. URL http://///.porrrreeeeeeeee///rrs
Zhong Li, Haotian Jiang, and Qianxiao Li. On the approximation properties of recurrent encoder- decoder architectures. In International Conference on Learning Representations, 2022. URL
https://openrevier net forumrid XDIIvxxXDII
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: A hybrid transformer-mamba language model, 2024a.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: hybrid transformer-mambe language model, 2024b. URL /Rttt/t:://axxx
org/ abs//403..344// 2403.19887777 19887
Binghin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference on Learning Representations, 2023. URL https::///.mm....t.. htttps://openzeviiee. http://ww......rroo.oeeeeeeee///
Shengjie Luo, Shanda Li, Tianle Cai, Di He, Dinglan Peng, Shuxin Zheng, Guolin Ke, Liwei Wang, and Tie- Yan Liu. Stable, fast and accurate: Kernelized attention with relative positional encoding, 2021.
Shahar Lutati, Itamar Zimerman, and Lior Wolf. Focus your attention (with adaptive iir filters), 2023.
William Merrill and Ashish Sabharwal. The Parallelism Tradcoff: Limitations Log-Precision Transformers. Transactions of the Association for Computational Linguistics, 11:531-545 06 2023. ISSN 2307-387X. doi: 10.1 10.1162/tacl_ a a_00562. /ttl:tttp://doo.orze tacl 00562]
William Merrill, Ashish Sabharwal, and Noah A. Smith. Saturated transformers are constant-depth threshold circuits, 2022.
J Ian Munro and Mike S Paterson. Selection and sorting with limited storage. Theoretical computer science, 12(3):31-3223 1980.
Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: Scratchpads for intermediate computation with language models, 2021.
Matanel Oren, Michael Hassid, Yossi Adi, and Roy Schwartz. Transformers are multi-state mns, 2024.
Jongho Park, Jaeseung Park, Zheyang Xiong, Nayoung Lee, Jaewoong Cho, Samet Oymak, Kang- wook Lec, and Dimitris Papailiopoulos. Can mamba learn how to learn? a comparative study on in-context learning tasks, 2024.
13

---

Published as a conference paper at ICLR 2025
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao. Xin Cheng, Michael Chung, Matteo Grella, Kranthi Kiran GV, Xuzheng He, Haowen Hou, Jiaju Lin, Przemyslaw Kazienko, Jan Kocon, Jiaming Kong, Bartlomiej Koptyra, Hayden Lau, Krishna Sri Ipsit Mantri, Ferdinand Mom, Atsushi Saito, Guangyu Song, Xiangru Tang, Bolun Wang, Johan S. Wind, Stanislaw Wozniak, Ruichong Zhang, Zhenyuan Zhang, Qihang Zhao, Peng Zhou, Qinghua Zhou, Jian Zhu, and Rui-Jie Zhu. Rwkv: Reinventing rnns for the transformer era, 2023.
Hao Peng, Nikolaos Pappas, Dani Yogatama, Roy Schwartz, Noah A. Smith, and Lingpeng Kong. Random feature attention, 2021.
Michael Poli, Stefano Massaroli, Eric Nguyen, Daniel Y. Fu, Tri Dao, Stephen Baccus, Yoshua Bengio, Stefano Ermon, and Christopher Ré. Hyena hierarchy: Towards larger convolutional language models, 2023.
Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. Samba: Simple hybrid state space models for efficient unlimited context language modeling, 2024. URL BtLhttss
// rxiv. org/ abs/ /200244662200 07522
Ohad Rubin and Jonathan Berant. Long-range language modeling with self-retrieval, 2023.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers, 2023.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarith- mic depth, 2024.
Jianlin Su, Murtadha Ahmed, Yu Lu, Shengfeng Pan, Wen Bo, and Yunfeng Liu. Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 568: 127063, 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma. Yuqing Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor to transformer for large language models, 2023.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaci, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, Dan Bikel, Lukas Blecher, Cris- tian Canton Ferrer, Moya Chen, Guillem Cucurull, David Esiobu, Jude Fernandes, Jeremy Fu, Wenyin Fu, Brian Fuller, Cynthia Gao, Vedanuj Goswami, Naman Goyal, Anthony Hartshorn, Saghar Hosseini, Rui Hou, Hakan Inan, Marcin Kardas, Viktor Kerkez, Madian Khabsa, Isabel Kloumann, Artem Korenev, Punit Singh Koura, Marie-Anne Lachaux, Thibaut Lavril, Jenya Lee, Diana Liskovich, Yinghai Lu, Yuning Mao, Xavier Martinet, Todor Mihaylov, Pushkar Mishra, Igor Molybog, Yixin Nie, Andrew Poulton, Jeremy Reizenstein, Rashi Rungta, Kalyan Saladi, Alan Schelten, Ruan Silva, Eric Michael Smith, Ranjan Subramanian, Xiaoqing Ellen Tan, Binh Tang, Ross Taylor, Adina Williams, Jian Xiang Kuan, Puxin Xu, Zheng Yan, Iliyan Zarov, Yuchen Zhang, Angela Fan, Melanie Kambadur, Sharan Narang, Aurelien Rodriguez, Robert Stojnic, Sergey Edunov, and Thomas Scialom. Llama 2: Open foundation and fine-tuned chat models, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, E. ukasz. Kaiser, and Illia Polosukhin. Attention all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Ad- vances in Neural Information Processing Systems, volume 30. Curran Associates, Inc., 2017. URL http::::///prrcceeennnne.eeerrii /tll////e222244470ee000111 Tal.://fee4444eeee111o0ooaaaaaaa
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, Garvit Kulshreshtha, Vartika Singh, Jared Casper, Jan Kautz, Mohammad Shoeybi, and Bryan Catanzaro. An empirical study of mamba-based language models, 2024. URL htttt:tts//aaaae wwtp://gggga/////////e 07887
Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, 2020.
Xuezhi Wang and Denny Zhou. Chain-of-thought reasoning without prompting, 2024.
14

---

Published as a conference paper at ICLR 2025
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Brian lchter, Fei Xia, Ed Chi, Quoe Le, and Denny Zhou. Chain-of -f-thuuhtt prompting elicits reasoning in large language models, 2023.
David J Willshaw, O Peter Buneman, and Hugh Christopher Longuet-Higgins. Non-holographic associative memory. Nature, 127):997)960-962 1969.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks, 2023.
Kai Yang, Jan Ackermann, Zhenyu He, Guhao Feng, Bohang Zhang, Yunzhen Feng, Qiwei Ye, Di He, and Liwei Wang. Do efficient transformers really save computation?. 2024.
Zhilin Yang, Peng Qi, Saizheng Zhang, Yoshua Bengio, William W. Cohen, Ruslan Salakhutdinov, and Christopher D. Manning. Hotpotqa: dataset for diverse, explainable multi-hop question answering, 2018.
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention networks can process bounded hierarchical languages, 2023.
Simiao Zuo, Xiaodong Liu, Jian Jiao, Denis Charles, Eren Manavoglu, Tuo Zhao, and Jianfeng Gao. Efficient long sequence modeling via state space augmented transformer, 2022.
15

---

Published as a conference paper at ICLR 2025
CONTENTS
Introduction
2
Related Works
3
3
Preliminaries
How Well Does CoT Improve the Representation Power of RNNS?
CoT Strictly Improves RNNs
CoT Cannot Close the Representation Gap with Transformers
4.2.1
Simple Problems In-Context Retrieval
4.2.2
Limitations of RNNs Beyond Simple In-Context Retrieval Problems
4.2.3
Transformers (with CoT) are Strictly More Expressive Than RNNs
5
Enhancing the In-Context Retrieval Capability Closes the Gap
5.1
Explicit Retrieval Through Regular Expression
5.2
Implicit Retrieval by Appending Just One Transtormer Layer
6
Empirical Validation
6.1
Validation on Synthetic Task: IsTree
6.2
Validation on Hotpot-QA
7
Conclusion and Discussion
10
Extended Related Works
18
Additional Definitions
19
B.1 Reasoning Tasks on Graphs.
19
B 2
More on Numeric Precisions.
19
Models
19
B.4 Language Models for Reasoning,
21
Omitted Discussion on In-Context Retrieval Examples
22
Omitted Experiment Details
22
Omitted Proof
23
E.1
Building Blocks of FENs Construction
23
Building Blocks of Transformers Construction
24
Building Blocks of RNNs Construction
31
Proof of Theorem 4.6
32
Proof of Theorem 4.1
33
16

---

Published as a conference paper at ICLR 2025
E,6
Proof of Theorem 4.7
34
E.6.1
Proof of Lemmal E.26
34
E,6.2
Proof of Lemma E.27
35
Proof of Theorem
37
E,8
AR function calls is not enough for closing the representation gap
.........
39
Proof of Theorem
40
E.
10
Proof of Theorem
40
Proof of Theorem 5.2
42
12 Proof of Theorem E.35
43
13 Proof of Theorem E 36
43
14 Proof of Theorem 5.4
44
17

---

Published as a conference paper at ICLR 2025
A EXTENDED RELATED WORKS
Comparison Between Transformers and RNNs (without CoT). A line of works focused on the comparison between RNNs and Transformers in terms of recognizing or generating formal languages (Bhattamishra et al., 2020; Hahn| 2020; Merrill et al. 2022). These works show that the lack of recurrent structure in Transformers makes them fail to recognize some formal languages that RNNs can recognize. However, Liu et al. (023); y Yao et al. (202)); Hao et al. (2022) show that such limitation can be mitigated when consider bounded length of input or bounded grammar depth. Our work differs from these works in that we consider the expressive power of RNNs and Transformers with CoT and show that in this case, the gap between RNNs and Transformers is one-sided (Theorem (.8). Another thread of recent works focuses on understanding why RNNs or SSMs are inferior to Transformers (Arora et al.) 2023 Jelassi et al. 2024 Sanford et al.) 2023 2024).
Our work differs from these prior works from two perspectives: (1) we identify the common reason behind a wide range of tasks where RNNs fall behind Transformers; (2) we show two potential solutions that can bridge the representation gap by addressing this common reason (the incapability to retrieve from context).
Retrieval Augmented Generation. Our work proposes to use retrieval augmentation to close the representation gap between RNNs and Transformers. This is consistent with the recent trend of retrieval augmented generation (Guu et al. 2020; Borgeaud et al. 2022, Rubin & Berant 2023). Empirically, retrieval augmented generation has been shown to improve the performance of recurrent models in various tasks (Kuratov et al. 2024 Akyürek et al. 2024) and our work provides theoretical foundation for this phenomenon. Our work also shows that an attention layer can be used to simulate the retrieval process, which is consistent with the finding that attention can improve the performance of RNNS (Vaswani et al. 2017; Arora et al. [2023] Park et al.) 2024 Peng et al., 2023 Hao et al.,
2019). It has also been shown empirically that attention can be used simulate complex retrieval process (Jiang et al. 2022).
Comparison Between Transformers and RNNs Prior work (Arora et al., 2023) has shown that input-independent gating SSMs are inferior to Transformers in the task called associative recall (Willshaw et al. 1969 Hoptild 1982, Hinton & Anderson 2014). The task requires the model to recall a previously seen pattern given a partial input. They show that iuput-dependers gating SSMs have better performance in associative recall and also propose a hybrid architecture that combines input-independent state space machines with attention to achieve better performance. Our work differs from this work in the following ways: (1) Our work studies associative recall from a theoretical perspective and proves formal lower bounds on the memory size of RNNs necessary for solving associative recall and other retrieval tasks; (2) We also study hybrid architectures but we provide a proof that appending single Transformer layer to RNNs can make them expressive enough; (3) Our theory applies to not only input-independent gating SSMs but also all RNNs with bbit memory.
Prior work (Jelassi al. 2024) proves representation gap between RNNs and Transformers in repeating a long sequence, which can be seen as a retrieval task. They show that RNNs have difficulty performing the task due to their limited memory. Our work further proves that RNNs are limited in solving many other retrieval tasks, even with COT. Technically, a key ingredient in their proof counting argument on the output sequence show limited memory size is not enough to produce too many different output sequences, but our proof can handle retrieval tasks that only require outputting a single token.
Notably, Sanford et al. (223) apply communication complexity to prove circuit size or memory size lower bounds for RNNs and Transformers on the task of sparse averaging. [Sanford al. (2024) extend this technique to another task called hopa, generalization of the associative recall task. Our technique is similar to theirs since our proof is also based on communication complexity. But we consider a broader range of tasks including seemingly irrelevant reasoning tasks such as IsTree, and further explore various ways to close the representation gap.
Notably.,Sanford
Alberti al., ) studies the universal approximation power of RNNs. They show that the upper bound of the approximation power of linear RNNs will be constrained by the dimension of the hidden states. Their works on the high level are consistent with our findings but are not directly comparable because we considering finite precision compute models with the assistance of CoT In-Context RAG.
Representation Theory of RNNs. Another line of works (Li et al. 2021 2022:
18

---

Published as a conference paper at ICLR 2025
B ADDITIONAL DEFINITIONS
We will now define some definitions used in the proofs.
B.1 REASONING TASKS ON GRAPHS.
When resaoning graphs, without otherwise specified, we will use n as the number of vertices and m as the number of edges. Without loss of generality, we will assume the vertices are labeled by [ni].
We will focus decision problems on graphs, which are defined follows: Definition B.1 (Decision Problem Graphs). A decision problem graphs is function (YES, NO), where G is the set of all possible graphs.
We will use the following decision problem our main example: Definition B.2 (IsTree). IsTree (G) YES if G is tree, and IsTree(Gi) NO otherwise.
One can view IsTree as a minimal example of reasoning task2s. One of the classical solutions to IsTree is running Depth First Search and this algorithm takes O(n) time.
B.2 MORE ON NUMERIC PRECISIONS.
We will ROUND(a, p) to denote the rounding function that rounds æ to the nearest number in Ep. We will assume p is an odd number without loss of generality.
Rp
(2bp
Σ 6,2(P 122--11/2-))
vi E lp), b 10.1] E0.11
(1)
For calculation over R.p we will assume the calculation exact and the result rounded to Hp the end, that is, for operator cD we will have
ROUND (z, p) @G ROUND(V. -ROUND (ROUND(z; p) @ ROUND(Y, p).p))
We will additionally define Z, the set of all integers that can be represented by p-bit floating point numbers. We will define 1/ 1/mm] as the set of unit fractional { [e[mm)- Further, we will define ROUND(1/(m),p) as the rounding of 1/[m] to Ep. We will additionally define for any real number x E {0, 1). next(x)=
m+I where m arg minkez |z
B.3 MODELS
Tokenization. To tokenize a string, we will tokenize all the words separated by the space character into sequence of tokens. To tokenize graph G, we will order its edges E = {(1, v.) < 4 randomly and tokenize into the following string:
Tokenize(G) = {<>>> u1, U1, Umm ~o
(2)
We hereby assume there are constant more special tokens that are not the same as any number token, which are listed below:

; >>: the first special token, indicating the start of sentence.
~; the second special token, indicating an edge.
YES: the third special token, indicating the answer is yes.
NO: the fourth special token, indicating the answer is no.
SStatteeacc>>: the fifth special token, indicating the start of search query.
<EndSeach>: the sixth special token, indicating the end of a search query.
We will denote the total number of special tokens ng and the total vocabulary size [V| = n will further define the detokenization function Detokenize,
Detokenize(S) = "S, S₂ Sy". Here each S; is either number special token, which will treat word.
19

---

Published as a conference paper at ICLR 2025
Embedding Functions. We will use d to denote the dimension of the embedding and w, to denote the i-th coordinate vector in R
We will separate the embedding function into two parts: the word embedding and the position embedding. For the word embedding. we will use iw E R to represent the embedding of the vertice in the tokenization of the graph. For the k-th special token, will use 122+ to represent its embedding. For example, the embedding of ~ is W2. We will denote the word embedding matrix as ((F) E [
For the position embedding, we will use Iwa represent the position embedding of the 1-th token in the tokenization of the graph, which is hyperparameter. The final embedding of any token sequence is the of the word embedding and the position embedding. We will use Emb to denote the embedding function.
This embedding function will be fixed and shared across all models consider in this paper and will not be learned during training, hence we will not consider it as part of the model parameters.
Language Modeling. We use V and V to denote the set of all finite sequences and all non-emple finite sequences of tokens in V, respectively. We study language models that predict the next token given a prefix of tokens. For this, we define a language model (LM) as a function M V+ + Pv that maps non-empty sequence
probability distribution over the next token, where Pv the probability simplex over V. We specifically study the case where the language model is realized by deep neural networks: first map the input sequence into a sequence of embeddings Emb(S), and then apply neural network, such as a Transformer or RNN, to process the embeddings and output the probability distribution. We would call series of parameterized models with increasing input size a family of models.
Transformer We will first define the Transformer architecture used in the theoretical analysis in this paper.
Definition B.3 (Transformer Block). Let X E Rdxd be the input matrix, where I the sequence length. The output of Transformer block f is defined as:
f(x) = x + A(X) + g(X + A(X)), (INIKA))) w tc). Vd
A(X) Σ w Xsoftmax
where is column-wiss ReGLU feed-forward network ! with width w and output dimen- sion d, A the scaled dot-product attention, softmax the column-wise softmax function, W(KA) W(Q.)) W(Vh) are the learnable parameters and И the number of heads, and
0
0
07
0
0
C
E R!x! is a mask to prevent the attention from attending to future
0
tokens.
In the context of language modeling, given a sequence of tokens S, Transformer T(S) defined as:
Definition B.4 (Transformer). Let S E IVI't be the tokenized input sequence, the output of Transformer is defined as:
T(S) softmax (W(E) (fL (... f1 (Emb (8))))
(4)
where softmax is the column- wise softmax function, fi is the i-th Transformer block. We will call the i-th Transformer block the i-th layer of the Transformer and denote its forward layer and attention layer as di and A; respectively.
RRGGUU means (x) ReL + b1) (War by), this is surrogate for the commonly used SwiGLU activation and allows the model perform multiplication of coordinates.
20

---

Published as a conference paper at ICLR 2025
RNNs We formally define the RNN architecture here.
Definition B.5 (RNN), An RNN architecture is characterized by two functions: state transition function t + (R x R + R and output function : θ - (RR - 1Ry), where A is the dimension of the state and θ is the parameter space. Let S E |V be the input sequence, the output of recurrent neural network with parameter θ is defined as:
Ra(S) softmax oo()))
Ak E [1], 8/: to ( Emb(S),,
where Su E A is a vector determined by and ( is the word embedding matrix. We will omit the subscript when it is clear from the context.
We define the following RNN architecture for the proof of our upper bound.
Definition B.6 (RNN block). A Linear RNN block defined as follows:
f(X) = + LU(X) + g(X LU(X)),
where g is a column-wise ReGLU feed-forward network with width w and output dimension d and LU is linear unit, defined as
hg 0, h Ah,,-1 + BX,,,, Lu) h;
Definition B.7 (Linear RNN). A Linear RNN is recurrent neural network
R(S) = softmax (W(E) (fz (... (Emb (5))))
(5)
where softmax is the column- softmax function, fs the i-th Linear RNN block. We will call the i-th Linear RNN block the i-th layer of the Linear RNN and denote its feed-forward layer and linear unit layer as and T.U, respectively.
Hybrid RNNs.
We formally define the hybrid RNN architecture here.
Definition B.8 (Hybrid RNN). A hybrid RNN is model that consists of an RNN with transition and output function t. o and one Transformer layer f. the output of the RNN is used the input of the Transformer layer and the output of the Transformer layer is used to produce the next token. Concretely, given the input sequence Sin, the output of the hybrid architecture is:
H(S) = softmax (wr))) (lo(s))RkEH))
Vk [1], 8k t($a-1,, Emb(S)...)
B.4 LANGUAGE MODELS FOR REASONING.
Chain of Thought. We will now define how we use language models to solve reasoning tasks utilizing the following technique called chain of thought.
Definition B.9 (Chain of Thought). Given language model, M with vocabulary V and the tokenized input sequence S. E [V]',, chain of thought (CoT) generates the following sequence tokenized sequence:
So Sin: next
(6)
8
arg max M (S,)(i) jEV
(7)
S S, sp next Vi 0.
(8)
The process terminates at S; when arg maxjev M(S) [s] is YES or NO. The language model can solve the reasoning task within T steps of CoT if the process terminates at S, where i < T and the final output is correct. We will call the special case where the language model solves the reasoning task within 0 steps of CoT solving the reasoning task without COT.
21

---

Published as a conference paper at ICLR 2025
Retrieval Augmentation. We will show in this paper that retrieval augmentation is necessary technique to solve reasoning tasks for recurrent neural networks. We will define retrieval augmentation as follows:
Definition B.10 (Retrieval Augmented Generation). Given an LM M with vocabulary V (containing two additional special tokens, <StartSeeach>> and <EEndearch>> and the tokenized input sequence S E [V|',, the LM M with In-Context RAG generates following sequence of tokenized sequence:
So = Sms
next 8
arg max M (S))
JEV
S i S
if sToRt * <EndSearch>>>
S
S + stext + RETRIEVE (S,),
otherwise.
Here RETRIEVE looks for the last occurrence of <StartSearch> at position I, and <EndSearch> S at position L, and treat Detokenize(St,... ) as a regular expression, where Detokenize maps the tokenized sequence back to the string, inserting space between every pair of adjacent tokens. The algorithm then runs a regular expression matching on Detokenize(S1..-1 substring, and returns the first capturing group according to the regular expression (i.e., content embraced by pair bracket in the regular expression). While there are many grammar standards of regular expressions, we adhere the standard specified in the re library of Python. That is, evaluate the following Python code to get the result of the regular expression matching:
finds the first matching
search (pattern, string) . group (1)
where Delokenize(St... (S,,,,, is the pattern and Detokenize((Sia.-. is the string.
We will note that assuming VV O(n) and every search query and the result is of length O(1), the regular expression evaluation can typically be evaluated in ((n) time.
IsTre In the context of language modeling, then the task of IsTree to determine whether is a tree by predicting a YES/NO token with or without COT. We use the following tokenization for the graph G:
write the graph G as sequence of tokens, and
Tokenize(G) = {< u1, ~", U1, U22 U2,---. Mm, ~ )m
(9)
where <>> and ~ are two special tokens representing the start of the sentence and an edge, and uy, 2/. are numbers denoting the nodes of the graph.
C OMITTED DISCUSSION ON IN-CONTEXT RETRIEVAL EXAMPLES
Index and AR are perhaps the most basic problems in retrieval, where Index asks for retrieving token from the input sequence viewed as a linear array tokens, and AR asks for retrieving a token from the input sequence viewed extensively by different communities. Index is a classic problem in streaming and communication complexity (Munro & Paterson 1980). known to be impossible solve with o(n) bits of memory for streaming algorithms. AR has been regarded as a fundamental problem that an artificial intelligence system should be able solve (Willshaw et al. 1969. Hopfield 1982 Hinton & Anderson, 2014 Graves et al., [0144] Ba et al.) (016). In the context of LLMs, AR has been observed to correlate with in-context learning performance (Elhage et al.) 2021) ) and has also been used extensively as
associative array. These two problems have been studied
synthetic surrogate tasks for pretraining performance (Fu et al. 2023 Poli et al. 2023 Lutati et al. 2023). Besides Index and AR, e-gram retrieval is natural extension of AR to the case where the query key can contain multiple tokens: instead of retrieving a token given single-token key, c-gram retrieval asks for retrieving token when the given key is (c- (c-1)-gramt 1)-gram. This task has been studied empirically, but not theoretically in Jelassi et al. (2024). Counting is problem that asks for the number of occurrences of token, thereby testing the model's capability to retrieve some statisties of relevant information from the input sequence.
D OMITTED EXPERIMENT DETAILS
We train three different architectures: (1) LLaMA architecture (Touvron et al., 2023) representing Transformers, (2) Mamba architecture (Gu & Dao 2023) representing RNNs, and (3) Mamba with
22

---

Published as a conference paper at ICLR 2025
additional layer of LLaMA block representing hybrid architectures. Following our theory, freeze and weight-tie the prediction head and word embedding in all the models. For ease of training. different embedding function mapping i-th token to [sin( Tooonzza ). cos( Tuanss7 (]je(d/a) with being the number different tokens and use standard RoPE (Su et al. 2024) as position embedding. We train every model with least 1M samples to guarantee convergence using Adam with cosine learning rate. If the model doesn't I converge, we retrain using 5M samples, After grid search over learning rates, we train all the Transformer models with learning rates le-3 and the rest of the models with learning rates 3e-4. We run all the experiments on a server with 8 A100s and the estimated time reproduce the results is within 2 days.
E OMITTED PROOF
E. BUILDING BLOCKS FENS CONSTRUCTION
We will first show some basic operations that multiple layers of feedforward neural networks with ReGLU activation can perform that will be used the following proofs.
Lemma E.1 (Multiplication). Given two dimensions i,, i2, there exists parameter configuration of I-layer feedforward neural network with ReGLU activation that for any inpul IR and constant width, computes the product of 24, and Fz2-
g(x) = [ri x Zi,, OF..,,0.
Proof. We can construct the following feedforward neural network with ReGLU activation:
Wi [Ti,
0
0]
W22 = [itz Z12
0 .... 0]
Wyh 11i +
0 0
0]
by = b2 =by- O.
((x) (W())LU(W112) & ReLUU(22)))
[mi, Figs 0,....01 O)T
Lemma E.2 (Linear Operations). Given a linear transformation W E ddd there exists parameter configuration of 1-layer feedforwardd neural network with ReGLU activation and width d that for any input x E R.d. computes Wx.
Proof.
b1 = 1",b b2 = 0, by = =0,
W = 0, W2 W,W, = laxd
Lemma E.3 (Indicator). Given a constant integer B S d and a dimension i, there exists a parameter configuration of a 1-layer feedforward neural network with ReGLU activation and width 4 that for any input x E idd computes the indicator function of x; B when I; is an integer.
Proof.
- - ... 10,
W2 = 0, Wix = -1;; -x;
ReLU(W1æ +b,) (Re 0.5), ReLU(e, B + 0.5), ReLU(B - -r; ReLU(B =; + 0.4)
W3 10 -1 -1 _T_
23

---

Published as a conference paper at ICLR 2025
Then,
(0x) = (Wa b1)
= 10 (ReLU(Z; B - 0.5) ReI.U B + 0.5))
+ (ReLU(B
-06) ReLU(B 2; + 0.4)) 10
x 0+ 10 -0
if < B
10 x -0.++ x 10
if B,
10 -1 -1 + 10
if > B.
Lemma E.4 (Lookup Table). For constant B and is such that kB S d. given lookup table which key tuple of k integers bounded by B, and value is a scalar in Rp. there exists a parameter configuration of layer feedfooward neural network with ReGLU activation with width O(B*) that for any input Rd, computes the value of the lookup table at the key In Digs... Fig.
Proof. We can calculate vi, + B Fig + B2 Fig functions to get the value of the lookup table at the key Izy Fig....,,
Bk-1 x Fig and then scale B indicator
Lemma E,5 (Threshold). Given any threshold u and constant e > there exists a parameter configuration of a 1-layer feedforward neural network with ReGLU activation and width 2 that for any input x E Rd, computes the indicator function x; > u on z; E [-00, u e] U [u, oo].
Proof.
62 = 1",b, = [-u+ 0..5e] O,
W2 0, W/ [;
ReLU(WW2 + b1) =
ReLU(r; u + 0.5e)]
2
W3 = [ -1 11
Then,
g(x) = WsReLU(W1r+ bb) b, 2/e (ReLU(I, -u-e) ReLU a + 0.5))
50 if < cu-e,
11 if u
Lemma E.6 (Interval). Given any constant u, and > O, there exists parameter configuration of a I-layer feedforward neural network with ReGLU activation and width 4 that for any input x E RR., computes the indicator function E [u, v] on I, |0 e] U [u, v] [e, oc).
Proof. The interval function here can be written as the difference of two threshold functions. We can use the same construction as Lemma E.5 to approximate the indicator function 2; > u and + c and then take the difference.
E.2 BUILDING BLOCKS TRANSFORMERS CONSTRUCTION
We will show in this section some construction for basic functionality using Transformer Blocks. This construction will be used in the following sections to prove the main results.
We will always X I Idxx as the input to the Transformer Block, where d is the dimension of the input, and is the length of the sequence. We will first outline all the building functionality and then show how implement them.
24

---

Published as a conference paper at ICLR 2025
Definition E.7 (Copying Function), For integer s, index set I1, 12 [d - 20] satisfying [41 [22) a copying function COPY 11, Iz] satisfies the following, AX E p
1dxt
then
COPY (2, 12. 21, max{f--.0) COPY|s, 11 (x) 1gg. =0 vr [m]]
4k S [m]]
Definition E.8 (Counting Function). For index set 1, C Id 20), [1 = [T2l s 10 and index i,
counting function COUNT [2, satisfies the following, if vo E II U 12. k E [4, N E Z, and
X,,k + 0, then
42. Σn 1[x X12.) +1 Ak [1]
vk E [4]
44
Definition E.9 (Matching Function). For index set L1, 12, 13, 14 C [d- 20 20,,111 44 = 422 E 10, 44 =
[4/. matching function Match(/1: 12, 1s, L4] satisfies the following, Ve E 1 U12 E [0) Zp, then
Match [T, 12, 13, (1) = XIa
4k E E [1]
min(h) XI, = X12th {h |1,, = X1a.k3 #0
where k-
1, otherwise
Definition E.10 (Matching Closest Function). For index set I, Iz. Is. I4 C [d [0), |1| [/] s
10, |/2| /La.. matching closest function Match(/1, 12, Ia, Ia] satisfies the following, if Ve 1 12, k E [1, X., k Ep. then
MatchClosc [, I2, Is, [1) Niak Vk
[7]
where k* = max(h) | X = Xlahth th X11.h = X12.33 40
otherwise
Definition E.11 (Matching Nearest Function). For index set I1, Is, I1 C [d 20), 11 Iz < 10. |4ا| = Lal and index i, matching nearest function MatchNearest [|, 1,, 13. L. satisfies the following, Vu E T T2, E [1], Xo,k E Zp. then
MatchNearest [I, I2, Is, Ia XIsk Wk [1]
arg min 1 Ne,eh {h X1,h x12,k3 0 otherwise
where k"
Definition E.12 (Matching Next Function). Given any interger constant A, assuming p >> 10.4 log n, for index set /1, 12. Jz. 1 C [d 20), |41 = V2l S 10, 173 // and a special counting index a, matching next function MatchNext/AA 41 1z, Is. I1, a] satisfies the following, if X satisfies the following condition:
1. Ve E TU X,, E Zyp
2. XaA E ROUND(1/H (ID)
For any & E [], given any k 2 k. the counting index multiset Sş {Xa, XI X1. takes consecutive and disjoint values in ROUNDOOOOO/O/E [/]//////]) [p)].) p). that is, there exists uky Uk ROUND(1/[na]os p) such that S₁ = [ux, rz] ROOND(1/[mA)]
then, have
MatchNext(11, T2. 1s, 14, ( = Nla Wk E [1] where A:* = arg
be X x (1(1)) |Xa, - (Xa,k))
min
Now we will show how to implement these functions using Transformer Blocks. The construction here is motivated by the construction in Feng al. (203) with some modifications.
25

---

Published as a conference paper at ICLR 2025
Lemma E.13 (Copying Blocks). For integer s, index set 11, 12 C [d- 10] satisfying [41] [12
copying function COPY[s, I, Iz] can be implemented with 1 feedforward block g and 1 attention block A with 1 attention head. Formally, when Xd,x k, it holds that
(g) + (x0
Proof of Lemma E.13 have
We will use the feedforward block calculate X and 1 (Lemma E.1 and
(g(X) - x) = =
(g(X) - x) =1
1,d- Vi E 1d-1,d 23, (o(X) + X)1, Xi,K-
We will use X' to denote g(X) + X. Then we will choose W(K), ((Q) such that
1
w(0)
k
1.22
- (k2 + s² 2sk)
X' =
2k: 28
-1
Hence
(w)w))) T. (wxl2)x)
W.k
n (k k((k - 28) + 2.2 + 8² 25k)
n(k: - k²²²
Hence we have
(IXIO))) T (112Ix))
arg max K'ck
maxfk - 8, 0}.
Also, for any 1' & k. + maxfk - 8, 0}, we have
((WIR)) (wiee))) - -((WIR))) T (w(o)))
max(k==,0),. (((())))(((())))
< <n
Hence after the column-wise softmax and rounding to = Oloog)) bit, we have
softmax ((wrk))) T (wx10x))
= I[K' maxfk - 8, 0}}
We will then choose WVV) such that
w(O X = X1, = X
Ak' E [1].
w(Vx
Ak' E [[].
This then concludes that
A
I, (X).
The proof is complete.
Lemma E.14 (Counting Blocks). For index set I C [d - 20 satisfying //1= |12| < 10, counting function COUNT I1, Iz] can be approximated with feedfforward block and attention block A
with attention head. Formally, when Xa,e k and X3,5 1[k 11, 0, it holds that
A (X) + X) ROUND (COUNC)OOSTT I, (2)
Ag (X) + x) 0.
26

---

Published as a conference paper at ICLR 2025
Proof of Lemma E.14 We will use the feedforward block to calculate X and have
11 U12 (Lemma E.1
(x(X) -) 1 [|/]].
(o(X) + x = 2 [|]].
(x(X) + a =1
E (d i 1 e [2 + 1]}, (o(X) x)i,k Xi,
We will use X' to denote g(X) X. Then will choose W(K), w(Q) such that
1 + 1 1k' 1]
x =n
X 110,,
iE(l]
12
1
xE(1]
Hence,
((N())) (Ix())x)
W,K
11
121
x/2 1/11,,10 - (2X111))) + x221111) + nnlk [kk = =1] -
x7 77 [
=-n n
i=1
|7|
(K, x1111)) + +n1kk'
1] Σ
x2
12
Hence we have
(Ixl)))) T (w(l)x)
max k'ck
=0
Equality holds when k' or X1, x = X1z0,, for all E [[,]]]
Also, for any 1' S k. k # or N N + (11b,, for some E [[ we have
T
nn
Hence after the column-wise softmax and rounding p Ollog n) bit, we have
(softmax ((w()))) T (wr(0)))
1
ROUND
Ch=1 1(X1, X1g + p
Here the O (+*, term comes from the fact that the softmax is rounded p O(log n) bit.
We will then choose W(V) such that
( (0Xx X
XA,N 1[k' = AK [!].
W(X) = AK E [1]
This then concludes that
ROUND A(g (X) + X) ROUND (COUNT|,, 4, (2)()).
(X) + +x)x,, 0.
27

---

We will discuss several cases for the distribution of ((W(K))) (((())) ki,k" It always holds Published as a conference paper at ICLR 2025
Lemma E.15 (Matching Blocks). Given any constant c, for index set 11, 12,12,1,, [d-20) 20]. [1
[Ill & 10, [/4] = L1), a matching function Match[[1, I2, I3, I1] can be implemented with 1 feed-
forward block g and attention block A with 1 attention head. Formally, when Xa,k k X3,x
[k = [1,= 1,,11 = and S n., holds that
(x) + X) Match(/i, 12, 13. Ia)(X)
x2 Proof. We will use the feedforward block to calculate ,²,x² o,d v E Uh U 12 as in the proof of Lemmas E.13 and E.14
We then choose W(K), w(Q) such that
(( (w)())))
=- n
4c+1
[11(), - {1ob)).a) - nkk²
1=1
11|
+ 1|k 1]
4ct1
X I1bl,k + - -
The detailed construction of W(K), ((Q) is omitted here since it is similar to the proof of Lem- mas M.s.E.1. and E.
that ((W(Kx) (wil)x)), 1.k = 2022
1. If there doesn't exists W, such that Xweh ((W()))) T (w())x)) i,
Xa, then for any i
have
-ntot1
2. If there exists W, such that Xw ((W(k))) T (wil)x)) W. -nke2 -n2 The rest of the entries are all smaller than -ntot1 Finally, the largest ke satisfying that X = X will corre- sponds to ((W(k))) T (wr()))) ki,k that is at least n larger than the second largest ((w(K)x) T (w()))))
Xg. then for such k, we have
as in the proof of Lemma E.13 ".
Concluding the above discussion, we have after the column-wise softmax and rounding to O(logn) n) bit,
(softmax (111x)))((N11)) ki.k
[[k min X = N1.. {h) XI = Nl #
1l 1). otherwise
Further, will choose (V(V) such that
w(O X 1s.k'
X = -N,, AK' E [T]
W = =0 AK' E [[]
This then concludes that
A + X) Match[T, Iz, Is. (X)
This concludes the proof.
28

---

Published as a conference paper at ICLR 2025
Lemma E.16 (Matching Closest Blocks). Given any constant for index 11, 12, 13. Ia (d [d20 I1| = [/] < 10, 13 = L1), a matching closest function MatchClose[[1, I2, 13. La] can be implemented with feedforward block and 1 attention block A with 1 attention head. Formally, when Xa,k k, X3,k l = 11, X1, = and k S n., it holds that
(g X) MatchClose 12, 13, (x)(
Proof. The proof is similar to the proof of Lemma E.15 and one can design the attention pattern as ((V1))))
n 4c+1
(X1M x n(k -nnk n(k-kF
i=1
[||
+ 1[K'
n
4c+1
x 1.
nn(h 1)2 n2 2022
The rest of the proof is omitted here.
Lemma E.17 (Matching Nearest Blocks). Given any constant c. for index set 12, Is, Ia C [d 20], [Zl
[Lal and index i a matching nearest function MetchNearet (1, 12, 13 4. can be implemented with feedforward block g and attention block A with attention head. Formally, when Xd,k = k. Na,k 1[k: = 1}, X1,,1 0 and k S n
12 < [11
holds that
A (X)
MatchNNcaret//4, I2, Is, I4, (X)
Proof. The proof is similar to the proof of Lemma E.15 and one can design the attention pattern as
((w()))) T
|r|
=- n 40+1
(Kniui, X(su),,)² -n(xix k2²
u=l
|7|
+ 1[K' =
n
40+1
X2 12juj,k
+ +n(1 -x,x)² - n 2012
The of the proof omitted here.
Lemma E.18 (Matching Next Blocks). Given any constant A, for index set I I2, Is, Ls C [d 20), [ = [I2 < 10. [al = [1a and a special counting index a, a matching next function MatchNext J2, T3, T4, a) can implement with 1 feedforward block g and attention block A with 1 attention head. Formally, when Xd,k = k. Xa,5 = 1[k = 1), X = 0 and $ n², holds that
Agg X) MatchNext([1, [41 12, Ia, (Xx)
Proof. We will use the feedforward block to calculate the following next function, where
I 7 and
>
x
next (z)
The value can be arbitrary for x E [1] 3 feedforward block through combination of Lemmas
3
This function is achievable by a
5:
E.
and
E.6
The purpose of this to approximate the next function for x E ROUND(1/[r<)) [m and we have the following lemma.
29

---

Published as a conference paper at ICLR 2025
Lemma E.19. For large enough and any E ROUND(L/nn).pp) we have
next(x)) - next(x)) S nxxt(x)² +0 10.
Proof. We always have ROUND(1/1(<",pp) n (L 10
3)
0. We will discuss
20
several cases for æ E ROUND(1//n=),p))
1. If 2 To then next(x)) (2cxt(2)).
2. If r <
7 it holds that I - 1/ml I 1/n 104 20
2 3, then
1
1
(ex x xx-
+
m
m
m3
11
1
1
+
m 1
m
mm
+ 1)
This then concludes the proof.
We then choose W(K), W(() such that
(()()) T (wrlex))
k',k
AA+3 Σ (X
x111),,) nAA+1 (next(Xaa,) Nx,
+ 1K 1] n4443
x2 1 + + m444+12 4A42
Again, the detailed construction of W(K), w(Q) is omitted here since it is similar to the proof of Lemmas E.13 and E.14
We will discuss several cases for the distribution of ((W())))) (w(+))) W
It always holds
that ((W(K))) (W())x)), 1. 4442
Nk,12- then for any > we have If there doesn't exists W, such that NR. ((W(K)x) (w(()x)) i, 4443
If there exists K, such that Xw.. Nk.a then for such W. we have ((W()))) (wll)x)) Wi -neM(next(Xax) Na, 4141 The rest of the entries are all smaller than 4442
for k' satisfying X It remains discuss the distribution of ((W(k))) (wil)x)) Ki
Xx.1.- When X satisfies the condition in Definition E.12 we have that St = (Xa,k* X = p). Hence, [Sk] > 2, suppose ([
Xg takes consecutive and disjoint values ROUND(1/[rz
y, E S₁ satisfies that
next(X (Na,b)) = min Jxc next next((X..a) #ESt
2 next((Xa.)))
min
Jxc - next((Xa,k))
We will discuss several cases for y,
30

---

Published as a conference paper at ICLR 2025
If y (ext((Xa,)) and nnxtt (X.,)) are both negative, then y have,
(u - next (Xx,(())))) (e..((X..))) = (y - =) + 2next (Xax))
1
- - z)² AndA
If y - (ex,(Xa,t) and x next(Xa,t) both positive, then y < 2, and same as above have
(y (exx)((X)) ()) (y 2) 2next (nex)((aa.)))
1
S - - 2)² S 4n
If next(XX..tt and 2 next( next(Xa..t (Xa.k) have different signs, then according to Lemma E. we have, y ROUND(next( (Xa,k) ).) p) because Sy. takes consecutive and disjoint values in ROUND(L/|n p). This then implies that
(v next (xak)( (L=-nosee (x,xt((a,:)))*
5 1
n7 1 +
(nex + )
1
next&(Xalt
4A
Concluding, we always have for any k" + k* arg maxio ((Wik)x) T (wil))))
W.
T (m((0)) (Ix(13))(((x00))
((NIA))) T (m(l)x) (
kiK
Concluding the above discussion, we have after the column- wise softmax and rounding p O(logn) n) bit,
(softmax ((WiRi)) (
= k' = arg
min
|aa. next
bE(MX (1))
Further, we will choose WVV) such that
w(V X ba
X X1a.K'
Ak' E [|].
V1 = - AK' E [1]
This then concludes the proof.
E.3 BUILDING BLOCKS OF RNNS CONSTRUCTION
We will now describe the building blocks Linear RNNs construction. We will introduce some basic operations that will be used to build more complex RNNs family.
Lemma E.20 (Recent Input Memorizing). Given any constant k and C, there exists a parameter configuration of linear wnit that maintains C' dimensions of last l input vectors in the state.
Proof. Suppose the input sequence is Fia Rd, and the dimensions that the state should memorize are d1, d2...., de. We can construct the following linear unit:
h [1--1dd
I1-1,de
f4-1,1 ... h-11,0((-11) h4-1,0xk41
hx-1dd
Lemma E.21 (Summation). Given any constant and C. there exists a parameter configuration of linear unit that maintains the of one dimension of the last l input vectors in the state.
31

---

Published as a conference paper at ICLR 2025
Proof. Suppose WLOG the input sequence is #1: IR, and the dimension that the state should memorize is 1. We can construct the following linear unit:
he = 11-1,11 11-1,1 1,2... 11
Lemma E.22 (Special Position Memorizing). Given any constant and C. there exists parameter configuration of linear unit and a FFN with Ck width that maintains the C' dimensions of the input vector position
in the state.
Proof. This is direct combination of Lemma E.21 and Lemmas F.1 and E.4 The FFN can assign all the input vectors with position greater than l to 0, and permute the corresponding dimensions of first & input vectors to the first Ck dimensions of the state. The linear unit can then maintain the state.
Lemma E.23 (Recite Fixed Sequence). Given any constant integer and C, there exists FFN with width kC that can output fixed sequence of scalars that takes values in [C] on a fixed set of positions 11....l lk-
Proof. This is direct consequence of Lemma E.4
E.4 PROOF OF THEOREM 4.6
We will first restate the theorem for clarity.
Theorem 4.6. For task T {Index, AR, c-gram retrieval, Counting), there exist constant-size Trans- formers that can solve T. On the other hand, any RNN with o(n)-bit memory cannot solve T of size with any length of CoT for large enough n.
Proof. We will discuss by cases.
When 7 is Index, we will first show why RNN cannot solve the Index question without (m)) memory. The key observation is that the recurrent form of RNNs allowed the algorithm to be run in streaming fashion with o(n) bit memory. Here streaming means that the algorithm gets to look each bit of the memory sequentially and can only update constant size of memory.
Lemma E.24. Consider the following two-party game, where Alice receives string z E 1] and Bob receives an integer k and Bob wants to know the value of æg. If only Alice is allowed to send signal Bob, then (2()) bit communication complexity is required.
Proof of Lemma E.24 Suppose there exists communication protocol where B only receives o(n) bit and can perfectly decide Fy. Because Alice doesn't know k. the protocol must send the same message to Bob for all k. Hence Bob can reconstruct the whole string with bit with o(n) bit communication. This is contradiction.
Now if RNN can solve the Index problem with o(n) bit memory, then it can also solve the Index problem with o(n) bit communication complexity. This is because Alice can simply run the RNN on input z and send the hidden state to Bob. Then Bob can run the RNN with the hidden state and le to get the answer. This is contradiction to Lemma E.24 Hence RNN cannot solve the Index problem with o(n) bit memory.
On the other hand, we will show that Transformers can solve the Index problem with O(log n) bit parameters. This is because using layers of Transformer, we will implement Match Block (Lemma E.15) that can match the last query token with the position of the previous token and retrieve the type of the matched token to the query token.
When T is AR, wthout loss of generality, we assume that n is even. The proof is similar to the proof of the proof of theIndex problem. As there are n different types of tokens, we can label them as [n]. Now for any boolean sequence x E {0, 1//2 solving AR for the following input equivalent to solving the Index problem for 2::
Sn 
, 1, n/2,2,42 n/2....,
.., n//2, xn/2 2 + n/2,k
32

---

Published as a conference paper at ICLR 2025
This then implies that RNN cannot solve AR with o(n) bit memory. Transformers, the other hand, can still solve AR with Ollog n) bit parameters, we will use one layer of copying function to copy each token's 's previous token's type to it. Then token with the position of the previous token and retrieve the type of the matched token to the query token.
use the Match Block to match the last query
When T is c-gram retrieval, without loss of generality, we assume that n is a multiple of c. The proof similar to the proof of Theorem 4.6 As there are different types of tokens, we can label them as [n]] Now for any boolean sequence {0, 13m/?, solving AR for the following imput is equivalent to solving the Index problem for x:
S.n 
, 1. n/c, 2,..., 12 + +n/c, .............................. en/.....,/// n/t....../............... Inje n/c, .
This then implies that RNN cannot solve o-gram retrieval with o(n) bit memory. Transformers, on the other hand, can still solve c-gram retrieval with Ollog n) bit parameters, we will use one layer of copying functiin to copy each token's 's previous c 1 tokens' type it. Then we can use the Match Block to match the last query token with the position of the previous token and retrieve the type of the matched token to the query token.
When T is Counting, we will first show why RNN cannot solve the Counting question without (()) memory. Consider the following setup, given any x E {0, 13", the input string is j1.22.... ... where {ji... ju (i xrj 1). then solving the Counting question for this input string for queried threshold is equivalent to solving the Index problem for x. This then implies that RNN cannot solve the Counting question with o(n) bit memory.
On the other hand, we will show that Transformers can solve the Counting question with (log bit parameters. This is because using? 2 layers of Transformer, copy the last query token to the token corresponds to the threshold, and then use a COUNT block (Lemma E4 that count the number m of the appearance of the last query token in the input sequence, and then write 1/(m + 1) to one of the dimension. Finally, we can use the Feed Forward Network on the last layer to multiply threshold +1 with this value and compare the result to get the answer.
first use COPY block to
E.5 PROOF OF THEOREM 4.1
We will first prove a lemma assuming PSPACE & P/poly.
Lemma E.25. If PSPACE & P/poly, then there exists Turing machine M with linear space complexity that cannot be simulated by a polynomial-siz circuit family.
Proof. We will prove this by contradiction. Assuming for every Turing machine M with linear space complexity, there exists polynomial-size circuit family {Cn, that can simulate M. We will construct a polynomial-size circuit family {C%) } that can decide PSPACE, which contradicts the assumption that PSPACE P/poly.
Given any language L E PSPACE, we can decide L by a Turing machine Mr with space O(n²) for some constant k. We can consider another language L' [x1|=* | x L}. We decide L' by a Turing machine My with linear space complexity by checking the length of the input and then simulating Mr By the assumption, there exists polynomial-size circuit family {C..) that can simulate My.. We can then construct a polynomial-size circuit family {Cr) that can decide L by first counting the length of the input and then simulating C₂ on the extended input. This contradicts the assumption that PSPACE & P/poly.
Now we are ready to prove our theorem.
Theorem 4.1. Assuming PSPACE & P/poly, there exists an algorithmic problem such that (1) there exist constant-sie size Linear RNNs that can solve the problem with polynomial length CoT; and (2) any constant-size regular RNNs cannot solve the problem without CoT.
33

---

Published as a conference paper at ICLR 2025
Binary Message: 0101, Index:
Euler Tour: 1,5,3,2, 6, 3, Tree: Yes
8
1/2
Takenize(G)
< 2, m, 1,,,,,,,,,,
3, ,,,
10
US
d6
Ds
DFS on C
Figure 6: An example of the graph constructed from the binary sequence 0101 and the index and the corresponding DFS tour.
Proof. By Lemma E.25 we know that if PSPACE ş P/poly, then there exists a Turing machine M with linear space complexity that cannot be simulated by a polynomial-size circuit family, We will use this result to prove Theorem 4.1
We design the task follows, for any nI, m [loge n]., for any boolean input x of length m. will choose input sequence as Sin- ("-mp and the label otherwise.
YES if (x) 1 and y NO
Because we are considering regular RNN with O(m) memory. we know that we can compute the result of RNN without CoT through circuit family with size Poly(m). However, we know that M cannot be simulated by polynomial-size circuit family. Hence, no RNN family with ((m) memory can solve the task for all n.
On the other hand, we can simulate M by the RNN by maintaining the state, the pointer, and the tape of the M inside the state of the RNN, The RNN can then maintain the transition function of the Turing machine its output function as lookup table Lemma B.4 and write down the updated state, the direction of the pointer movement, and the updated cell value the pointer in its context. Paired with the ability to memorize the recent input Lemma E.20 the RNN can then simulate the rumning of the Turing machine.
Because the space complexity is linear in m, the time complexity of M exp(O(m)) which is polynomial in Hence, solve the task by an RNN with CoT and (()) memory and polynomial-size circuit family.
E.6 PROOF THEOREM
We will proceed prove main theorem, which states that Transformers with chain-of-thought can solve IsTree perfectly, while RNNs cannot. We will first restate the theorem here.
Theorem 4.7. There exist constant-size Transformers that can solve IsTree with CoT of length O(n). On the other hand, any RNN with o(n)-br memory cannot solve IsTree with any length of Co'..
Proof of Theorem 4.7
We will prove this theorem by proving the following lemmas.
Lemma E.26. For any n and RNN R with o(n) memory, R cannot perfectly solve IsTree of size n.
Lemma E.27. There exists a Transformer T with constant depth and width, and O(log n) precision. that can solve IsTree of size n perfectly with Chain of Thought.
This proof is direct combination of Lemmas E.26 and E.27
E. PROOF OF LEMMA E.26
We first reduce another problem in communication complexity to IsTree.
34

---

Published as a conference paper at ICLR 2025
Lemma E.28. Consider the following two-party game, where Alice receives string E 1 and Bob receives an integer k and Bob wants to know whether 2y = 2k-1. only Alice is allowed to send information, () bit communication complexity is required.
Proof of Lemma E.28 We can reduce this problem to the problem Lemma E.24 Considering the
game in Lemma E.24 given any I E {0, 1.", we can construct ã; Σ string of length n with only 0 and 1. Moreover, 2n = zk-1 if and only if 2g = -g-1. Hence, if Bob
mod 2. Then T is
solve the problem in Lemma E.28 with o(n) bit, he can solve the problem in Lemma E.24 This
a contradiction.
Proof of Lemma E.26 Now suppose that have a streaming algorithm for IsTree with only o(n) memory. We shall prove Alice and Bob in Lemma E.28 can use it to solve the original question with o(n) memory.
Consider the following graph with n 1--+2 2 nodes. There is a node i corresponding to each x; for i E [n] and two special nodes 1, + 2. Node i will be connected to n + 1 if 2; = 0 and n + if Iy 1. Moreover, & - 1 and k will be connected. Now the original answer Bob wants is False if and only if the graph is Tree. Hence, given access to the streaming algorithm, Alice can run on the edges that she knows exist and send the memory to Bob. Bob can then run it on the edges that he knows exist. Combining they will be able to solve the original problem. This a contradiction.
Moreover, as RNN with CoT is also streaming algorithm, it also requires ((n) memory.
E.6.2 PROOF OF LEMMA
Proof of Lemma E.27 The proof is two-folded. We will first define an algorithm that can solve IsTree by generating a sequence of vertices of length (n)), and then we will show that this sequence can be generated by a Transformer with constant depth and width, and O(log n) precision as a Chain
of Thought.
Algorithm Depth-First Search Algorithm
Require: A graph G = (V, E) with n vertices and has an ordering 61,..., Cy-
1: Initialize two stacks of vertices S1, S2 with S = [m S2 0.
2: while S₁ not empty do
3: Let v be the top of Sj. Yield v.
4: if there exists a neighbor u of e not in S₁ U S2 then
5: Choose such that edge (u, v) has the smallest possible order and push to S₁.
6: else
7: Pop from S and push to S2.
8: end if
9:
end while
Algorithm for IsTree. We define Algorithm 1 as a depth-first search algorithm that generate sequence of vertices of length O(n) that can be used to solve IsTree. We will use two stacks S1, S2 to store the vertices. 5₁ will be used to store the vertices that are not yet visited, and S2 will be used to store the vertices that are already visited. The algorithm will start with S1 = [0]] and Sy 0. At each step, the algorithm will pop the top of S, and push it to S2. Then it will push all the neighbors of the popped vertex that are not in S₁ U S2 to Sj. The algorithm will terminate when S is empty. We will denote the yielded vertice sequence for G as A(G). The following lemma shows the connection between the result of the algorithm and the IsTree problem.
Lemma E.29. For any graph G, A(G) is tree traversal of spanning tree of the connected component of G containing V1. Hence IsTreo is True if and only if G has n 1 edges and A(G) contains 2n vertices.
Proof of Lemma E.29 First, every vertex in the connected component of G containing in will be visited. This is because the algorithm will always push all the neighbors of the popped vertex that
35

---

Published as a conference paper at ICLR 2025
not in S1 U S2 to S₁. Hence, the algorithm will terminate when all the vertices in the connected component of G containing U1 are visited.
Second, every two consecutive vertices in the yielded sequence will be connected by edge. This is because the algorithm will always push one of the neighbors of the popped vertex that is not in Si U22 S1- Hence, every two consecutive vertices in the yielded sequence will be connected by an edge. On the other hand, the combination of these edges will form a tree because the algorithm will never push vertex that is already in S S2 St. Hence, the yielded sequence is tree traversal of a spanning tree of the connected component of G containing Vi.
Construction of Transformer. We will now show that the yielded sequence of Algorithm can be generated by a Transformer with constant depth and width, and (log precision as a Chain of Thought. The Transformer will generate valid yielded sequence but can terminate early if the graph is not a tree. We will now describe the Transformer in detail. We will assume the input token sequence S is follows,
S Tokenize(G'), 01....
(10)
for some 2 0 and 01 ... u, is valid yielded sequence. The length of Tokenize(G)) is 3n 2 with 3 tokens for each edges and 1 special token 
, We will further denote the input to the first layer X as Emb(S). We will similarly denote the input to layer as x(1), We will also denote the output of the last layer as Xou.
1. Layer and Layer Attention. The attention at Layer will output zero and the FFN at Layer 1 and Attention at Layer 2 will implement counting function (Definition E.8) to count the number of vertices n appears in the previous token sequence and write ROUND PP) in a new dimension i result.
2. Layer 2 FFN and Layer 3 Attention. The FFN Layer 2 and Attention Layer 3 will implement a copying function (Definition F.7) copying the first dimension and the counting dimension ij of each token to its successor at two new dimensions i2 and ig. For each edge, this moves the type of the first vertice and the number of times the first vertice appears to ~u, For every vertice in the chain of thought, this moves the type of the previous vertice them.
3. Layer 3 FFN and Layer Attention. The FFN at Layer 3 and Attention at Layer 4 will implement another copying function, copying the dimensions i2 and is of each token to its successor at two new dimensions ia and is. Especially, for each edge, this moves the type of the first vertice and the number of times the first vertice appears to the position corresponding the second vertices.
Layer 4 FFN. This FFN will process the information gathered from the previous layer and prepare for the next layer. It will make sure the following properties hold for (S)
For every token, the position number, its square, and will be kept in the last three dimensions.
For the first vertices in each edges, ~ and 
 The rest dimension will be zero
For the second vertices of each edges (a, b). there will be four dimensions ie, i7, is, in with value a, and Naaee The,. where Th, = ROUND( Happears to current edge 1).
For vertice oy vi,...,Up. Uy. there will be four dimensions i10: 411, 412, d1s with value vn W--1 and u/, uf_ (No = 0).
5. Layer 5 Attention. Combining with the previous Layer 4 FFN layer, will implement two match functions with two attention heads matching (410, (11) or (611. (10) with (is, 47) at Layer Attention, i.e. finding the edge in input for each step in the chain of thought, we will then copy (m) to dimensions is and ig.
6. Layer 6. We will use Layer 5 FFN and Layer 6 Attention to implement the match function that matches dimension 110 of the current token to i the previous token. This will match i to the first appearance of ey in the chain of thought and we will copy i of the matched token to igg. This dimension will be the first predecessor of U/ in the chain of thought (0 for 01). We will denote this predecessor of by to be f(1) it is the father of in the tree. Now we will need to split into two cases depending on whether 1 is f1x). If 01-1 = f(r)) or 1 = 0 (for vy) we will set dimension is to be 1 and ig to be O. Otherwise, we will keep dimension and ig (wnnn))
36

---

Published as a conference paper at ICLR 2025
7. Layer 7. Now will use Layer 6 FFN and Laver 7 Attention with two attention heads to implement two MatchNext function (Defmittonn... and match 2t at i1n to ig 87 respectively. We will then copy dimensions is to ig of the matched tokens to i14 to i21 (because there will be two of them).
which use ig or ig as the counting index,
The match next function will be able retrieve the first edge containing 01. For any i 2 2, one of the matches next function will be able to retrieve the next edge containing V; after (v,, 1++1) if it exists. If it doesn't It exist, the corresponding counting dimension will either be zero or no smaller than (m..(mmmmn) We will use Layer 6 FFN to decide whether the next edge exists and set dimension 814 of the output of Layer otherwise. and 815. 116 of the output of layer 6 to be the counting dimension of the next edge if it exists, or otherwise. For each edge in the original input, we will also set dimension V15 è16 to be the counting dimension of the edge.
be the other edge in the next edge if exists, or
8. Layer& 8 Attention We will grab the next edge again, in the same manner as Layer 6, but this time using dimension i15 and 116- The necessity of this step that the next edge containing (15-1, 0i) in the original graph can be the same as the (f(u), 15) and in such case need to check whether the next edge after this edge.
9. Layer 8 FFN. We now have, at each position corresponding to vt, the first edge (f()), vi) in the yielded sequence containing 2/ and the other vertex in the edge containing vy that hasn't 't been visited if it exists. If they don't exist, the corresponding dimension will be zero. This allows us to use Layer 8s FFN to decide the next vertex in the yielded sequence, which is exactly the first vertex different with f(r) in the two edges if they exist, or f(r) otherwise. We will use Layer 8 FFN calculate the potential next vertex and put it in dimension i23 and its square in 224-
10. Layer 9 Attention. Combining with Layer FFN, we will match i23 of the current token to 110 of the previous token to find the first appearance of the potential next vertex in the chain of thought. We will then copy dimension of the matched token to i25. This value being will imply this vertex has never appeared in the chain of thought before and any other value will imply the vertex has appeared before.
11. Layer FFN. We can now check several cases,
If the potential next vertex v is either f(u,) + 0 or never appears in the chain-of-thought sequence, then Layer 9 will output
IJ. which will decodes to If the potential next vertex is not f(, and appears in the Mhain-of-thoughe sequence, then Layer 9 will output nw6. which will decodes to NO, because the chain thought has already visited and hence the graph is not a tree.
If and the potential next vertex is f(u,) = this means the chain of thought has finished. In this case, layer 9 FFN will check whether the position is 3n + 2n - 1 =
5n 3 and output nws if it is, or output narg otherwise, which will decode to YES and NO respectively.
This concludes the construction of the Transformer.
E.7 PROOF OF THEOREM 4.8
The theorem is in the same vein as the recent work on the CoT for Transformer (Li et al. 2024).
which shows the constant size and constant precision Transformer with a polynomial embedding can simulate any polynomial size circuit. The major difference of our theorem is that (1) we consider Transformer with fixed word and position embedding, hence allowing the parameter number to be logarithmic in the input size, and (2) we consider simulating RNNs, which is a special kind of circuit family and hence we can use more succinct representation utilizing the structural property attached to the recursive process.
position
We will now prove Theorem 4.8, We will first restate the theorem for convenience.
Theorem 4.8. Given input length n, let R is an RNN with word embedding (E) E 18 tns where p- (log n) is the precision, the constant ng is the number of special symbols in the
"The constant A here in Definition E.12
37

---

Published as a conference paper at ICLR 2025
vocabulary the constant d is the embedding dimension. If each recurvent iteration can be computed by a circuit of size ((n) S 2P/2, and if the RNN produces the final answer after running at most n steps of CoT for some constant A >> O, then there exist Transformers with (loo n) -iii precision. O(C(n)) parameters and word embedding [W(E) ((*+ns)xd) that can produce the same final answer after running (C(n) 1)n steps of CoT.
Proof. The proof is inspired by Theorem E. from Li al. (2024).
Theorem E.30 (Theorem 3.3 of Li et al. (2024)). For any and any circuit with size T'n)) and input size n, there exists Transformer with constant depth and precision, Ollog n) width, and position embedding with size (T log n),, such that for any input S of length n, the Transformer computes
the output of the circuit on S using T(n) steps.
However, direct utilization of Theorem E.30 is not feasible because we are interested (1) O(log n) precision Transformer, and (2) simulating the RNN for n step, which would correspond to circuit with n C()) in size. However, as the calculation is recurrent, we can encode the RNN circuit in O(C(n)) parameter instead.
To do so, we will unroll the circuit of each recurrent step of the RNN into C(n) gates. We will then assign each gate a unique id in [C(n)] and assume the circuit is calculated in the order of the gate id in the following manner.
1. Each gate has type (()), which is either constant gate outputting 1, an input gate, a hidden state gate, an AND gate, or an XOR gate.
2. Each gate i's output depends on two values ((i) and r()). If ((i) is a constant gate, then ((i) and r(i) are assigned to be O. When it is input gate, ((i) will be assigned to be the coordinate of the input embedding and r(i) will be assigned to be the index of the bit of the value at ((i) coordinate. When is hidden state gate, ((i) will be assigned to the coordinate of the hidden state embedding, and ri)) will be assigned to be the index of the bit of the value at I(i) coordinate. If it is an AND gate or an XOR gate, I(1) and r(i) will be assigned to be the id of the two gates that it depends on.
We will further assume without loss of generality that the hidden state gate is the first Pd gate. The output of the last Py gate will be the next hidden state. We will also assume that the last p(A + d) to pA - gates are the output gates. We will now first describe the chain of thought that the Transformer will output and then construct the Transformer.
Chain of thought Taking any input S with length n, the Transformer will output a sequence of 0 and 1 tokens. The first tokens will be the same as the input sequence. For each a 0 and b E (C(m) + 1], the n + a (C(n) + 1) + b token is
1. the output of gate 6 when RNN circuit is calculating the output at position plus if < C(n).
2. the n + a + 1 token in the RNN chain of thought, if = C(n) +
Construction of the Transformer.
1. Layer 1. The first attention layer will output zero and the first FFN layer will be of width OCC()). encoding all the gate information. The output of the first layer at position n + a (C(n) + 1) + will have the following coordinate:
The input will be encoded in the first dimensions. a. a, b, 62 will be encoded in four different dimensions.
The gate type t(s(b)) will be encoded in the next dimension, where s(b) = (b + 1) mod (C(n) + 1) If C(n) then the gate type will be encoded 0.
The necessary dependence ((s(b)), (=bb)) and r(s(b)), (=((b) will be encoded in the next two dimensions.
38

---

Published as a conference paper at ICLR 2025
A constant 1 will be encoded in the next dimension.
2. Layer 2 Attention. Together with the Layer FFN, the Layer 2 Attention will implement two match functions (Definition E.9 to copy the output of gate ((b + 1) and r(6 + 1) when RNN circuit is calculating the output position. When the type of gate 1 is not AND or XOR, the result will not be used in the following calculation.
3. Layer 2 FEN Layer FEN will be of width O(1). The output of the layer will be
When ((n) and ((s(b) is AND XOR constant, dimension of the output will be the output of gate + when RNN circuit is calculating the output at a position.
When b ((n) and (((()) is an input or hidden state gate or b C(n) 1, one dimension of the output will be the position in the current chain of thought where the input bit hidden
state bit copied from and the other dimension will be the square of that position
When b C(n), the output remains the same the input to Layer 2 FFN.
4. Layer Attention. Layer 3 Attention will be of width O(1). Together with Layer 2 FFN, Layer 3 Attention will implement Match heads (Definition E.9) to copy the output at the position where the input bit or hidden state bit is copied from. When the type of gate + is not input or hidden state gate, the result will not be used in the following calculation.
5. Layer 3 FEN Layer FFN will be of width O(1). The output of the layer will be
When * C(n). one dimension of the output will output of gate s(b) when RNN eircuit is calculating the output + 1 C(n) + position.
When b Cn)), the output remains the same the input to Layer 3 FFN.
6. Layer 4 Layer 4 Attention will have p heads and each head will be of width O(1). Head h E [p 1] will copy the first dimension of the output of Layer FFN at position n + a a(P(n) 1) + - (p- h) and weight each of them by 2 2 (p-11/2 1)/2 and add them in one dimension. The Layer 4 FFN will calculate r when the first dimension of the input is 1 and -r otherwise. Hence, hp, h E [A A + d] token contains a dimension i which dimension of the output of the RNN at position a.
for each 2 0, the n (C(n) +
is the
7. Layer 5 Layer 5 Attention will have d + 1 heads and each head will be of width O(1). Head à E [d + will copy the dimension i of the output of Layer 4 FFN position n + a( ((n)) 1) b- (h + A)p to a disjoint dimension f1+1- The Layer 5 FFN will then make sure the output of Layer satisfies the following:
When b * C(n), one dimension of the output will be n times the output of gate s(b) when RNN circuit is calculating the output at + 1 C(n) + position plus 1, which will decode to the corresponding value.
When = (nn), the first d dimension of the output will be the same as the output of the RNN at position and the dimension will be which will decode to the same token as the chain of thought of the RNN at position + 1.
This concludes the construction the Transformer.
E.8 AR FUNCTION CALLS IS NOT ENOUGH FOR CLOSING THE REPRESENTATION GAP
Proposition E.31. For any RNN family with Olog n) bit memory and Olog n) parameter with an oracle to receive results for the AR problem (Definition 4.3) for any queries, for large enough n, the RNN can't solve the index problem (Definition) 4.2) with length n in any CoT steps.
Proof. Consider a special type of index problem where every token at the even position of the input sequence is special token x and the rest of the tokens are uniformly random. Then the oracle for the AR problem can be simulated by the RNN by simply outputting the k when the query is not k and outputting the third token when the query is However, following similar proof of Theorem 4.6 we can show that the RNN can't solve this special form of index problem with length n in any CoT steps.
39

---

Published as a conference paper at ICLR 2025
E.9 PROOF OF THEOREM E.32
We will first state the theorem for clarity.
Theorem E.32. For task T {Index, AR, c-gram retrieval, Counting), there exists Linear RNN family with (log n) bit memory and O(logn) parameter, that can solve T with In- Tn-Context RAG in O(1) CoT steps.
Proof. When task T Index, given an input sequence that ends with a query token k, RNN will generate the following search query sequence: SStrttSerh,,, ^ (?: As\ss (k 1 (N5), <EndSearch>.
Then the regular expression will match the k-th token in the input sequence. The RNN needs to recite the format of the query, remember the index k and calculate k - to generate the regular expression. As have shown in Lemmas E.20 and E.23 RNN can recite fixed sequence at fixed position and memorize the recent input, the above sequence can be generated by an RNN. The explicit construction of the RNN is omitted here.
When task 7' is AR, given an input sequence that ends with query token k; RNN will generate the following search query sequence: <StartScarch>, \b \b (\S+) \b, <EndScarch>.
Then the regular expression will match the next token after the first occurrence of k in the input sequence. The RNN needs recite the format the query and remember the query token & to generate the regular expression. The explicit construction of the RNN is omitted here.
When task T is c-gram retrieval, given an input sequence that ends with query tokens, RNN will generate the following search query sequence: <StartSearch>, <EndSearch>>
.. \b (VS+) \b,
Then the regular expression will match the next token after the first occurrence of ki,.. he- in the input sequence. The RNN needs to recite the format the query and remember the query tokens ki,... 1 to generate the regular expression. The explicit construction of the RNN is omitted here.
When task is Counting, given an input sequence that ends with query token v and a query threshold k. RNN will generate the following search query sequence <StartSearch>, (Nb e \b) {k + 1), <EndScarch>.
Then the regular expression will match the k-th occurrence of in the input sequence. The RNN needs to recite the format of the query and remember the query token v and the threshold & to generate the regular expression. The RNN can then check whether the retrieval result is FAILED to determine the count is less than k. The explicit construction of the RNN is omitted here.
E.10 PROOF OF THEOREM E.33
In this section, we will prove Theorem
E333
We will first state the theorem for convenience.
Theorem E.33, There exists a Linear RNN family with O n) bit memory and O(log n) parameter, that can solve IsTree of size n with In-Context RAG in O(n) CoT steps.
Proof of Theorem E.33
We will first define the sequence that the eutrievl-auugmente RNN will generate and then construct an RNN that can generate such a sequence.
Sequence Generation. We will use a variant of Algorithm to generate the sequence and we will use the concatenation of the tokenization of the sequence returned by Algorithm as the sequence that the retrieval augmented RNN will generate.
RNN Construction. We can use similar proof in Theorem 5.2 by having the RNN memorize local sequences and determine the phase of Algorithm 2 it is in. The RNN will maintain the length of S₂ (Lemma E.21 and the top of S₁ in the state (Lemma E.20) and it is easy to check that the retrieval function will retrieve the correct result for each search query. The way to determine the next vertex in the stack is the same as in the proof of Lemma E.27 We will omit the simple but tedious detailed construction here.
40

---

Published as a conference paper at ICLR 2025
Algorithm 2 Depth-First Search Algorithm with Retrieving
Require: A graph G E) with vertices and E has an ordering ci...., Cm-
1: Initialize two stacks of vertices Si, S2 with Si [vi], S2 = 0, list L with = 0, and a vertex
of FAILED.
2: while S1 is not empty do
3: Let v be the top of St. Push L.
4: Generate the regular expression 122 = \b ( (1S+) \b - \bu\b.
5: Let f()) be the predecessor of v in S1 for the first time and FAILED when 01-
6: Push <StartSearch>, r2, EEnSeaacch>,fC to L.
7: if o' + [(w) then
8:
Generate the regular expression
(vbr'bb Abvib
Abelb> \bo'1b). *? .,?(1x(88)) lbbl I
\bvlb> \b (1S+)
9:
else
10:
Generate the regular expression 11 = =bb (\S+) \b> \bulb
\b
11:
end if
12:
Push <StartSearch>, T1, <EndSearch> to L.
13:
if there exists a neighbor
such that (u, 2) has larger order than (v. e') when v' * f(u) there exists neighbor of such that + f() when v' f(v) then
14:
Choose such that edge (u, v) has the smallest possible order and push a to L. Let v" = u. else
15:
16:
Push FAILED to L. Let v" FAILED.
end if
18:
if v f(v) * FAILED then
Generate the regular expression
(Vbu"\b~ \brlb
\bulb> (bolbmbb(()). *? .>2(1((S5)) \b~ \bolb
F3
\bu\ba (b())))
20:
Push <StartSearch>, Ts, <EndScarch>> to L.
21:
there exists a neighbor u of o such that (u, v) has larger order than (v, r") then
22:
Choose such that edge (u, v) has the smallest possible order and push u to L. Let v" u.
23:
else
24:
Push FAILED to L. Let FAILED.
25:
end if
26:
end if
if 0
FAILED then
28:
Pop v from Si. Push v to S2. Let 0' =.
else
30:
Generate the regular expression = (1S+) \b Ibu"tb Push <StartSearch>, F4, <EndSearch>, to if o" is not in S1 then
32:
33:
Push FAILED, U. -, L. Push o" to St. Let =
34:
35:
else
36:
Let f(u") be the predecessor of u in S₁ for the first time and FAILED when e U1. Push f(u"), NO to
37:
38:
return L.
39:
end if
40: end if
41: end while
42: if S2 has n vertices then
43: Push YES to L.
return L.
45: else
Push NO to L.
47: return L.
48: end if
41

---

Published as a conference paper at ICLR 2025
E.11 PROOF OF THEOREM 5.2
In this section, we will prove Theorem 5.2 We will first restate the theorem for convenience.
Theorem 5.2. Given A, B O, for all polynomial- -imee Turing machines E TIME(n with B states and vocabulary size B, there exist Linear RNNs with B special symbols, O A log n precision and memory, and ((AB)) parameters that can output the result Tf steps of CoT with In-Context RAG.
T by running ((na)
We will denote the state of T as I,..., B (we will use 1 as the initial state) and the vocabulary T as 1..... B. We will assume T operates on an infinite tape TAPE, which sequence of cells indexed by Z. We will also assume that the tape is imitialized with all cells being 0 except for the n cell starting tape. The pointer is initialized to 1. At each time step, the Turing machine reads the cell pointed by POINTER and updates the cell pointed by POINTER and the pointer p according to the transition function [B x [B] [B] [B] {-1, which takes the current state and the current cell value (could be empty, which corresponds to + 1) as input and outputs the new state, the new cell value and the direction move the pointer. The Turing machine halts when the state is B. Because E (IME(n A) the Turing machine will always halt n steps. We will use TAPE i| as the value the i-th cell TAPE before the / timestep. We will use POINTER[[] the value of the pointer before the timestep and St] as the state of the Turing machine before the timestep. We will further use Direction(t) as the direction the pointer movement before the timestep.
Proof of Theorem 5.2
The Turing machine also has a pointer p that points to a cell in the
We will first define the sequence that the vetieevl-auggmented RNN will generate and then construct an RNN that can generate such a sequence.
Sequence generation. The input token sequence Sin will be as followed,
S. = 
, TAPE[I, 1], TAPE[L, TAPE[L, ni]
Here all the symbols the tape are represented by one special symbol in the vocabulary. Given this input token sequence, the retrieval augmented RNN will generate the following output token sequence.
S = Sin, StarttSettee,r,, 6 (?: (1) (As), <EndSearch>,
TAPE[1,1] 1]
1, [T,
<StartSearch>, 6 (2:\5\a*)) (n) (AS), <EndSearch>,
TAPE[1, m]
TAPE[[, n, <StartSearch>, ()) POINTER[]] .*?$)), <EndSearch>,
SearchResult(1),
POINTER[J], TAPE[2, POINTERIIJ State|2], Directoon!) (2).
<StartSearch>, (POINTER (.) POINTER([] +?$), <EndSeach,,,
SearchResult(t)
POINTER[(), TAPE(t + 1, [OINTERR]]] POINTERRI] Stateft + 1], Direction it 1).
Here Resucc)esult(t)) is defined
FAILED; if POINTER[]] is empty cell before TAPE(1 [
SearchResult(t)
The output token sequence simulates the Turing machine T on the tape TAPE due to the following lemma.
42

---

Published as a conference paper at ICLR 2025
Lemma E.34. Given any t E [n]] and i 1, the last string in S that contains i, i as substring i, TAPElt, i), TAPE is not empty and the empty string otherwise.
Proof. The proof is by induction, for i 1, the result holds. For any t 2, we only need notice that POINTER(t is the only cell that can be updated at time t 1.
Construction of RNN Given the input, the RNN can first iterate over 1 n and generate the first search queries and results by maintaining counter in its state and memorizing the most recent search result (Lemma E. Then it easy to see that the retrieval oracle will generate the correct Search given the input S. Therefore, we will only need to construct an RNN that can generate the rest part of S.
We will assume the RNN maintains the state and pointer of the Turing machine in its state and show that they can be updated.
Based on Lemma E.20 the RNN can maintain constant recent token types in its state, we will assume the RNN memorize the last tokens up to the most recent <StartSearch>> and also calculate the position relative to the most recent <StartSearch>, By lookup table in the FFN Lemma E,4| the RNN can output the fixed format of the search query. Similarly, RNN can output the POINTER[[] State(t). Direction(t]. the RNN can
To generate the update TAPE(t + 1, I,PONTEER with O (B²) width to memorize the transition function the Turing machine and output the update. Then, the RNN can the memorized recent input update the state and the pointer of the Turing machine at the next <StartSearch>, The proof is then complete.
FFN
E.12 PROOF OF THEOREM E.35
Theorem E.35. For task T (Index, AR, c-gram retrieval, Counting), there exists a hybrid Linear RNN (Definitions) B.7 and B.8 family with Ollog n) bit memory and O n.) parameter, that can solve T without COT.
Proof. The proof here is essentially the same the construction of the Transformer in Theorem 4.6 We would use the same Transformer layer to solve T. The only difference is that we would use the output of the RNN, instead of FFN, as the input of the Transformer layer.Alee Also for Counting, instead of using a COPY function, we write the query token in the state the RNN (Lemma E.20
E.13 PROOF OF THEOREM E.36
Theorem E.36. There exists a hybrid Linear RNN with O n.) bit memory and O(log parameter, that solve IsTree of size with chain of thought length On ((nlogn) logn).
However, instead of using regular expressions retrieve the next neighbor and parent, we will need use the Transformer layer. The Transformer layer can retrieve the parent through an attention head implementing the match closest head (Lemma E6 if the RNN part maintains the predecessor of each node in the chain of thought.
Proof. The proof is similar to the proof of Theorem E.33
Retrieving the next neighbor is more complicated and we will use (llog steps of the chain of thought to do that. Given an edge (v, e'), we will first use one match head retrieve the position p of (v, e²) in the input sequence and write it to the chain of thought. Then we will use two MatchClose heads to retrieve the edge that contains v and is closest to + for i = 0, I,..., log, n iteratively until the heads return edge that is not (v, v') or reaches log, Here 2 can be computed through doubling one of the dimensions in the state of the RNN and reset that dimension to after termination.
We will then compare the retrieved edge with the father of check if it is the same. If it is the same, we will search the next neighbor of e after the parent of in the same way. The other part of the proof is similar to the proof of Theorem R.33
43

---

Published as a conference paper at ICLR 2025
E.14 PROOF OF THEOREM 5.4
Theorem 5.4. Given A, B >> for all polynomial-time Turing machines T E TIME(n with B stales and vocabulary size B, there exist constant-size hybrid Linear RNNs with B special symbols, O(A log -bib precision and memory, and O(AB²) parameters that can outpu the result of by running (nn²) steps of CoI.
Proof. Sequence Generation. Under the same formulation of proof of the Theorem 5.2
The hybrid
RNN will output the following sequence.
Sm, POINTER(() TAPE|2, POINTER[1), State|2), Direction[]..
POINTER POINTER([) TAPE + 1, PONNER, State|t + 1], Direction + 1],
Note that Lemma E.34 still holds. We only need to prove that the hybrid architecture generate the above sequence.
Hybrid Construction. The way RNN maintains the pointers and the states is the same the proof of Theorem 5.2 Given each pointer value, we can retrieve the last value of the cell at the pointer through the layer of attention by implementing match closest head (Lemma 1,
L
44


Reviews:
[{"review_id": "RyHCoJjLH1", "reviewer": "Reviewer_bMYq", "summary": "This paper investigates the representational capabilities of RNNs and Transformers, both with and without Chain-of-Thought (CoT). Assuming a limited RNN state size of O(log n), the authors show that RNNs augmented with CoT have greater representational power than standard RNNs, but are still less powerful than Transformers with CoT. To bridge this gap, the authors propose adding in-context retrieval mechanisms to RNNs, showing that this extension enables RNNs with CoT to solve polynomial-time problems, effectively matching the capabilities of Transformers with CoT.\n\n\nInformally, the relationships can be summarized as:\n- Standard RNN < RNN + CoT < Transformer + CoT\n- RNN + In-Context Retrieval + CoT ~= Transformer + CoT", "questions": "Please address points from weaknesses as questions if authors do not agree.", "limitations": "", "rating": 3, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 1, "strengths": "- The paper is well-structured and clearly written, main ideas and results are easy to follow.\n    \n- The paper constructs building blocks for Transformers that can be used to solve various tasks, including indexing, associative recall, c-gram retrieval, and counting. It provides detailed analyses of how these problems can be addressed using different architectures and shows their limitations.\n    \n- The authors support their theoretical results with empirical evaluations on tasks such as IsTree and HotPotQA, thus strengthening their claims.\n    \n- Show that in-context retrieval, similar to attention mechanism, is sufficient to close the gap between RNNs and Transformers with CoT, and suggest how to build more expressive RNN architectures.", "weaknesses": "1. The authors compare RNNs and Transformers under different assumptions regarding their state sizes. Specifically, they limit RNNs to a limited hidden state size of O(log n), while allowing Transformers to have an effective state size that grows linearly with the input length due to the accumulation of key-value pairs in self-attention. A decoder-only Transformer can be viewed as an RNN with a recurrent state that increases linearly -- adding each processed input to the key-value cache:  in transformer h_t = Transformer(x_t, [h_1, …, h_{t-1}]), in RNN h_{t}  = RNN(x_t, h_{t-1}); In other words, the results obtained are equivalent to comparing RNNs with limited state sizes to RNNs with linearly growing state sizes. There is nothing inherently specific to Transformers in this comparison. The key differentiator in this comparison is the state size rather than any unique property of Transformers. Therefore, concluding that models with larger state sizes have more expressive power than those with limited state sizes is unsurprising.\n    \n2. Violation of initial assumptions with in-context retrieval. By introducing in-context retrieval to RNNs, the model effectively needs to store the entire input sequence, resulting in O(n) state size requirements - similar to Transformers. This breaks the initial assumption of a limited-size RNN state O(log n) used to state RNNs limitations. As a result, the proposed solution does not address the original problem under the same constraints.\n    \n3. The proposed solution of adding in-context retrieval to RNNs essentially increases the state size linearly, effectively transforming the RNN into an RNN with linearly growing state size, which resembles a decoder-only Transformer. This leads to the somewhat tautological conclusion that to close the gap between RNNs and Transformers, one should make a Transformer-like model out of an RNN.\n    \n4. Overlap with existing concepts is not discussed. The relationship between RNNs with CoT and Adaptive Computation Time (ACT) mechanisms is not fully explored. Since CoT can be seen as a form of ACT, and RNNs with ACT have been previously studied, the novelty of the findings may be limited. A deeper exploration of how this work builds upon or differs from existing research on RNNs with ACT could strengthen the contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "aqfbO76A4W", "reviewer": "Reviewer_8brv", "summary": "This paper investigates the representational limitations of Recurrent Neural Networks (RNNs) compared to Transformers in handling sequence modeling tasks. It particularly examines RNNs’ in-context retrieval capabilities, which are essential for tasks requiring memory and context comprehension. Through theoretical analysis and empirical validation, the authors identify the inherent weaknesses of RNNs in retrieving information across long sequences and propose potential methods to bridge the performance gap between RNNs and Transformers.", "questions": "Why in-context retrieval can not be deemed as a part of COT?\nSince I tried RWKV6 -7B and it can retrieve in context words without the help of regular express function but with the help of CoT. \nI tried api on https://huggingface.co/spaces/BlinkDL/RWKV-Gradio-2\nMy input is: \n\nAssistant: ”How can we craft an engaging story featuring vampires on Mars? Let's think step by step and provide an expert response. “ What's the word between 'featuring' and 'on'. Think step by step.\n\nThe result is unstable but primarily consist of 'vampire' and 'mars‘， which means RWKV has somewhat retrieval ability.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "Originality:\nThe paper tackles a unique question—whether RNNs can be enhanced to achieve Transformer-like in-context retrieval. By applying communication complexity and information-theoretic methods, it provides a fresh perspective on RNN limitations. Proposing hybrid architectures and Retrieval-Augmented Generation (RAG) is an innovative approach to bridging this gap.\n\nQuality:\nThe paper is theoretically rigorous, with strong proofs demonstrating RNN limitations. The authors conduct well-designed empirical experiments that validate the theoretical claims on both synthetic and real-world tasks, making the results robust and credible.\n\nClarity:\nThe structure is logical, with clear definitions and effective visuals that aid understanding. Complex proofs could benefit from simplified explanations, but overall, the writing is accessible and well-organized.\n\nSignificance:\nThis work addresses a critical issue in model efficiency, with practical implications for memory-constrained applications. It lays a solid theoretical foundation for future research on enhancing RNNs, potentially impacting both low-resource applications and the development of hybrid model architectures.", "weaknesses": "1.This paper spends a large section explaining the well-known concepts in AI community like Language Models and CoT while ignoring the explanation of communication complexity and other less-known concepts that are also important in the proof. It would be better to shift the focus of concept explanation. \n\n2. In some cases, with enough training, RNN may have the ability to better organize information based on the original O(n) bit. For example, based on the case of in-context retrieval, it may be possible to finetune the RNN to enable it to achieve the information retrieval process with the deepening of layers. e.g. the hidden states in shallow layers store the raw information and the deeper layers add organized information (like the regular expression for retrieval ) into the hidden state.  Please disprove this possibility.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "cFeLxfHwCk", "reviewer": "Reviewer_4LG8", "summary": "This paper compares the performance differences between RNNs and Transformers when enhanced with chain of thought (CoT) reasoning. The term RNN is used in a broad sense to refer to the class of architectures that are linear time and constant memory with respect to sequence length. This paper demonstrates, both theoretically and empirically, that while CoT enhances the performance of RNNs, it is not sufficient to close the gap between RNNs and Transformers. To close this gap, the authors demonstrate that retrieval-augmented generation or adding a single transformer layer is sufficient.", "questions": "1. Why were the word embedding and decoding heads frozen in the experiments for Figures 2 and 4?\n2. How do the author's define a \"RNN with o(n)-bit memory\"?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. This is a well-designed study with a focused question, along with theoretical and empirical support for its conclusions.\n2. This paper’s theoretical analysis is a strong contribution, given the numerous prior works that empirically find hybridization and RAG help close the performance gap between transformers and RNNs.", "weaknesses": "1. Perhaps this is a misunderstanding, but line 297 seems to contradict Arora et al. (2023), who observe that while transformers outperform RNNs on associative recall, certain RNNs can perform associative recall. Claiming that “constant-size RNNs cannot solve any of the four tasks” seems too strong, given that Theorem 4.6 is only true for “large enough $n$.” This paper would benefit from the addition of a paragraph reconciling this claim with the empirical results of prior work.\n2. The use of IsTree as a synthetic task could be better motivated. As currently written, it seems that this task was chosen simply because it does not obviously require retrieval to solve which is not unique to IsTree. It would be helpful to expand on the authors' rationale.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "4kaBiMqoq5", "reviewer": "Reviewer_aZ5D", "summary": "This paper investigates the gap in representation power between Transformers and Recurrent Neural Networks (RNNs) and examines whether RNNs, with techniques like Chain-of-Thought (CoT) prompting, can achieve comparable performance. Through theoretical analysis, the authors demonstrate that RNNs lack the expressiveness needed to handle in-context retrieval tasks, such as associative recall and graph recognition, that Transformers handle with ease. They further support these findings with empirical evidence from both synthetic tasks and natural language experiments, showing that while CoT prompting improves RNNs, it alone is insufficient to bridge the gap.\nThe authors propose that enhancing RNNs with Retrieval-Augmented Generation (RAG) or adding a Transformer layer can significantly close this performance gap, bringing RNNs closer to Transformer-level capabilities. However, they also acknowledge that practical challenges, such as computational costs and scalability, may arise when implementing these improvements, suggesting a trade-off between efficiency and retrieval effectiveness.", "questions": "None", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "A key strength of this paper is its use of theoretical lower bounds to rigorously show that RNNs lack the expressiveness needed for certain in-context retrieval tasks, unlike Transformers. This analytical approach confirms that RNN limitations are inherent, not merely practical, underscoring the need for hybrid architectures or additional layers to bridge the performance gap. Additionally, the paper validates these findings through empirical experiments across various tasks (such as associative recall and graph recognition), demonstrating that the theoretical limitations are evident in real applications as well.\nThe authors also provide practical techniques—like Retrieval-Augmented Generation (RAG) and hybrid architectures involving Transformer layers—that make their findings actionable and relevant for architecture design in large language models (LLMs). By identifying specific ways to improve RNNs, the paper offers a path toward resource-efficient alternatives to Transformers, valuable for applications in memory-constrained or mobile environments. Moreover, their analysis sheds light on the mechanisms behind in-context learning, clarifying what capabilities make Transformers excel and highlighting specific limitations in RNNs. This deeper understanding could inform future research on in-context retrieval and neural network architecture design.", "weaknesses": "Readability and Structure: The paper is hard to follow due to poor language flow, making the argument difficult to track.\nRedundant content, especially on Chain of Thought (CoT) prompting, appears in both the Related Work and Preliminaries sections. Consolidating these would improve readability. Some definitions are introduced but not utilized in the main text, adding unnecessary complexity.\n\nDepth of Discussion and Empirical Validation: The paper discusses two main ideas: first, that Chain of Thought (CoT) prompting can improve RNNs but is insufficient to fully bridge the representation gap with Transformers; and second, that enhancing the in-context retrieval capability of RNNs can potentially close this gap. However, each of these points would benefit from deeper discussion and more practical experiments. The empirical validation is insufficient; model comparisons are not “apple-to-apple, limiting the clarity and robustness of results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)