RNNs are not Transformers (Yet): The Key Bottleneck on In-Context Retrieval

We thank the reviewer for stating our findings as actionable and relevant. We would like to address some of the concerns here.

W1: Need to improve readability.

We appreciate the reviewer's feedback on the readability of the paper and would like to address the concerns as follows.

W1.1: Poor language flow.

We have carefully proofread the paper to ensure clarity, but the comment on "poor language flow" lacks specificity for us to act on. Could the reviewer provide concrete sections or paragraphs where the language flow is problematic?

W1.2: Redundant content on CoT

The two mentions of CoT in our paper serve different purposes. The first mention in the related works reviews prior works, and the second mention in the preliminaries introduces a formal definition of CoT. We agree with the reviewer that it is not necessary to cite previous works on CoT prompting in the second mention and have removed the corresponding sentence.

W1.3: Unused definitions

We checked the text carefully and could not identify any definitions introduced but left unused in the main text. Could the reviewer clarify which definitions they believe were unnecessary?

We would appreciate the opportunity to further discuss with the reviewer to refine the paper based on clarified and more actionable suggestions.

W2: The depth of discussion and empirical validation is not sufficient and is not ‘apple-to-apple’.

We disagree with the reviewer that the empirical validation is not sufficient to support our arguments and is not ‘apple-to-apple’.

1. On the synthetic side, we controlled the parameter size of each architecture rigorously on the IsTree task. We observe that Transformers and hybrid architectures outperform RNNs significantly on this task when the number of vertices increases. We also observed that chain-of-thought and retrieval augmented generation improved all the architectures we tested, further enhancing our theory.
2. In the real-world experiments on HotpotQA, due to the limited compute resources, we could not afford to pretrain models from scratch and had to use open-source pretrained models. However, to maximally guarantee we are close to making an apple-to-apple comparison, we tested on RNNs and hybrid models distilled from the pretrained Phi-1.5 models. In this way, we can guarantee that the models are of similar parameter sizes and go through approximately the same training dataset. The observed performance is highly consistent with our theoretical analysis.

Combining the experiments in both settings, we believe the experiments conducted so far strongly supported our theoretical results.

W1: Does the current result contradict Arora et al. (2023) which states that empirically some RNNs can perform associative recall?

There is no contradiction because we are proving the performance when the context length is long enough and Arora et al. are only testing on a maximal context length of 512. Some related empirical studies can justify our results: in $1$ , both pretrained Mamba and Mamba-2 7B models are shown to have significantly worse Phonebook-retrieval capabilities on 1K context length than Transformers with the same size and trained on the same data. We have added a paragraph discussing how our results connect with prior empirical observations.

$1$ An Empirical Study of Mamba-based Language Models

W2: What is the reason why we pick the task IsTree?

Here, we briefly argue why we prefer the IsTree task over other synthetic tasks.

1. As the reviewer mentioned, IsTree is a simple and well-known algorithmic problem that is seemingly irrelevant to in-context retrieval. However, as a classical problem studied in the literature of streaming algorithms, IsTree has been proven to require $O(n)$-bit memory for streaming algorithms, where $n$ is the number of nodes.
2. IsTree serves as a minimal abstraction of logical reasoning suggested in line 987. In a logical reasoning setting, if we view theorems as vertices and equivalences as edges, when given n theorems and n−1 relations, checking the equivalence of all the theorems reduces to verifying whether the graph forms a tree. This connection highlights IsTree's relevance as a foundational task for studying logical reasoning.

Q1: Why were the word embedding and decoding heads frozen in the experiments for Figures 2 and 4?

In line with our theoretical setup, we consider a large vocabulary that contains all the numbers in $1$ and $n$ (where n is the maximum input length). Asymptotically, this will lead to a model size linear in the input length for large $n$. To ensure that the model does not have too many parameters, we fix the word embedding and decoding heads as a function that deterministically maps numbers to vectors. We believe our experiment results still hold even if we use a more standard tokenizer that decodes numbers into several digits.

Q2: How do the authors define a "RNN with o(n)-bit memory"?

We say that an RNN has $M$-bit memory if we need $M$ bits to represent the internal state of the RNN. In other words, if there are $S$ floating-point numbers in the state and each number is represented by $f$ bits, then the RNN has M \= S \\cdot f memory.

We thank the reviewer for praising the originality and clarity of this paper and would like to address the concerns here.

W1: The paper should shift to explain communication complexity and other less-known concepts.

We would like to thank the reviewer for the valuable suggestion. We have updated our related work section to explain the relationship between communication complexity and streaming algorithms which are important for the proof. We will also add a more detailed explanation in future versions.

W2: RNNs can learn to organize and compress the context to break the theoretical lower bound.

While RNNs can and will learn an efficient compression for some prefixes, our paper argues that there exist important atomic tasks like in-context retrieval that require a state size that grows linearly with the context length. This lower bound will hold regardless of the depths of RNNs and couldn’t be solved through finetuning. At a high level, this is because the RNNs could only write the information of the prefix into the constant-size state but information-theoretically in-context retrieval requires a large state size.

Q1: Why in-context retrieval can’t be a part of CoT as models like RWKV can retrieve from context?

We do not claim that RNNs couldn’t retrieve information from context of any length. Instead, we show that the context length that RNNs can retrieve the context perfectly from is linearly scaling with its state size. Therefore, it is not a surprise that empirically RNNs can retrieve from short context as provided in the question because the state size of the model may be large enough to compress the short context into the state already. Our theory further shows that this limitation is not amenable by any type of chain-of-thought, including asking the model to explicitly retrieve from the context.

We thank the reviewer for acknowledging this paper as being well-written and would like to address the reviewer’s concern here.

W1.1: The authors compare RNNs and Transformers under different assumptions regarding their state sizes.

The reason that we compare RNNs and Transformers with different state sizes is that RNNs indeed have smaller state sizes than Transformers with the same number of parameters.

1. Having a smaller state size is a key feature of RNN, and it has motivated the recent empirical trend of developing modern RNNs (such as Mamba) to avoid the high memory cost of Transformers on long context inputs.
   1. For example, in $1$ , it is argued that “RNNs require less memory, particularly for handling long sequences.”
   2. Our work reveals a fundamental limitation of this approach, showing that such models with high memory efficiency are incapable of doing in-context retrieval.
   3. To overcome this limitation, the state size of RNN needs to scale linearly with respect to the context length for in-context retrieval.
   4. Given Transformer’s state size scales linearly with the context length, it already reaches the optimal asymptotic rate when the context length grows so it is impossible to improve the memory efficiency beyond a constant factor.
2. We compare RNNs and Transformers with the same order of number of parameters because it is typical in current empirical studies to test modern RNN architectures on in-context retrieval tasks with Transformers with the same amount of parameters $1,2,3,4$ . More specifically, we assume the parameter size of both architectures to be constant and the precision is of order $\\log n$ when the context length is $n$. This is a standard assumption in previous works considering the representation power of Transformers $5, 6$ .

$1$ RWKV: Reinventing RNNs for the Transformer Era

$2$ Mamba: Linear-Time Sequence Modeling with Selective State Spaces

$3$ Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality

$4$ Parallelizing Linear Transformers with the Delta Rule over Sequence Length

$5$ Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

$6$ Self-Attention Networks Can Process Bounded Hierarchical Languages

W1.2: Transformers can be viewed as RNNs with linear state size, and hence showing Transformers outperforming RNNs with state size $O(\\log n)$ is not surprising.

1. Models with larger state sizes do not automatically imply a stronger representation power. To prove our results, some key details in the Transformer architecture are crucial. For example, if the softmax mechanism is removed, the linear attention mechanism will not be able to perform in-context retrieval with context length $n$. This is because such Transformers are not effectively using their large state sizes. In fact, one can rewrite them as RNNs with $O(\log n)$ state sizes.
2. Can we rewrite Transformers with softmax attention as small RNNs? This is unclear a priori, but our work is able to show that Transformers with softmax attention are more expressive than RNNs, especially on in-context retrieval tasks. So there is something inherently specific to Transformers (with softmax attention) in this comparison.
3. We not only show that RNNs are worse than Transformers on the simple in-context retrieval tasks but also show that even if we use CoT, RNNs are still worse than Transformers (without CoT).

W2&3: Augmenting the model to retrieve from context violates the initial assumption as these augmentations enlarge the state size. This makes the argument tautological.
We respectfully disagree that these conclusions are tautological for the following reasons.

1. First, we would like to clarify our main results on enhancing the in-context retrieval capability of RNNs, as we suspect that the reviewers may have misunderstood our results. We proposed two approaches: In-context RAG and adding one Transformer layer.
   1. For the first approach, we are still considering RNNs with O(log n) state size, in contrast to the linear state size mentioned by the reviewer. However, the reviewer is right because the whole inference algorithm needs O(n) memory to implement the In-context RAG externally.
   2. For the second approach, it is obvious that adding 1000 Transformer layers to an RNN can greatly enhance its representation power, so as mentioned in the introduction (Line 137), we are interested in the minimal possible change to the architecture that can close the representation gap. Our final result shows that adding 1 layer can just work.
2. Both the above approaches use O(n) memory, but we argue that with o(n) memory, no approach can perform in-context retrieval. The reason is exactly what we have shown in the first part of our paper (Section 4).
3. Our results are not tautological because not all approaches with O(n) memory can perform in-context retrieval.
   1. A concrete counterexample is RNN augmented with an associative recall oracle. As stated in line 382 and Appendix E.8 of our paper, this cannot solve all the in-context retrieval tasks.
4. The proposed augmentations are closely related to the current empirical studies and can shed light on explaining empirical observations.
   1. Theorem 5.2 (explicit retrieval augmentation) is an idealized version of a Retrieval-Augmented-Generation (RAG) system. Our result justifies theoretically that this method can enhance the representation power when the backbone model is an RNN.
   2. Theorem 5.4 (hybrid model with single attention layer) shows that few attention layers can greatly enhance the retrieval capability of RNNs. This is consistent with the empirical studies showing that merely 3 layers of attention can significantly improve the in-context retrieval of a 56-layer hybrid model $1$ .

$1$ An Empirical Study of Mamba-based Language Models

W4: Lack of discussion & comparison with RNN with ACT (Adaptive Computation Time)

We would like to thank the reviewer for pointing out the related work and we will discuss ACT in the related work section. However, a key difference between ACT and CoT is that ACT will not output intermediate thought tokens. This means that CoT may be a much more powerful decoding strategy that can significantly enhance the representation power. It has been shown by many previous theoretical works that CoT can enhance the representation power of Transformers $1,2$ , but we are not aware of similar analyses for ACT. The effectiveness of CoT on Transformers motivated us to study its effectiveness on RNNs, which had not yet been thoroughly studied before our work. Therefore, we disagree that our novelty is reduced by previous studies on ACT.

$1$ Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective

$2$ Chain of Thought Empowers Transformers to Solve Inherently Serial Problems

Dear reviewer,

As the deadline approaches, we kindly ask if our responses have adequately addressed your concerns. We would greatly appreciate your feedback to ensure we have fully resolved any outstanding issues. In our response, we have clarified our contribution and we have especially stressed that we are comparing different architectures with the same number of parameters. We have also justified in detail why our results are not trivial.

Thank you for your time and consideration.