Separations in the Representational Capabilities of Transformers and Recurrent Architectures

Summary

This paper demonstrates theoretical separations in the representational abilities of Transformer and Recurrent Architectures on selected synthetic tasks, including index lookup, nearest neighbor, recognizing bounded Dyck languages, and string equality. The class of recurrent architectures examined includes popular networks such as RNNs, State Space Models, and Mamba. The results reveal significant differences in the required model sizes for each architecture to effectively represent these tasks. Additionally, the authors present some experiments that highlight the optimization efficiencies of architectures belonging to both classes on the selected tasks on practical-size sequences.

Strengths

1) This is a theoretically solid paper. The proof techniques are novel and well explained, and they might be useful for further analysis of modern deep learning architectures. 2) The architectures considered are of practical interest, and the analysis contributes to a deeper understanding of the power and limitations of these architectures. 3) The paper extends previously known separation results to a set of diverse and more real-world applicable tasks.

Weaknesses

1) The experimental part is fairly brief. The analysis of computational or statistical learning complexity is left to future work.

Questions

- Do you have any guess of which Boolean functions are conjectured to be hard to compute for 2-layers Tranformers of poly-logarithmic size? - It seems your separation results are worst-case over the input space. Can you argue whether all/some separations still hold for typical inputs, if for instance a uniform distribution over the input space is assumed?

Rating

7

Confidence

3

Soundness

3

Presentation

4

Contribution

3

Limitations

Limitations and societal impact are well addressed.

Summary

This works studies the differences between Transformers and recurrent models with respect to 4 tasks: index lookup, associative recall, string equality and Bounded Dyck languages. The authors prove that for index lookup and nearest neighbor recall, there exists a 1-layer Transformer that needs poly-logarithmic width in the number of inputs while recurrent architectures require at least linear width. On the other hand, 1-layer Transformers require at least linear width for checking string equality and Bounded-Dyck languages. The authors strengthen their claims with experiments on synthetic datasets.

Strengths

I found the following techniques were interesting: (1) the usage of Johnson-Lindenstrauss lemma to construct solutions with poly-logarithmic width in the number of inputs. (2) Using results from communication complexity to prove results for the lower-bound for the model width. While the ideas relating relating to communication complexity have been used in Sanford et al. (which is acknowledged by the authors), I nevertheless found the presentation of the lower-bound results to be interesting. I also found the writing to be refreshing; the authors are careful when interpreting the results and do not overstate the importance of their results while being precise in their claims.

Weaknesses

Overall, I worry that the contributions of the work don't make progress towards the motivating questions in the introduction: why Transformers have supplanted LSTMs (L29) in many problems. The analysis is specific to synthetic tasks with 1-layer recurrent models and Transformers (https://arxiv.org/abs/2105.11115, https://arxiv.org/abs/2208.01066). The theorems comment on the width of 1-layer Transformers but they do not seem to have any relevance to Transformers with multiple layers or when trained on other tasks. With the goal of differentiating Transformers from LSTMs, I wonder why these particular tasks were chosen and what makes them important. It would help to motivate why these tasks are important and for example, explains why Transformers are better for language modelling. The different tasks are often motivated for specific contexts such as in-context learning (associative recall) or natural language itself (Dyck languages - https://arxiv.org/abs/2105.11115). Furthermore, the results also suggest that scaling up both Transformers and recurrent networks suffices for tackling any of the tasks. Another concern is that the results are primarily for 1/2-layer Transformers. However, it is more relevant to study Transformers with multiple layers and it is unclear how to extend these results or if they have any implications for deeper Transformers. For example, we know that Dyck-languages can be represented by deeper Transformers.

Questions

To reiterate the points raised above 1. How do we extend these results to deeper networks? 2. Are the results relevant for other task or other domains and is the relevance of these results very narrow? 3. Is there a way to motivate the selection of the tasks considered in this work?

Rating

7

Confidence

4

Soundness

4

Presentation

3

Contribution

2

Limitations

The authors sufficiently acknowledge limitations. In particular they discuss that existence proofs do not guarantee such a solution is always found using gradient descent. The authors also comment on the limitations of their results to 1-layer Transformers and how they do not extend to 2-layer Transformers.

Summary

This paper analyzes the differences in terms of representations between Transformers and recurrent architectures. They highlight multiple cases: a) a setting where 1-layer Transformer can represent the task with a log number of parameters but not RNNs (index lookups) b) a case where RNNs can represent the task with a log number of parameters but not a 1-layer Transformer (Dyck language) c) a case where both 1-layer Transformer and RNNs cannot represent the task (boolean tasks like string equality). Then finally show that a 2-layer Transformer can represent the string equality task and other associative recall tasks such as nearest neighbors. They end up with some experiments to validate their theorems: they show that 1-layer transformers are great at learning the index lookup task and that RNNs learn very quickly the bounded dyck languages.

Strengths

I think that the paper is of a great value to the community. There have been many papers in the community that the study the cases where Transformers are superior over RNNs [1] and where RNNs are superior to Transformers [2]. It is great that this paper gives a unified picture of the strengths and weaknesses of both models. Besides, their analysis is clean: the theorems and proof sketch are very easy to follow. For this reason, I advocate for acceptance for this paper. [1] Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032. [2] Liu, B., Ash, J., Goel, S., Krishnamurthy, A., & Zhang, C. (2024). Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36.

Weaknesses

There are not a lot of flaws for this paper in my opinion. I will just give some suggestions: - I believe that the authors could further improve the presentation of their results. In particular, it is not clear at the beginning why they choose the tasks they propose and it sounds a bit like a "catalog". I think the authors should clearly say that they consider a case where a 1-layer transformer with log(L) params fail but not RNN, another case where RNN with log(L) params fail but not a 1-layer etc. Maybe it may be worth adding a table where the column headers are the tasks and the row headers are the two models and each entry contains the complexity of each model at each task. - Besides, I liked a lot the discussion from lines 210 to 218 and I believe that this should be earlier in the paper (even in the introduction). I think this point is central to understand one difference between Transformers and RNNs. I think that the point raised by the authors is not totally novel since a similar behavior has been reported by [1] in the case of the copy task. - Regarding the superiority of RNNs over Transformers in the case of the bounded Dyck languages, do the authors have a similar discussion to add? If yes, could they add it? Is the explanation similar to the one advanced by [2]? - Lastly, I don't know if it is possible to give any intuition about the limitation of 1-layer Transformers and RNNs at representing boolean functions? - I am not sure to follow the experiments for the bounded Dyck languages. The authors say that the 1-layer Transformer achieve near perfect accuracy up to lengths 100 but the red curve in figure 2 right is around 60% for all the iterations. Did I miss something? [1] Jelassi, S., Brandfonbrener, D., Kakade, S. M., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032. [2] Liu, B., Ash, J., Goel, S., Krishnamurthy, A., & Zhang, C. (2024). Exposing attention glitches with flip-flop language modeling. Advances in Neural Information Processing Systems, 36.

Questions

I mentioned my questions in the weaknesses section.

Rating

8

Confidence

4

Soundness

4

Presentation

3

Contribution

4

Limitations

The authors mention the limitations of their work.

Summary

In this paper, the authors study the representational separation results about two widely used classes of language models: transformers and RNNs. For a set of practically well-motivated tasks, they establish lower and upper bounds for attention based one-layer (and some for two) transformers and arbitrary RNNs using and building upon techniques from communication complexity. They also empirically validate their conclusions.

Strengths

First of all, I want to appreciate the authors for clearly elucidating the main ideas and the intuitions in a very accessible manner even for non-experts. Such style of writing is a rarity these days and the authors deserve the credit for this. Overall, it's a very important topic and problem of research interest to understand the differences between the two widely used models in the form of transformers and RNNs. They present a set of intereseting results for well-motivated tasks, and establish theoretical results for transformers needing to be of linear size in the input dimension, whereas RNNs could get away with logarithmic size, and vice versa.

Weaknesses

While I appreciate the technical results of the paper, I am not fully sure about how meaningful these results are inorder to decipher the fundamental differences between RNNs and transformers. In particular, for the index lookup task, it makes sense that the only way RNNs can retrieve the symbol $s_p$ at an unknown location $p$ revealed at the end is only via storing all the past information in its hidden state and hence $m \geq N/p$ is logical. However, if you feed the sequence in the order $p, s_1, \ldots, s_N$, would the same result hold? I believe you could get away with $\log N$ here too though I could be wrong. Shedding light on things like this could yield more insights about how these architectures are fundamentally different? Because for an index lookup task, if all we care about is retrieving a symbol at some position, we don't really care about which order you feed the input sequence right? Also on a minor note, as the authors themselves acknowledged, these results might not have full bearing on the learnability settings. It would be interesting to see how these results hold for SGD learnability on the same tasks. Food for future thought.

Questions

1. I realized that the soft-max attention in line 121 is non-causal. Would your results change if it's causal attention?

Rating

7

Confidence

4

Soundness

4

Presentation

4

Contribution

3

Limitations

Yes.

Thank you for the rebuttal

Thank you for the detailed response, particularly on the challenges of extending these results to multi-layer Transformers and for clarifying that the setup includes multi-layer RNNs. Most of my concerns were addressed and I am hence raising my score.

Thank you

I appreciate the authors' thoughtful response, and I will be raising my score accordingly.

Acknowledgement of the rebuttal

I am satisifed with the authors' response which addressed my concern. Happy to raise my score to 7.

Potential useful reference

I only followed the above discussion cursorily and there is a recent work that tries to provde lower bounds on depth for the transformer as a function of the Markovian order of the input process. Sharing it here in case it's of interest/relevance to the authors/reviewers: https://arxiv.org/abs/2407.17686v1